onnxruntime

62aced3b - [WebGPU] Initial PagedAttention implementation (1/n) (#31611)

[WebGPU] Initial PagedAttention implementation (1/n) (#31611) ### Description Wires up the WebGPU **PagedAttention** kernel end-to-end for continuous-batching / variable-Q-length workloads. Dispatch path: ``` → scatter K/V into paged cache (RunScatterKVToPagedCache) → gather paged K/V into padded BNSH scratch (RunGatherKV) → unpack packed varlen Q into LEFT-aligned BSNH scratch (RunUnpackQuery) → ApplyFlashAttention over padded scratch → repack padded output back to (token_count, hidden_size) (RunRepackOutput) ``` This is Phase 1 of a planned 4-phase rollout; see the roadmap at the bottom for the delivery plan. The v1 kernel is `MLFloat16`-only; `softcap`, `local_window_size`, and `bfloat16` are explicitly rejected with `NOT_IMPLEMENTED`. ### FlashAttention: optional `seqlens_q` input The existing FA shader clamps `past_sequence_length = total_kv_b − max_seqlen_q` to zero on underflow. That clamp is only correct for LEFT-aligned Q with `past = 0` (the GQA `BatchedRightPaddedRotaryPrefill` scenario). Under continuous batching, `past_b > 0` while `q_len_b < max_seqlen_q` is common, and the clamp silently under-counts `past_len_b` — causing real Q tokens to leak future KV positions through the causal mask (~85% output mismatch in the mixed-`q_len` test). The fix adds an **optional per-batch new-Q-length** input to FA: - `FlashAttentionProgram` and `FlashAttentionDecodeQKVProgram` gain a `use_seqlens_q_` template-conditional gate and a `seqlens_q` shader input. - When bound, the shader computes `past_sequence_length_b = total_kv_b − seqlens_q[b] = past_len_b` — always non-negative and correct for any `(past, q_len)` combination. - Non-PA callers (GQA / MHA / Attention) pass `nullptr`, `use_seqlens_q_ = false`, and the shader takes the byte-identical `#else` branch with the pre-existing clamp. Zero regression risk. - `use_seqlens_q_` is included in each program's `CacheHint` to prevent pipeline-cache collision. ### PagedAttention: LEFT-aligned Q layout `RunUnpackQuery` places real Q tokens at padded slots `[0, q_len_b)` with padding at `[q_len_b, max_seqlen_q)`; `RunRepackOutput` mirrors that layout. Matches GQA's existing convention and lets FA's `use_seqlens_q` path compute the correct `past_len_b`. ### Test coverage | Suite | Result | Notes | |---|---|---| | `TestPagedAttentionWebGpu` (Python parity) | **32 / 32** | MHA + GQA, packed on/off, `batch_size ∈ {1, 2}`, `sequence_length ∈ {1, 4, 16}`, `total_sequence_length ∈ {32, 64}`, `block_size = 256` | | `WebGpuPagedAttention.EndToEnd_*` (C++) | **5 / 5** | Includes `EndToEnd_MixedPrefillDecode_MultiBatch_VariablePast`, the exact bug-fix path | | `GroupQueryAttentionTest.*_WebGPU` (regression) | **31 / 31** | Includes `BatchedRightPaddedRotaryPrefill_WebGPU` and `BatchedRightPaddedRotaryPrefillFlashAttention_WebGPU` — unchanged since GQA does not pass `seqlens_q` | | MHA / Attention / TurboQuant / QKNorm WebGPU tests | **22 / 22** | No regressions | ### Files changed - **New:** `paged_attention_gather_kv.wgsl.template`, `paged_attention_unpack_query.wgsl.template`, `paged_attention_repack_output.wgsl.template`. - **Modified:** `flash_attention.{cc,h,wgsl.template}`, `flash_attention_decode_qkv.wgsl.template`, `paged_attention.{cc,h}`, C++ + Python tests, design doc. - **Renamed:** `test_paged_attention_cuda.py → test_paged_attention.py` (adds `TestPagedAttentionWebGpu`). ### Roadmap Work will be split into 4 phases and delivered incrementally as time permits (not back-to-back). - **Phase 1 (this PR).** Initial functionally-correct implementation using the gather-then-flash fallback path described above. Goal is early review and unblocking downstream work. - **Phase 2.** Optimized implementation: KV-page-aware prefill and decode FlashAttention kernels that back the PagedAttention op directly, with the Phase 1 gather-then-flash path retained as a fallback for cases FlashAttention doesn't cover. This also lifts `softcap` and `local_window_size` inside `FlashAttentionProgram`, which also drops GQA's `CanApplyFlashAttention` bailouts. - **Phase 3.** Support the features introduced by microsoft/onnxruntime#29912 — quantized KV cache (`T_CACHE`), head-sink, QK-Norm — starting with quantized KV. - **Phase 4.** Tuning driven by real-world model traces. Not planned for this PR: `T = bfloat16` (blocked on Dawn stability), graph-capture with `attention_metadata` sizing bound (design-doc §4.4), MLA / LATENT layout (Phase 4+ as customer need materializes). CI coverage note: `TestPagedAttentionWebGpu` currently runs on **zero** CI legs — the two WebGPU CI workflows are build-only, and `nightly_webgpu.yml` / macos-ci run `--test` without `--enable_transformers_tool_test`. The `WebGpuPagedAttention.EndToEnd_*` C++ gtests DO run on `nightly_webgpu` and macos-ci. Wiring the Python parity suite into a WebGPU CI leg is a small follow-up. --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>