[WebGPU] Add int8 kv cache block quantization

[Sushanth Rajasankar (sushraja-msft)]

Description

Extends configurable KV-cache quantization for the WebGPU Execution Provider, to support Q8:

• Q4 TurboQuant
• Q8 symmetric block quantization
• FP16/unquantized operation

The quantization mode is selected through:

ep.webgpuexecutionprovider.kvCacheQuantizationBits

Supported values are:

Value	Behavior
0	Disable KV-cache quantization - FP16/FP32 KV Cache
4	Enable Q4 TurboQuant
8	Enable Q8 symmetric block quantization

Implementation

• Generalizes Flash Attention to consume quantized KV caches based on the configured bit width.
• Adds Q8 symmetric per-vector quantization with one FP32 scale followed by packed INT8 values for each KV head.
• Preserves Q4 TurboQuant's centroid-based dequantization.
• Shares quantized-cache sizing and dequantization logic between Q4 and Q8.

KV-Cache Layout

Each quantized KV head uses the following packed layout:

[FP32 scale][packed quantized values]

The compressed size is computed as:

1 + (head_size * bit_width / 32) u32 words

Performance

KV cache	1K prefill	1K generation	4K prefill	4K generation	GPU memory at 4K
FP16	2,807 tok/s	159.39 tok/s	1,708 tok/s	127.20 tok/s	3.07 GiB
Q4	2,754 tok/s (-1.9%)	140.94 tok/s (-11.6%)	1,671 tok/s (-2.2%)	117.67 tok/s (-7.5%)	2.66 GiB (-13.4%)
Q8 block	2,798 tok/s (-0.3%)	148.32 tok/s (-6.9%)	1,672 tok/s (-2.1%)	124.55 tok/s (-2.1%)	2.83 GiB (-7.8%)

Performance summary

• All measurements are on a RTX4070 with a 2B class text model.
• Prefill performance remains within 2.2% of FP16 for both Q4 and Q8.
• Q8 retains 97.9% of FP16 generation throughput at 4K while reducing GPU memory by 7.8%.
• Q4 provides the largest memory reduction at 13.4%, with a 7.5% generation-throughput reduction at 4K.
• Alternative design with TurboQuant Q8 was considered but block Q8 increases generation throughput from 89.42 tok/s to 139.12 tok/s (+55.6%) and reduces end-to-end latency by 21.1%.

IFEval Results

IFEval results are comparable across all the quantization techniques

KV cache	Strict prompt vs. FP16	Strict instruction vs. FP16
FP16	Baseline	Baseline
Q4	+1.4% improvement	+2.3% improvement
Q8 block	-2.0% drop	-0.5% drop

Prompt Evaluation Benchmark

A 200-sample prompt evaluation benchmark was used to compare the quantized KV-cache configurations with the FP16 baseline for web scenarios. Results are reported as relative changes.

KV cache	Overall score vs. FP16	High-score rate vs. FP16
FP16	Baseline	Baseline
Q4	-2.6%	-9.1%
Q8 block	+1.1%	+9.1%

Q4 outcome distribution

Compared with FP16 across the 200 samples:

• 45 samples improved.
• 105 samples were unchanged.
• 50 samples regressed.
• Of the regressions, 33 were minor one-point changes and 17 were two points or greater.

The observed Q4 regressions primarily involved:

• Incomplete generative responses, including responses that stopped after an opening sentence.
• Repetitive generation loops.
• Missing requested items or incomplete coverage of prompt requirements.
• Responses containing incorrect details, excessive length, or repetition.
• Occasional classification errors where the response intent or sentiment was interpreted incorrectly.

Overall, Q8 block quantization remained comparable to or slightly better than FP16 on this benchmark. Q4 showed a small aggregate score reduction, with most regressions appearing in longer-form generative tasks.

[Copilot]

Pull request overview

Adds Q8 symmetric block-quantized KV-cache support to the WebGPU GroupQueryAttention and FlashAttention paths.

