onnxruntime

92b4c663 - webgpu: Enable FlashAttention for batched GQA with right-padded prompts (#29247)

webgpu: Enable FlashAttention for batched GQA with right-padded prompts (#29247) ## Summary Lift WebGPU FlashAttention's `batch_size == 1` restriction so batched GQA with right-padded prompts (the common GenAI batched-prefill shape) takes the fused FlashAttention path instead of falling back to `ApplyAttention`. - **Per-batch seqlens in FlashAttention shaders.** Prefill, decode split-reduce, CopyKVCache, and the fused rotary-and-copyKV template now read `seqlens_k[batch_idx]` instead of hardcoding `seqlens_k[0]`. All `past_X = total_X - new_X` subtractions are clamped to avoid u32 underflow when a short batch's per-batch total is less than the batch-wide `sequence_length`. - **Indirect-dispatch sizing uses GQA's `total_sequence_length` input.** `CopyKVCache`, `SplitPackedQKVWithRotaryEmbeddingAndCopyKV`, and `FlashAttentionDecodeQKV` now take a new `total_sequence_length_input` binding (GQA input #6, GPU-resident under graph capture) for the indirect-dispatch grid sizing. This is the global max KV span across the batch by construction, replacing the previous `seqlens_k[0] + 1u` that under-dispatched whenever batch 0 wasn't the longest. Per-batch `seqlens_k[batch] + 1` still drives causal masking and K/V bounds inside the kernels. GQA now enforces `graph_capture_enabled -> past_present_share_buffer_` so the host-side `use_indirect_dispatch` predicate stays simple. - **Decoupled attention_bias stride from per-batch OOB.** `attention_bias` is still allocated to the global max `total_sequence_length`; only the causal-mask / softmax tile loops are gated by the per-batch total. The one-past-end fallback was tightened to clamp inside the same row (`offset_base + stride_total_seq - 1u`). - **Decode workgroup grid stays at global max.** `decode_qkv` keeps a workgroup grid sized to the global max tile count to keep `workgroup_idx` slicing consistent across batches, with neutral `(-inf, 0)` early-exit for tiles beyond a short batch's per-batch total so the `VxReduce` online softmax rescaling is not skewed. - **New `use_seqlen_k` template parameter** (separate from `use_indirect_dispatch` which still requires graph capture). It is enabled whenever `seqlen_k` is provided and (`graph_capture || batch_size_ > 1`). - **Rotary fix prerequisite** (`webgpu: fix GQA batched right-padded prefill with do_rotary`, 591df5b1): clamps `past_seqlen` to 0 in `RotaryEmbeddingProgram`, `FusedQKRotaryEmbeddingProgram`, and `split_packed_qkv_with_rotary_embedding`, which previously produced gibberish for the shorter batches. ## Motivation GenAI's batched prefill right-pads short prompts to the batch max and reports each batch's real length via `seqlens_k[b] = real_len[b] - 1`. The previous FlashAttention gate forced every batched call onto the slower `ApplyAttention` path, and the rotary shaders underflowed `u32` for any batch shorter than the batch-wide `sequence_length`, producing garbage Q/K positions and gibberish output text for the shorter batches. ## Test plan - [x] All `GroupQueryAttentionTest.WebGPU_*` op tests pass, including `BatchedRightPaddedRotaryPrefill` (FlashAttention path) and the new `BatchedRightPaddedRotaryPrefillFlashAttentionLargeSpread_WebGPU` covering a `real_lens` spread > tile_size - [x] phi4-prune three-prompt batched generation: coherent outputs on WebGPU matching CPU reference (3 prompts, 384 tokens, 173 tps) - [x] phi4-prune single-prompt generation regression: coherent - [x] phi4-graph-prune (graph capture enabled): `verify_model_correctness.py` 4/4 PASS; `verify_multi_gen.py` sequential + overlapping both PASS - [x] whisper-tiny-int4 transcription regression: 2/2 byte-exact with CPU - [x] Lintrunner clean on all changed files