webgpu: Fix TurboQuant quantized KV cache for batch>1 with per-batch seqlens - #29752
Conversation
…seqlens The TurboQuant copy-to-quantized-KV-cache kernels previously read seqlen_k[0] for every batch, so batches 1..N-1 used the wrong past sequence length and produced corrupted output. genai decode runs batch_size==1 so this was not caught, but right-padded batched GQA (batch>1) needs per-batch seqlens. Fixes: - turbo_quant_hadamard.cc / turbo_quant_hadamard.wgsl.template: remove the batch_size==1 restriction and read seqlen_k[batch]. batch/head/seq are unflattened from the uniform copy sequence length (matching the host dispatch layout), then total_seq_length is derived per batch from seqlen_k[batch]. - turbo_quant_fused_rotary_hadamard.wgsl.template: compute the batch id (per Q/K/V workgroup type) before accessing seqlen_k, then read seqlen_k[batch]. Tests (WebGPU, TurboQuant-4bit EP): - WebGPU_TurboQuant_Decode_MultiBatch_UsesPerBatchSeqlensK (rotary path, fused rotary+Hadamard kernel) - WebGPU_TurboQuant_Decode_MultiBatch_NoRotary_UsesPerBatchSeqlensK (plain Hadamard kernel) Both use a swap-invariance check: running batches [A,B] with seqlens [sA,sB] and the physically-swapped [B,A] with [sB,sA] must yield swapped outputs; this fails if the kernel reads seqlen_k[0] for all batches.
There was a problem hiding this comment.
Pull request overview
Fixes the WebGPU TurboQuant (quantized KV cache) path to correctly honor per-batch seqlens_k[b] when batch_size > 1, removing the previous hard restriction to batch_size == 1 and adding regression coverage to catch seqlen indexing mistakes.
Changes:
- Update TurboQuant WGSL kernels to index
seqlen_k[batch](notseqlen_k[0]) and to computebatch/head/seqconsistently with the host dispatch layout. - Remove the
batch_size == 1validation guard in TurboQuant WebGPU host code. - Add multi-batch “swap-invariance” decode tests to validate per-batch
seqlens_kbehavior for both rotary and non-rotary TurboQuant copy paths.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | Adds swap-invariance multi-batch decode tests for TurboQuant rotary and non-rotary paths. |
| onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.wgsl.template | Unflattens (batch, head, seq) using uniform dispatch layout and switches to seqlen_k[batch] for per-batch totals. |
| onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc | Removes the previous batch_size == 1 rejection when seqlen_k is provided. |
| onnxruntime/contrib_ops/webgpu/bert/turbo_quant_fused_rotary_hadamard.wgsl.template | Computes batch id before reading seqlen_k, switching total length to seqlen_k[batch]. |
Clamp the per-batch past sequence length calculation so right-padded prefills cannot underflow u32 when their valid sequence length is shorter than the padded K/V input. Add multi-batch static-cache regression coverage for both the plain Hadamard copy path and the fused rotary path.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
Suppressed comments (1)
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc:3841
RunTurboQuantMultiBatchSwapInvariance(do_rotary=true)is intended to cover the fused TurboQuant rotary+Hadamard kernel, but it currently feeds separate Q/K/V inputs. In the WebGPU GQA implementation, non-packed inputs take the path that applies rotary to Q/K before callingApplyFlashAttention(...)withcos_cache/sin_cache == nullptr, so TurboQuant uses the non-rotary copy kernel instead. To actually test the fused rotary+Hadamard shader in the multi-batch/per-batch-seqlens scenario, provide packed QKV input and mark key/value as optional whendo_rotaryis true (similar to the existing packed-QKV tests).
tester.AddInput<float>("query", {batch_size, sequence_length, hidden_size}, concat(q0, q1));
tester.AddInput<float>("key", {batch_size, sequence_length, kv_hidden_size}, concat(k0, k1));
tester.AddInput<float>("value", {batch_size, sequence_length, kv_hidden_size}, concat(v0, v1));
tester.AddInput<float>("past_key", {batch_size, kv_num_heads, past_seq_len, kv_head_dim}, concat(pk0, pk1));
tester.AddInput<float>("past_value", {batch_size, kv_num_heads, past_seq_len, kv_head_dim}, concat(pv0, pv1));
Use per-batch sequence lengths in both TurboQuant copy shaders while preserving uniform host dispatch indexing and clamping right-padded prompt lengths. Size graph-capture indirect dispatch from the batch-wide total sequence length input, and pass the physical past-cache stride from past_key shape instead of reconstructing it in WGSL. Add rotary and non-rotary regression coverage for multi-batch cache addressing, right padding, and static-cache graph capture.
Select the concatenated rotary cache bank from the batch-global total sequence length while retaining per-batch lengths for positioning and padding. Read the GPU total-length input during graph capture across standard and TurboQuant paths, and add coverage for static and non-static packed QKV variants.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (3)
onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc:244
- TurboQuantApplyRotaryAndCopyToQuantizedKVCache adds total_seqlen as a ProgramInput whenever prepare_indirect_dispatch is true, but does not enforce that total_seqlen is non-null. A null total_seqlen would be dereferenced by ProgramInput construction.
if (prepare_indirect_dispatch) {
program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None});
}
onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc:113
- When prepare_indirect_dispatch is enabled, total_seqlen is unconditionally added as a ProgramInput, but the code does not validate total_seqlen is non-null. ProgramInput constructors dereference the Tensor pointer, so a missing total_sequence_length input would crash instead of returning a clean error.
This issue also appears on line 242 of the same file.
if (prepare_indirect_dispatch) {
program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None});
}
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc:102
- If use_total_sequence_length_input is true, total_seqlen is added as a ProgramInput without checking it is non-null. ProgramInput dereferences the Tensor pointer in its ctor, so this can crash if the optional total_sequence_length input is omitted.
if (use_total_sequence_length_input) {
program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None});
}
Review: PR #29752 — webgpu: Fix TurboQuant quantized KV cache for batch>1 with per-batch seqlens (head
|
Summary
seqlen_k[batch]instead ofseqlen_k[0], enabling correct quantized KV-cache handling forbatch_size > 1.total_sequence_length_inputwhen preparing indirect dispatch.batch_size == 1restriction.Motivation
The TurboQuant KV-cache copy kernels previously used
seqlen_k[0]for every batch. Consequently, batches1..N-1could use the wrong past sequence length and write to incorrect cache locations.Right-padded prompts introduce another case where a batch’s logical total length can be shorter than the padded K/V input length. Direct unsigned subtraction would underflow in that case.
Graph capture also requires special handling because the host-side total sequence length uniform can remain zero while the current value is supplied through a GPU input. This GPU value must be used for indirect dispatch sizing and for selecting the concatenated multi-RoPE cache bank. The multi-RoPE selection is batch-global, while rotary positions and padding checks remain per-batch.
Test plan
Regression coverage added for:
WebGPU_TurboQuant_Decode_MultiBatch_UsesPerBatchSeqlensKWebGPU_TurboQuant_Decode_MultiBatch_NoRotary_UsesPerBatchSeqlensKWebGPU_TurboQuant_Prefill_MultiBatch_RightPadding_NoRotaryWebGPU_TurboQuant_Prefill_MultiBatch_RightPadding_RotaryWebGPU_TurboQuant_IndirectDispatch_UsesGlobalLength_NoRotaryWebGPU_TurboQuant_IndirectDispatch_UsesGlobalLength_RotaryWebGPU_IndirectDispatch_MultiRotaryCache_UsesGlobalLengthWebGPU_TurboQuant_IndirectDispatch_MultiRotaryCache_UsesGlobalLengthWebGPU_MultiRotaryCache_UsesGlobalLength_NonStaticCacheVerification:
GroupQueryAttentionTest.WebGPU_TurboQuant*: 22 passedgit diff --check