Changes:

• Adds Q8 quantization, packed storage, and dequantization shaders.
• Generalizes cache sizing and FlashAttention handling across Q4/Q8.
• Adds provider-option and WebGPU coverage.

Reviewed changes

Copilot reviewed 23 out of 23 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
.github/skills/ort-build/SKILL.md	Documents Windows WebGPU builds.
onnxruntime/test/providers/webgpu/webgpu_context_test.cc	Tests Q8 option parsing.
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc	Adds Q8 GQA tests.
onnxruntime/core/providers/webgpu/webgpu_provider_options.h	Defines the Q8 option value.
onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc	Parses Q8 configuration.
onnxruntime/core/providers/webgpu/compute_context.h	Updates quantization documentation.
onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.wgsl.template	Clarifies Q4 shader logic.
onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.h	Uses shared cache sizing.
onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc	Integrates shared Q4 sizing.
onnxruntime/contrib_ops/webgpu/bert/turbo_quant_fused_rotary_hadamard.wgsl.template	Refines fused Q4 shader.
onnxruntime/contrib_ops/webgpu/bert/turbo_quant_dequant.wgsl.template	Removes superseded Q4 helper.
onnxruntime/contrib_ops/webgpu/bert/kv_cache_quantization.h	Adds shared sizing helpers.
onnxruntime/contrib_ops/webgpu/bert/kv_cache_quantization_dequant.wgsl.template	Adds shared Q4/Q8 unpacking.
onnxruntime/contrib_ops/webgpu/bert/kv_cache_block_quant_int8.wgsl.template	Implements Q8 cache writes.
onnxruntime/contrib_ops/webgpu/bert/kv_cache_block_quant_int8.h	Declares Q8 programs.
onnxruntime/contrib_ops/webgpu/bert/kv_cache_block_quant_int8.cc	Dispatches Q8 quantization.
onnxruntime/contrib_ops/webgpu/bert/kv_cache_block_quant_int8_fused_rotary.wgsl.template	Fuses rotary and Q8 writes.
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc	Allocates and validates Q8 caches.
onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template	Generalizes cache dequantization.
onnxruntime/contrib_ops/webgpu/bert/flash_attention.h	Carries quantization bit width.
onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc	Routes Q4/Q8 attention paths.
onnxruntime/contrib_ops/webgpu/bert/flash_attention_paged_decode_qkv.wgsl.template	Generalizes paged unpacking.
onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template	Generalizes decode unpacking.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template:307

• This FP16 conversion can zero Q8 V rows with small but representable values: scale = max_abs / 127 may be below the FP16 subnormal rounding threshold even when quantized_value * scale is representable. Multiply the unpacked vector by the f32 scale first, then convert the final vector to q_value_t.

          let v_val =
              kv_cache_quant_unpack_vec4(packed >> quantized_shift) * q_element_t(scale);

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

onnxruntime/test/contrib_ops/group_query_attention_op_test.cc:4329

