webgpu: Fix TurboQuant quantized KV cache for batch>1 with per-batch seqlens

[qjia7]

Summary

• Fix the TurboQuant copy-to-quantized-KV-cache WebGPU kernels to read seqlen_k[batch] instead of seqlen_k[0], so each batch uses its own past sequence length.
• Remove the batch_size == 1 restriction in TurboQuantCopyToQuantizedKVCache and TurboQuantApplyRotaryAndCopyToQuantizedKVCache.
• In turbo_quant_hadamard.wgsl.template, unflatten batch/head/seq from the uniform copy sequence length (matching the host dispatch layout), then derive per-batch total_seq_length from seqlen_k[batch].
• In turbo_quant_fused_rotary_hadamard.wgsl.template, compute the batch id (per Q/K/V workgroup type) before accessing seqlen_k.
• Add two multi-batch swap-invariance tests (rotary and non-rotary paths).

Motivation

PR #29247 enabled FlashAttention for batched GQA with right-padded prompts, and the later TurboQuant work added a quantized KV cache path. However, the TurboQuant KV-cache copy kernels read seqlen_k[0] for every batch, so batches 1..N-1 used the wrong past length and produced corrupted output. This was not caught because genai decode runs batch_size == 1; right-padded batched GQA (batch > 1) with quantized KV cache exercises the broken path.

Test plan

• WebGPU_TurboQuant_Decode_MultiBatch_UsesPerBatchSeqlensK (rotary → fused rotary+Hadamard kernel)
• WebGPU_TurboQuant_Decode_MultiBatch_NoRotary_UsesPerBatchSeqlensK (non-rotary → 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.
• Full GroupQueryAttentionTest suite passes: 72 passed, 14 skipped (CUDA-only, no CUDA device); 0 failures.
• lintrunner -a: no issues.
• Release build (Windows, D3D12 WebGPU).

[Copilot]

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] (not seqlen_k[0]) and to compute batch/head/seq consistently with the host dispatch layout.
• Remove the batch_size == 1 validation guard in TurboQuant WebGPU host code.
• Add multi-batch “swap-invariance” decode tests to validate per-batch seqlens_k behavior 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].