[WebGPU] Initial PagedAttention implementation (1/n)

[Hariharan Seshadri (hariharans29)]

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 [CUDA] PagedAttention: quantized KV cache, XQA decode, MLA, QK-Norm and head sink #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.

[Copilot]

Pull request overview

This PR wires up an initial end-to-end WebGPU implementation of com.microsoft::PagedAttention using a gather→unpack→FlashAttention→repack fallback pipeline for continuous-batching / variable-Q-length workloads, and extends the WebGPU FlashAttention shaders with an optional seqlens_q input to correctly compute per-batch past_sequence_length for LEFT-aligned variable q_len.

Changes:

• Add WebGPU PagedAttention kernel implementation + WGSL programs (scatter, split packed QKV, rotary, gather KV, unpack Q, repack output) and register the kernel with the WebGPU EP.
• Update WebGPU FlashAttention (prefill + decode shaders/programs) to optionally consume seqlens_q and include it in program cache keys.
• Add WebGPU paged-attention end-to-end gtests and extend Python parity tests to cover WebGPU.

Reviewed changes

Copilot reviewed 17 out of 18 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_paged_attention.py	Adds WebGPU parity suite and EP/device plumbing; pins ONNX opset to avoid exceeding ORT max opset.
onnxruntime/test/providers/webgpu/paged_attention_test.cc	New WebGPU end-to-end correctness gtests for PagedAttention vs CPU reference.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc	Registers PagedAttention kernel for WebGPU EP.
onnxruntime/contrib_ops/webgpu/bert/paged_attention.h	Declares WebGPU PagedAttention kernel and supporting Program classes.
onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc	Implements WebGPU PagedAttention pipeline and host-side length derivation for FA.
onnxruntime/contrib_ops/webgpu/bert/paged_attention_unpack_query.wgsl.template	Unpacks packed varlen Q to padded BSNH scratch (LEFT-aligned).
onnxruntime/contrib_ops/webgpu/bert/paged_attention_split_packed_qkv.wgsl.template	Splits packed QKV into separate Q/K/V tensors.
onnxruntime/contrib_ops/webgpu/bert/paged_attention_scatter_kv.wgsl.template	Scatters new K/V tokens into the paged KV cache.
onnxruntime/contrib_ops/webgpu/bert/paged_attention_rotary.wgsl.template	Applies rotary embedding to packed 2D token tensors.
onnxruntime/contrib_ops/webgpu/bert/paged_attention_repack_output.wgsl.template	Repacks padded BSNH output back to packed varlen output.
onnxruntime/contrib_ops/webgpu/bert/paged_attention_gather_kv.wgsl.template	Gathers paged KV cache to padded contiguous BNSH scratch.
onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template	Adds optional seqlens_q path for correct causal masking under variable q_len.
onnxruntime/contrib_ops/webgpu/bert/flash_attention.h	Extends FA program APIs and ApplyFlashAttention signature with optional seqlens_q.
onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc	Wires seqlens_q through program inputs + cache hints for both prefill and decode paths.
onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template	Adds optional seqlens_q causal-bound computation in decode QKV shader.
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc	Switches to provider-neutral helper header for input validation.
onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h	Introduces shared input-validation helper for PagedAttention.
docs/design/webgpu_paged_attention.md	Adds design doc describing scope, constraints, and phased rollout plan.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 18 changed files in this pull request and generated no new comments.

Suppressed comments (2)

onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc:483

• This rejects graphs that legally omit optional outputs 1/2 (key_cache_out/value_cache_out) even though the schema marks them optional. CUDA PagedAttention allows these outputs to be null (it only enforces aliasing when they’re present) and still performs the in-place cache update via the input cache buffers. Consider matching that behavior by falling back to writing into the input cache tensors when the optional outputs are omitted.

  Tensor* key_cache_out = context.Output(1, cache_shape);
  Tensor* value_cache_out = context.Output(2, cache_shape);
  // The schema marks these outputs Optional, but the scatter kernel needs a
  // destination buffer; require both to be bound.
  ORT_RETURN_IF(key_cache_out == nullptr || value_cache_out == nullptr,

docs/design/webgpu_paged_attention.md:201

• This roadmap section says Phase 0 ("this PR") adds a stub PagedAttention whose ComputeInternal returns NOT_IMPLEMENTED, but this PR actually contains the functional Phase 1 implementation. The doc should not describe the stub as part of the current PR, otherwise readers will be misled about what’s shipping.

### Phase 0 — Skeleton (this PR)

- Add `contrib_ops/webgpu/bert/paged_attention.{h,cc}` with the kernel class
  and `ComputeInternal` returning `NOT_IMPLEMENTED`. Same shape as the CPU
  stub in #29867. Purpose: register the op with the WebGPU EP so a model

[Tianlei Wu (tianleiwu)]

Summary

Nice first cut. The six-program pipeline (split-packed-QKV → rotary → scatter → gather → unpack-Q → ApplyFlashAttention → repack) is a sensible reuse strategy for v1, the seqlens_q FlashAttention extension is minimal and correctly gated behind a shader #param so existing GQA/MHA callers compile identically, and the WGSL address models match paged_attention_impl.cu where I checked them (rotary interleaved/split pairing + sign + h >= rotary_dim pass-through, scatter block addressing, GQA head mapping). The paged_attention_helper.h relocation is a pure move (verified with git diff -M).

My findings are concentrated in three areas: host-side validation of device-derived values, WebGPU resource limits, and doc/comment drift. Details inline; cross-cutting items below.

Cross-cutting

1. Two blocking D→H readbacks per node per Run. context.CopyTensor(gpu → cpu) lands in BufferManager::Download, which ends the compute pass, Flush()es the queue, creates a fresh staging buffer, and blocks in context_.Wait. On a 32-layer model that is 64 full pipeline flushes and 64 staging-buffer creations per decoded token — and a blocking Wait is not viable on the browser main thread, which is the primary WebGPU target. The design doc frames this only as a graph-capture blocker deferred to Phase 2; it is also a v1 latency/viability issue and deserves a note there. A cheap partial fix available now: concatenate cumulative_sequence_length and past_seqlens into one small GPU scratch with a tiny copy program and issue a single Download — halves the flushes for a few lines of code.

2. KV-cache quantization is not guarded. ApplyFlashAttention opens with turbo_quant_enabled = context.KvCacheQuantizationEnabled(). With that session option on, the padded fp16 k_padded/v_padded fail the TurboQuant shape check and the user gets "TurboQuant KV cache shape mismatch for present_key" from deep inside FA. Worth rejecting up front alongside the existing softcap / local_window_size guards.

3. Unused uniforms across all six programs — scatter: kv_hidden_size, max_num_blocks_per_seq; gather: max_num_blocks_per_seq; unpack-query: batch_size, hidden_size; repack-output: max_seqlen_q, token_count; rotary: token_count. They inflate the uniform buffer and the program cache key for no benefit. (max_num_blocks_per_seq is the interesting one — see the scatter inline comment.)

4. Three copies of the token→batch linear scan in scatter, rotary, and repack_output. A shared #use helper (or a token→batch map computed once) would remove the duplication.

5. CanApplyFlashAttention is never consulted. Its preconditions (!is_packed_qkv_, head_size_ == v_head_size_, head_size % 4 == 0 / % 8 on Qualcomm) are all implied by the helper's head_size % 8 == 0 check and the gqa_params construction, so this is not a live bug — noting it only so the coupling is deliberate rather than accidental.

6. C++ test coverage gaps. paged_attention_test.cc does not cover do_rotary=1, packed-QKV, token_count == 0, or the softcap / local_window_size NOT_IMPLEMENTED guards — the last two are precisely the host-side enforce paths design doc §9 says are runnable under lavapipe, so they are cheap to add and would run in the Linux leg.

Things I specifically checked and found correct: the cache copy-through happens before the token_count == 0 early return; the gather reads the post-scatter key_cache_out/value_cache_out; kv_sequence_length = 0 deliberately hits FA's kv_empty aliasing path so FA never writes the padded scratch back; the batch linear scan handles zero-length batches (picks the first i with cum[i+1] > token_idx); seq_causal_length = past_sequence_length + q_idx_global + 1 is correct for LEFT-aligned Q once past_sequence_length is per-batch; and use_seqlens_q is threaded into both CacheHint calls so compiled variants do not collide.

[Hariharan Seshadri (hariharans29)]

Summary

Nice first cut. The six-program pipeline (split-packed-QKV → rotary → scatter → gather → unpack-Q → ApplyFlashAttention → repack) is a sensible reuse strategy for v1, the seqlens_q FlashAttention extension is minimal and correctly gated behind a shader #param so existing GQA/MHA callers compile identically, and the WGSL address models match paged_attention_impl.cu where I checked them (rotary interleaved/split pairing + sign + h >= rotary_dim pass-through, scatter block addressing, GQA head mapping). The paged_attention_helper.h relocation is a pure move (verified with git diff -M).

My findings are concentrated in three areas: host-side validation of device-derived values, WebGPU resource limits, and doc/comment drift. Details inline; cross-cutting items below.

Cross-cutting

1. Two blocking D→H readbacks per node per Run. context.CopyTensor(gpu → cpu) lands in BufferManager::Download, which ends the compute pass, Flush()es the queue, creates a fresh staging buffer, and blocks in context_.Wait. On a 32-layer model that is 64 full pipeline flushes and 64 staging-buffer creations per decoded token — and a blocking Wait is not viable on the browser main thread, which is the primary WebGPU target. The design doc frames this only as a graph-capture blocker deferred to Phase 2; it is also a v1 latency/viability issue and deserves a note there. A cheap partial fix available now: concatenate cumulative_sequence_length and past_seqlens into one small GPU scratch with a tiny copy program and issue a single Download — halves the flushes for a few lines of code.

2. KV-cache quantization is not guarded. ApplyFlashAttention opens with turbo_quant_enabled = context.KvCacheQuantizationEnabled(). With that session option on, the padded fp16 k_padded/v_padded fail the TurboQuant shape check and the user gets "TurboQuant KV cache shape mismatch for present_key" from deep inside FA. Worth rejecting up front alongside the existing softcap / local_window_size guards.

3. Unused uniforms across all six programs — scatter: kv_hidden_size, max_num_blocks_per_seq; gather: max_num_blocks_per_seq; unpack-query: batch_size, hidden_size; repack-output: max_seqlen_q, token_count; rotary: token_count. They inflate the uniform buffer and the program cache key for no benefit. (max_num_blocks_per_seq is the interesting one — see the scatter inline comment.)

4. Three copies of the token→batch linear scan in scatter, rotary, and repack_output. A shared #use helper (or a token→batch map computed once) would remove the duplication.

5. CanApplyFlashAttention is never consulted. Its preconditions (!is_packed_qkv_, head_size_ == v_head_size_, head_size % 4 == 0 / % 8 on Qualcomm) are all implied by the helper's head_size % 8 == 0 check and the gqa_params construction, so this is not a live bug — noting it only so the coupling is deliberate rather than accidental.

6. C++ test coverage gaps. paged_attention_test.cc does not cover do_rotary=1, packed-QKV, token_count == 0, or the softcap / local_window_size NOT_IMPLEMENTED guards — the last two are precisely the host-side enforce paths design doc §9 says are runnable under lavapipe, so they are cheap to add and would run in the Linux leg.

Things I specifically checked and found correct: the cache copy-through happens before the token_count == 0 early return; the gather reads the post-scatter key_cache_out/value_cache_out; kv_sequence_length = 0 deliberately hits FA's kv_empty aliasing path so FA never writes the padded scratch back; the batch linear scan handles zero-length batches (picks the first i with cum[i+1] > token_idx); seq_causal_length = past_sequence_length + q_idx_global + 1 is correct for LEFT-aligned Q once past_sequence_length is per-batch; and use_seqlens_q is threaded into both CacheHint calls so compiled variants do not collide.

Thanks for the detailed review and for validating the WGSL address models against the CUDA implementation.

I agree with the main findings and addressed the following in the latest update:

1. Metadata readback: cumulative_sequence_length and past_seqlens are now packed into one small GPU buffer before a single D→H download. The design document now calls out the remaining blocking readback as a v1 latency limitation, in addition to being a graph-capture blocker.

2. KV-cache quantization: PagedAttention now rejects the KV-cache quantization session option up front with NOT_IMPLEMENTED, rather than failing later inside FlashAttention with a TurboQuant shape mismatch.

3. Unused uniforms: The unused uniforms identified in scatter, gather, unpack-query, rotary, and repack have been removed.

4. Token-to-batch lookup: I investigated factoring the repeated linear scan into a shared WGSL helper. That refactor introduced a runtime shader/program lookup issue, so I reverted it and retained the existing inline scans for this PR. The algorithm is unchanged; this can be revisited as a follow-up cleanup.

5. FlashAttention applicability: CanApplyFlashAttention is now checked explicitly at the PagedAttention FlashAttention call site, making the relationship between the input validation and FA preconditions deliberate.

6. C++ coverage: Added coverage for the token_count == 0 cache-copy-through path. Packed-QKV and rotary behavior remain covered by the existing Python parity matrix. Focused C++ tests for rotary, packed-QKV, and the unsupported softcap / local_window_size guards can be added separately.

I also corrected the documentation drift around:

• zero-filled padding in the unpack-query path
• the seqlen_k last-valid-index convention
• the actual FlashAttention K/V bindings
• LEFT-aligned Q and the seqlens_q causal-mask path
• the one-download v1 latency limitation

The resource-limit findings are valid as well. I’ll address the checked dispatch-size, scratch-buffer, and storage-binding-limit validation in the corresponding inline comments or a focused follow-up so that resource validation remains independently reviewable.

Validation against fresh build artifacts:

• WebGpuPagedAttention.*: 6/6 passed
• TestPagedAttentionWebGpu: 32/32 passed

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 18 out of 19 changed files in this pull request and generated no new comments.

Suppressed comments (4)

onnxruntime/test/providers/webgpu/paged_attention_test.cc:12

• This test uses std::numeric_limits but does not include , which will fail to compile on some toolchains.

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <numeric>
#include <vector>

onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc:510

• This kernel currently errors out when key_cache_out/value_cache_out outputs are omitted, even though the schema marks them Optional and the CUDA kernel supports omitting them by updating the input caches in-place. Requiring the optional outputs can break valid models that only request the attention output.

  // The schema marks these outputs Optional, but the scatter kernel needs a
  // destination buffer; require both to be bound.
  ORT_RETURN_IF(key_cache_out == nullptr || value_cache_out == nullptr,
                "PagedAttention (WebGPU): key_cache_out and value_cache_out outputs "
                "are required by this kernel (schema marks them Optional, but the "

onnxruntime/test/python/transformers/test_paged_attention.py:36

• This comment says the mapping is EP -> (torch_device_when_cuda_available, ort_iobinding_device), but the dict values are only the ORT device string. This is confusing given the new Config.torch_device/ort_device split below.
  docs/design/webgpu_paged_attention.md:247
• The Phase 1 plan says the WebGPU kernel should fail if cache outputs are not aliased, but the implementation now falls back to copying into separate outputs (with a warning). This doc section should be updated to match the implemented behavior.

6. **Cache-output aliasing.** Emit `key_cache_out`/`value_cache_out` and
   verify `MutableData<T>() == input->Data<T>()`. Fail INVALID_ARGUMENT
   otherwise, exactly as CUDA does. WebGPU EP allocator reuse should make
   this straightforward.

[github-actions]

You can commit the suggested changes from lintrunner.

[Jiajia Qin (qjia7)]

LGTM

[Tianlei Wu (tianleiwu)]

The latest head addresses the earlier metadata-validation, resource-limit, quantization, cache-copy, and documentation concerns. One blocking schema-compatibility issue remains: WebGPU rejects SEPARATE-mode nodes that omit both optional cache outputs, although the schema explicitly permits that form and CUDA updates the input caches in place. I continued the existing thread with the current-head details to avoid duplicate inline feedback. Please support omitted outputs and add a focused test.

It's fine that webgpu requires the output.