• This FP16 test uses only [-0.5, 0.5] random values and checks finite/nonzero output, so it does not exercise either repaired FP16 edge: a Q8 scale small enough to round to zero when prematurely cast, or a large-Q/small-K case whose unscaled dot overflows f16. Add deterministic inputs for those cases and compare against the reference path; otherwise both numerical fixes can regress while this test remains green.

  auto output = RunGQATurboQuant<MLFloat16>(

[Jiajia Qin (qjia7)]

Review frame

• Problem/feature validity: Validated. WebGPU already owns the Q4 quantized-cache boundary, and adding a lower-overhead symmetric Q8 format is a legitimate extension. The checked-out and GitHub head are both c4e2baf1bb3e73781ad57faefaa89e8c1150940b against base d67cc269518ab6e7e68a6dc58f7dcdce48ba51f8.
• Risk/scope: Deep. This changes the externally allocated KV-cache layout, provider-option contract, cache-write shaders, dense FlashAttention prefill/decode paths, FP16/FP32 numerical behavior, graph capture, and adapter workgroup-storage portability.
• Direction gate: Pass. Selecting Q4 TurboQuant versus symmetric Q8 at the existing WebGPU cache/FlashAttention boundary, while sharing packed-size and dequantization helpers, is the owner-level direction I would choose. The remaining issues are implementation and integration blockers rather than a reason to redesign the feature.

Confirmed findings

C1: Keep quantized FP16 decode scores in FP32 through softmax

onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template:268

The new Q8 path accumulates each dot product in FP32, but tile_qk[m][local_idx] = q_element_t(sum); immediately narrows it back to FP16 before max subtraction. With head size 128, default attention scaling, and valid FP16 Q/K values of 100, sum is about 113,100, above the maximum finite FP16 value. WGSL permits an out-of-range conversion to produce signed infinity or the nearest finite value: the former makes the following softmax evaluate inf - inf, while the latter saturates distinct scores before max subtraction. Either behavior defeats numerically stable softmax and can produce NaNs or incorrect probabilities. This is attributable to the new Q8 specialization: unquantized/Q4 narrowing existed before this PR, while this PR adds the Q8 FP32 accumulation and then narrows its result here.

The same cast also explains the fully masked-tile case from the previous C2 draft. This PR changes the quantized mask sentinel from the FP16-representable -65504 to -FLT_MAX; an implementation that converts it to -inf can produce an invalid normalization for a completely causal-masked tile. For example, new sequence length 31 and past length 100 select dense split-reduce, and the tile beginning at key 128 is entirely masked for the first query. This is not a separate root cause: retaining FP32 scores through softmax prevents that narrowing as well.

Please keep tile_qk in FP32 for quantized FP16 specializations through max, sum, and normalization, converting only where the normalized probability is consumed. Update QuantizedDecodeWorkgroupStorageBytes for the resulting FP32 array and add a zero-valid-key guard or neutral tile result so correctness does not depend on the conversion behavior. Apply the same correction to the duplicate paged template, although quantized paged dispatch is currently rejected by the host.

C2 (latent): Use FP32 QK/softmax state in the new Q8 FP16 flash-prefill path

onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template:102

This new load routes dequantized Q8 FP16 K values into the existing FP16 score accumulators (score around line 377 and the subgroup QK vectors around line 435); the subgroup load added around line 188 has the same behavior. For head size 128, default attention scaling, and valid FP16 Q/K values of 100, the score is about 113,100. The shared-memory and subgroup variants therefore overflow the finite FP16 range before max subtraction. Depending on the permitted conversion behavior, that produces infinity and NaNs or saturated scores and incorrect probabilities. The accumulator limitation predates this PR for other modes, but this PR exposes the newly supported Q8 specialization to it.

Please accumulate QK scores and maintain the online-softmax maximum and denominator in FP32 for quantized FP16 in both prefill implementations, converting only normalized weights or final values back to FP16.

S1: Align the advertised Q8 head-size contract with the GenAI allocator

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc:276

This PR deliberately accepts Q8 head sizes divisible by four and adds a head-size-96 success test. The companion microsoft/onnxruntime-genai allocator still applies head_size >= 8 && power_of_two to both 4- and 8-bit modes in src/models/kv_cache.cpp::ComputeQuantizedKvCacheHeadSize; it therefore rejects the new 96-wide configuration before ORT runs. The allocator also duplicates this PR's packed-size formula because ONNX shape inference cannot communicate the provider-specific cache layout. I found no open GenAI PR carrying the acknowledged follow-up.

This was already raised in the existing inline thread, so it should not be posted as a duplicate comment, but it remains unresolved. Please either retain the power-of-two Q8 restriction until the allocator contract is updated, or link a concrete GenAI change and establish merge/release ordering so the newly tested configuration is usable by the primary preallocating caller.

Clarifications

None.

Test coverage

The added FP32 cross-validation covers short decode, length-40 flash prefill, rotary split-reduce, packed storage, graph-capture variants, and head size 96. It does not cover the confirmed FP16 failure conditions above. Add deterministic FP16 reference comparisons for:

1. Q8 decode with large finite Q/K values whose correctly scaled score exceeds 65,504.
2. Q8 flash prefill with the same large-value condition, covering both shared-memory and subgroup-capable adapters.
3. Q8 causal split-reduce with at least one completely masked 64-token tile.

The current FP16 test uses random values in [-0.5, 0.5] and only checks finite/nonzero output, so it cannot regress either overflow path. All current WebGPU and WGSL CI checks passed, including the macOS WebGPU execution jobs, but none of the added tests use the triggering FP16 values or fully masked-tile shape.

Verdict

The Q8 feature is valid and the design direction is appropriate. C1 and C2 block merge because supported FP16 Q8 inputs exceed the precision range used before softmax in both decode and prefill; C1 also subsumes the previous standalone fully masked-tile comment. S1 remains an unresolved integration issue for the newly advertised non-power-of-two head sizes, but it already has an inline thread and should not be reposted as a duplicate. No clarification requests remain. Documentation wording such as the stale 0 = off, 4 = 4-bit comment in webgpu_execution_provider.h is cleanup only.

[Sushanth Rajasankar (sushraja-msft)]

Review frame

• Problem/feature validity: Validated. WebGPU already owns the Q4 quantized-cache boundary, and adding a lower-overhead symmetric Q8 format is a legitimate extension. The checked-out and GitHub head are both c4e2baf1bb3e73781ad57faefaa89e8c1150940b against base d67cc269518ab6e7e68a6dc58f7dcdce48ba51f8.
• Risk/scope: Deep. This changes the externally allocated KV-cache layout, provider-option contract, cache-write shaders, dense FlashAttention prefill/decode paths, FP16/FP32 numerical behavior, graph capture, and adapter workgroup-storage portability.
• Direction gate: Pass. Selecting Q4 TurboQuant versus symmetric Q8 at the existing WebGPU cache/FlashAttention boundary, while sharing packed-size and dequantization helpers, is the owner-level direction I would choose. The remaining issues are implementation and integration blockers rather than a reason to redesign the feature.

Confirmed findings

C1: Keep quantized FP16 decode scores in FP32 through softmax

onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template:268

The new Q8 path accumulates each dot product in FP32, but tile_qk[m][local_idx] = q_element_t(sum); immediately narrows it back to FP16 before max subtraction. With head size 128, default attention scaling, and valid FP16 Q/K values of 100, sum is about 113,100, above the maximum finite FP16 value. WGSL permits an out-of-range conversion to produce signed infinity or the nearest finite value: the former makes the following softmax evaluate inf - inf, while the latter saturates distinct scores before max subtraction. Either behavior defeats numerically stable softmax and can produce NaNs or incorrect probabilities. This is attributable to the new Q8 specialization: unquantized/Q4 narrowing existed before this PR, while this PR adds the Q8 FP32 accumulation and then narrows its result here.

The same cast also explains the fully masked-tile case from the previous C2 draft. This PR changes the quantized mask sentinel from the FP16-representable -65504 to -FLT_MAX; an implementation that converts it to -inf can produce an invalid normalization for a completely causal-masked tile. For example, new sequence length 31 and past length 100 select dense split-reduce, and the tile beginning at key 128 is entirely masked for the first query. This is not a separate root cause: retaining FP32 scores through softmax prevents that narrowing as well.

Please keep tile_qk in FP32 for quantized FP16 specializations through max, sum, and normalization, converting only where the normalized probability is consumed. Update QuantizedDecodeWorkgroupStorageBytes for the resulting FP32 array and add a zero-valid-key guard or neutral tile result so correctness does not depend on the conversion behavior. Apply the same correction to the duplicate paged template, although quantized paged dispatch is currently rejected by the host.

C2 (latent): Use FP32 QK/softmax state in the new Q8 FP16 flash-prefill path

onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template:102

This new load routes dequantized Q8 FP16 K values into the existing FP16 score accumulators (score around line 377 and the subgroup QK vectors around line 435); the subgroup load added around line 188 has the same behavior. For head size 128, default attention scaling, and valid FP16 Q/K values of 100, the score is about 113,100. The shared-memory and subgroup variants therefore overflow the finite FP16 range before max subtraction. Depending on the permitted conversion behavior, that produces infinity and NaNs or saturated scores and incorrect probabilities. The accumulator limitation predates this PR for other modes, but this PR exposes the newly supported Q8 specialization to it.

Please accumulate QK scores and maintain the online-softmax maximum and denominator in FP32 for quantized FP16 in both prefill implementations, converting only normalized weights or final values back to FP16.

S1: Align the advertised Q8 head-size contract with the GenAI allocator

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc:276

This PR deliberately accepts Q8 head sizes divisible by four and adds a head-size-96 success test. The companion microsoft/onnxruntime-genai allocator still applies head_size >= 8 && power_of_two to both 4- and 8-bit modes in src/models/kv_cache.cpp::ComputeQuantizedKvCacheHeadSize; it therefore rejects the new 96-wide configuration before ORT runs. The allocator also duplicates this PR's packed-size formula because ONNX shape inference cannot communicate the provider-specific cache layout. I found no open GenAI PR carrying the acknowledged follow-up.

This was already raised in the existing inline thread, so it should not be posted as a duplicate comment, but it remains unresolved. Please either retain the power-of-two Q8 restriction until the allocator contract is updated, or link a concrete GenAI change and establish merge/release ordering so the newly tested configuration is usable by the primary preallocating caller.

Clarifications

None.

Test coverage

The added FP32 cross-validation covers short decode, length-40 flash prefill, rotary split-reduce, packed storage, graph-capture variants, and head size 96. It does not cover the confirmed FP16 failure conditions above. Add deterministic FP16 reference comparisons for:

1. Q8 decode with large finite Q/K values whose correctly scaled score exceeds 65,504.
2. Q8 flash prefill with the same large-value condition, covering both shared-memory and subgroup-capable adapters.
3. Q8 causal split-reduce with at least one completely masked 64-token tile.

The current FP16 test uses random values in [-0.5, 0.5] and only checks finite/nonzero output, so it cannot regress either overflow path. All current WebGPU and WGSL CI checks passed, including the macOS WebGPU execution jobs, but none of the added tests use the triggering FP16 values or fully masked-tile shape.

Verdict

The Q8 feature is valid and the design direction is appropriate. C1 and C2 block merge because supported FP16 Q8 inputs exceed the precision range used before softmax in both decode and prefill; C1 also subsumes the previous standalone fully masked-tile comment. S1 remains an unresolved integration issue for the newly advertised non-power-of-two head sizes, but it already has an inline thread and should not be reposted as a duplicate. No clarification requests remain. Documentation wording such as the stale 0 = off, 4 = 4-bit comment in webgpu_execution_provider.h is cleanup only.

adressed both issues and bumped up precision of general flash attention qk accumulation to fp32 for all quantization levels of the KV cache.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated no new comments.

Suppressed comments (1)

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc:431

• This new fallback for quantized interleaved RoPE has no WebGPU test: the existing/new quantization cases all leave rotary_interleaved at its default, and no C++ WebGPU GQA test sets that attribute. Add packed-QKV static-cache coverage with do_rotary=1, rotary_interleaved=1, and quantization enabled (ideally Q4 and Q8) so this branch is verified against a reference; otherwise a regression can silently route back to the split-half fused shader and produce incorrect rotations.

    if (will_use_flash_attention && parameters.past_present_share_buffer_ &&
        (!kv_cache_quant || !parameters.rotary_interleaved_)) {