Skip to content

[WebGPU] Add int8 kv cache block quantization - #32284

Merged
Sushanth Rajasankar (sushraja-msft) merged 14 commits into
mainfrom
user/sushraja/q8_kv_cache
Sep 8, 2026
Merged

[WebGPU] Add int8 kv cache block quantization#32284
Sushanth Rajasankar (sushraja-msft) merged 14 commits into
mainfrom
user/sushraja/q8_kv_cache

Conversation

@sushraja-msft

Copy link
Copy Markdown
Contributor

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 AI balanced review requested due to automatic review settings August 26, 2026 17:25
Comment thread onnxruntime/contrib_ops/webgpu/bert/kv_cache_block_quant_int8.h Fixed

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread onnxruntime/core/providers/webgpu/webgpu_provider_options.h Outdated
Comment thread onnxruntime/core/providers/webgpu/compute_context.h Outdated
Comment thread onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc
Comment thread onnxruntime/test/contrib_ops/group_query_attention_op_test.cc Outdated
Comment thread onnxruntime/test/contrib_ops/group_query_attention_op_test.cc
Comment thread onnxruntime/test/contrib_ops/group_query_attention_op_test.cc Outdated

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

Comment thread onnxruntime/contrib_ops/webgpu/bert/kv_cache_block_quant_int8.h Outdated
Comment thread onnxruntime/contrib_ops/webgpu/bert/kv_cache_block_quant_int8.h Outdated
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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);

Comment thread onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template Outdated
Comment thread onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template Outdated
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>(

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

@qjia7 Jiajia Qin (qjia7) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@sushraja-msft

Copy link
Copy Markdown
Contributor Author

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 github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

Comment thread onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Comment thread onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc Outdated
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_)) {

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

The critical long-prefill workgroup-storage issue must be fixed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details
  • Files reviewed: 23/23 changed files
  • Comments generated: 1
  • Review effort level: Balanced

Comment thread onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

A critical WGSL validation failure, incomplete storage-limit validation, and Q8 test-coverage gaps remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (3)

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

onnxruntime/test/contrib_ops/group_query_attention_op_test.cc:3812

  • This fused-Q8 test only reaches the helper's finite-value, batch-swap, and untouched-padding checks, so a deterministic but numerically wrong—or all-zero—fused rotary implementation still passes. The later rotary cross-validation uses non-aliased outputs and therefore takes the generic split/rotate path. Add a numerical reference or packed-cache assertion that specifically validates this static-cache fused route.
    onnxruntime/test/contrib_ops/group_query_attention_op_test.cc:5326
  • This test is named Decode, but RunGQATurboQuantNoPast always uses past_seq_len=0; the other Q8 decode-shaped tests seed past caches with arbitrary float bit patterns and only check for finite/nonzero output. Consequently, no test verifies the core generation flow where a valid Q8 cache produced by one invocation is consumed and extended by the next. Add a two-step prefill-then-decode comparison against the unquantized reference so packed-past copying, scale decoding, and append offsets are covered.

onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc:937

  • Q8 validation permits any head size divisible by 4, but both new quantization shaders allocate 2 * head_size * sizeof(u32) + 64 * sizeof(f32) bytes of workgroup storage. For example, FP16 head_size=4096 can pass the split-reduce storage check on a 32 KiB adapter, then this copy shader requires 33,024 bytes and fails pipeline creation. Include the quantization shader's storage requirement in this host-side device-limit check, or reuse scratch storage so every accepted shape is dispatchable.
  if (use_q8_block_quant && (parameters.head_size_ < 4 || parameters.head_size_ % 4 != 0)) {
    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                           "Q8 block-quantized KV cache requires head_size to be divisible by 4. Got head_size=",
                           parameters.head_size_);
  • Files reviewed: 23/23 changed files
  • Comments generated: 1
  • Review effort level: Balanced

Comment thread onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Broad GPU kernel changes include unresolved moderate test coverage for the Q8 interleaved-RoPE fallback.

Review details

Suppressed comments (2)

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

  • The new quantized fallback for interleaved RoPE is untested: every added Q8 rotary case leaves rotary_interleaved at its default false value. Add a packed-QKV, static-cache Q8 test with rotary_interleaved=1 and compare it with an unquantized reference so this branch cannot regress into the unsupported fused split-half shader.
    // Quantized fused rotary shaders currently implement only split-half RoPE; use the generic
    // split/rotate path for interleaved RoPE.
    if (will_use_flash_attention && parameters.past_present_share_buffer_ &&
        (!kv_cache_quant || !parameters.rotary_interleaved_)) {

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

  • The new Q8 divisibility check has no negative-path test. The existing suite covers Q4 rejection and Q8 head size 96 success, but never verifies that a Q8 head size not divisible by 4 fails with this diagnostic. Add a WebGPU GQA test using (for example) head size 98 and bit_width=8 so this validation cannot regress.
    if (kv_cache_bits == 8 && (hs < 4 || hs % 4 != 0)) {
      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                             "Q8 block-quantized KV cache requires head_size to be divisible by 4. Got head_size=",
                             hs);
  • Files reviewed: 23/23 changed files
  • Comments generated: 0 new
  • Review effort level: Balanced

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

The broad, performance-sensitive shader and attention-path changes require final human review.

Review details
  • Files reviewed: 23/23 changed files
  • Comments generated: 0 new
  • Review effort level: Balanced

Comment thread onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template
Comment thread onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template

@qjia7 Jiajia Qin (qjia7) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI, I observed a slight prefill regression on the normal KV cache path, with throughput decreasing from 3918.29 tokens/s to 3814.61 tokens/s. It's likely caused by the qk precision from output type to f32 for softmax changes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants