[WebGPU] WIP: Optimized PagedAttention implementation (2/n) - #31727
Open
Hariharan Seshadri (hariharans29) wants to merge 36 commits into
Open
[WebGPU] WIP: Optimized PagedAttention implementation (2/n)#31727Hariharan Seshadri (hariharans29) wants to merge 36 commits into
Hariharan Seshadri (hariharans29) wants to merge 36 commits into
Conversation
Registers a NOT_IMPLEMENTED PagedAttention kernel for the WebGPU EP and lands the design doc describing the phased delivery plan. Follow-up PRs will implement the K/V writer, decode, and gather-then-flash prefill paths.
The helper is pure host code with no CUDA dependencies. Move it to contrib_ops/cpu/bert/ so it can be shared by other execution providers (CPU, WebGPU) without an EP-scope-violating include across contrib_ops/cuda/.
Replace the Phase 0 unconditional NOT_IMPLEMENTED with the full ComputeInternal control flow, minus the actual kernel launches: - Fetch all 10 inputs and route them through the shared paged_attention_helper::CheckInputs, populating a PagedAttentionParameters. - Populate the three non-helper fields (local_window_size, do_rotary, rotary_interleaved) from constructor state, matching the CUDA implementation. - Enforce the do_rotary => cos_cache && sin_cache invariant with a specific error. - Allocate output 0 with shape (token_count, hidden_size) and the two optional cache outputs with the paged shape (num_blocks, block_size, kv_num_heads, head_size). - Enforce the schema-declared alias between input caches and output caches at compute time via a raw-pointer equality check (matches CUDA; no Alias/MayInplace on the KernelDef for now). - Fast-path token_count == 0 to Status::OK. - Branch the final NOT_IMPLEMENTED into distinct decode-vs-prefill messages that reference the design doc phase, so failures are informative. Phase 1b (upcoming) will replace the two NOT_IMPLEMENTED tails with real WGSL kernel dispatch. See docs/design/webgpu_paged_attention.md §5.
… (Phase 1b.1) Adds the first per-program CUDA-parity kernel for the WebGPU PagedAttention op: a plain (non-packed, non-rotary) scatter of new K/V tokens into the block-based paged cache. * onnxruntime/contrib_ops/webgpu/bert/paged_attention_scatter_kv.wgsl.template: WGSL template. One invocation per (token, kv_head, dim); linear-scan cumulative_sequence_length to find seq_idx, then abs_slot = past_seqlens[seq] + local_tok, block_id = block_table[seq, abs_slot/block_size], slot = abs_slot%%block_size. * onnxruntime/contrib_ops/webgpu/bert/paged_attention.h: adds ScatterKVToPagedCacheProgram with 8 Uint32 uniforms. * onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc: wires the scatter program from ComputeInternal, adds .MayInplace(3,1).MayInplace(4,2) hints on the KernelDef for the aliased GenAI fast path, and a copy-fallback for the non-aliased OpTester path (mirrors GroupQueryAttention). Output tensor is zero-filled until the attention path lands in Phase 1b.3/1b.4. * onnxruntime/test/providers/webgpu/paged_attention_test.cc: 3 gtest cases covering single-token/no-past, multi-token with past, and multi-batch/multi-head with per-sequence past lengths and non-contiguous block_table. Phase 1b.1 of docs/design/webgpu_paged_attention.md.
… (Phase 1b.2)
Adds a rotary embedding WGSL program used by the WebGPU PagedAttention op
to rotate Q and K in the non-packed layout before scattering K/V into the
paged cache. Mirrors paged_attention_impl.cu::RotaryEmbeddingTNH: same
interleaved-vs-split math, same position_id = past_seqlens[b] + s formula,
and dims >= rotary_dim are copied through unchanged.
ComputeInternal flow when do_rotary=1:
1. Rotate query into output(0) (temporary layering until 1b.3 attention
lands and overwrites output with real attention results).
2. Rotate key into a GPU temp tensor.
3. Scatter rotated key + untouched value into the paged cache via the
existing ScatterKVToPagedCacheProgram from Phase 1b.1.
Value is not rotated. Packed-QKV + rotary path still returns NOT_IMPLEMENTED
(deferred to Phase 1b.2b).
Adds 3 gtests covering full-head non-interleaved, full-head interleaved,
and rotary_dim < head_size tail pass-through with multi-batch + GQA broadcast.
All 6 WebGpuPagedAttention.* tests pass.
Adds packed-QKV support to the WebGPU PagedAttention op. When `key` and `value` are absent and the `query` input carries all three projections concatenated per token (cols `[0, Q_hidden)` = Q, `[Q_hidden, Q_hidden + KV_hidden)` = K, `[Q_hidden + KV_hidden, Q_hidden + 2*KV_hidden)` = V), a new pre-pass kernel splits the packed tensor into three standalone Q, K, V tensors. The rest of the existing 1b.2 pipeline (optional non-interleaved / interleaved rotary followed by paged-KV scatter) then runs unchanged against the split tensors. Design: split-then-reuse is intentionally conservative for the first packed-QKV cut. It costs one extra full-tensor read/write in device memory per Q/K/V column relative to a fused approach, but avoids templating every downstream kernel on a packed-input layout and keeps the CPU-side output-shape and cache-mutation reasoning identical to the non-packed path. A fused rotary+scatter+packed variant can be revisited in Phase 1c when we have baseline perf numbers. Implementation: - `PagedAttentionSplitPackedQKVProgram` (new): one WGSL kernel, one invocation per input element. Dispatch is `ceil(token_count * packed_hidden_size / WORKGROUP_SIZE)` groups. Uniforms carry `token_count`, `q_hidden_size`, `kv_hidden_size`, `packed_hidden_size`, `dispatch_size`. - `paged_attention_split_packed_qkv.wgsl.template` (new): row-major linearization of `(token, packed_col)`, branching on the column range to route each element to the correct output tensor. - `PagedAttention::ComputeInternal` (edited): when `parameters.is_packed_qkv` is true, allocate three transient GPU tensors of shapes `(token_count, hidden_size)`, `(token_count, kv_hidden_size)`, `(token_count, kv_hidden_size)`, run the split kernel, and rebind `query`/`key`/`value` locally to the split outputs before falling through to the existing rotary + scatter path. Tests: extends the WebGPU PagedAttention test harness with a `bool is_packed` field on both `ScatterCase` and `RotaryCase`, a `PackQKV` helper (per-token concatenation of the reference float buffers), and three new tests exercising the packed path: `PackedQKV_NoRotary_MultiToken_SingleBatch`, `PackedQKV_Rotary_NonInterleaved_SingleToken`, `PackedQKV_Rotary_Interleaved_MultiBatch_GQA`. All 9 `WebGpuPagedAttention.*` tests pass.
…FA seqlens_q
Wires up the WebGPU PagedAttention kernel end-to-end for
continuous-batching / variable-Q-length workloads. Replaces the earlier
Phase 1a stub / Phase 1b.1-1b.2b sub-kernel scaffolding with the
production dispatch path:
scatter K/V into paged cache
-> 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)
## FlashAttention: optional seqlens_q input
The existing FA shader clamps
past_sequence_length = total_kv_b - max_seqlen_q to 0 on underflow.
That clamp is only correct for LEFT-aligned Q with past=0 (the GQA
"BatchedRightPaddedRotaryPrefill" scenario). For PagedAttention's
continuous-batching regime, past_b can be > 0 while q_len_b <
max_seqlen_q, and the clamp silently under-counts past_len_b, causing
real Q tokens to leak future KV positions through the causal mask
(observed as 85% mismatch in the s=16 packed=True test).
Introduces an optional per-batch new-Q-length input `seqlens_q` to
FA:
- `FlashAttentionProgram` / `FlashAttentionDecodeQKVProgram` gain a
`use_seqlens_q_` template-conditional gate + `seqlens_q` shader
input.
- When set, the shader computes
past_sequence_length_b = total_kv_b - seqlens_q[b] = past_len_b
which is always non-negative and correct for any (past, q_len)
combination.
- Non-PA callers (GQA / MHA / Attention) pass nullptr, leave
`use_seqlens_q_ = false`, and the shader takes the `#else` branch
that is byte-identical to the pre-patch clamp path. Zero regression
risk.
- `use_seqlens_q_` is included in the CacheHint for both programs to
avoid pipeline-cache collision.
## PagedAttention: LEFT-aligned Q layout
`RunUnpackQuery` now places real tokens at padded slots [0, q_len_b)
with padding at [q_len_b, max_seqlen_q). `RunRepackOutput` mirrors
by reading from s = local_tok directly. This matches GQA's convention
and enables the correct per-batch past_len_b via seqlens_q above.
## Test coverage
- **32 / 32** WebGPU parity configs pass in
`TestPagedAttentionWebGpu` (batch_size in {1,2}, sequence_length
in {1,4,16}, MHA + GQA, packed on/off, block_size=256). The
previously-failing test 25 (mixed q_len + past > 0) now passes.
- **5 / 5** C++ end-to-end tests
(`WebGpuPagedAttention.EndToEnd_*`), including
`EndToEnd_MixedPrefillDecode_MultiBatch_VariablePast`.
- **31 / 31** `GroupQueryAttention` WebGPU tests, including both
`BatchedRightPaddedRotaryPrefill_WebGPU` and
`BatchedRightPaddedRotaryPrefillFlashAttention_WebGPU`, unchanged
since GQA doesn't pass seqlens_q.
## Cleanup: removed transitional Phase 1b.1 / 1b.2 / 1b.2b scaffolding
- Removed `_debug_mode` schema attribute + all three mode
branches (unpack roundtrip, gather-slice verification, and
legacy output=zeros/rotated_q).
- Removed `PagedAttentionGatherVerifyProgram` + its .wgsl.template
+ `RunGatherVerify`.
- Deleted 13 transitional gtests (`ScatterOnly_*`, `Rotary_*`,
`PackedQKV_*`, `DebugMode_*`). The 5 `EndToEnd_*` tests cover the
same functionality end-to-end; Python
`TestPagedAttentionWebGpu` covers non-Linux platforms.
## Not in scope (deferred)
- `softcap != 0`: rejected with NOT_IMPLEMENTED.
- `local_window_size != -1`: rejected with NOT_IMPLEMENTED.
- `T = bfloat16`: only MLFloat16 registered.
- Graph capture (attention_metadata): documented as Phase 2 in
`docs/design/webgpu_paged_attention.md` §4.4.
- Quantized KV cache (T_CACHE), MLA / LATENT, head_sink / QK-Norm:
Phase 3 / 4 items from the design doc, tracked alongside CUDA
parity work.
## Follow-up work (later PRs)
- Rewrite C++ Rotary_* and PackedQKV_* transitional tests to
compare against an end-to-end reference so their coverage is
restored on non-Linux CI.
- Add coverage-gap tests for `block_size != 256`, empty query
(`token_count == 0`), and explicit non-default `scale`.
- Softcap + local_window_size in FlashAttentionProgram (also lifts
GQA's `CanApplyFlashAttention` bailouts).
- Wire `TestPagedAttentionWebGpu` into a WebGPU CI leg. Today the
Python parity suite runs on zero CI legs: the two WebGPU legs
(linux_webgpu.yml, windows_webgpu.yml) are build-only, and
nightly_webgpu.yml / macos-ci run `--test` but not
`--enable_transformers_tool_test`. The C++
`WebGpuPagedAttention.EndToEnd_*` gtests DO run on
nightly_webgpu (Windows A10) and macos-ci (Metal), which is
where CI protection sits today. A ~10-LOC follow-up to
nightly_webgpu.yml can add a targeted pytest step for this file.
- paged_attention_test.cc: add missing #include <limits> (uses std::numeric_limits<float>::infinity()). - paged_attention.cc: convert ORT_ENFORCE on the two optional cache outputs into ORT_RETURN_IF with a clearer error message (the scatter kernel needs both outputs, even though the schema marks them Optional). - paged_attention.cc: move the input-to-output cache copy above the token_count==0 fast path so that the non-aliased path (OpTester) leaves initialized cache outputs even when there is no scatter work to do.
Copilot review noted that the doc's Phase 0 section was labeled '(this PR)' but this PR actually delivers Phase 1. Update Phase 0 label to '(early commits in this PR)' and move the '(this PR)' marker to Phase 1, which is the final state delivered.
…ention # Conflicts: # onnxruntime/test/python/transformers/test_paged_attention.py
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
…o hari/webgpu_paged_attention
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
## Summary Update the CUDA plugin pipeline to publish release-ready Linux archives and split NuGet packages by canonical .NET RID. ## Key changes - Publish Linux `.tar.gz` archives in the `cuda_ep_cuda12_linux_gz` / `cuda_ep_cuda13_linux_gz` artifacts alongside the existing platform zip artifact. - Generate one NuGet package per enabled RID: - `Microsoft.ML.OnnxRuntime.EP.Cuda12/13.win-x64` - `Microsoft.ML.OnnxRuntime.EP.Cuda12/13.win-arm64` - `Microsoft.ML.OnnxRuntime.EP.Cuda12/13.linux-x64` - `Microsoft.ML.OnnxRuntime.EP.Cuda12/13.linux-arm64` - Keep a shared `.csproj`; `pack_nuget.py` selects the package ID and exact native RID contents at pack time. - Publish all RID-specific package IDs in pipeline metadata and update the Windows GPU NuGet test to consume the `win-x64` package. - Update packaging documentation and local examples to use the canonical RID names. ## Validation - YAML and project XML parsing passed. - Ruff check and format check passed for `pack_nuget.py`. - Editor diagnostics reported no errors in touched files. - End-to-end dry packing produced four packages, each containing only its matching runtime directory. - PowerShell/tar archive construction was tested locally. Azure pipeline execution was not run locally.
The helper is pure host code with no CUDA dependencies. Move it to contrib_ops/cpu/bert/ so it can be shared by other execution providers (CPU, WebGPU) without an EP-scope-violating include across contrib_ops/cuda/.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
TODO
Motivation and Context
TODO