[WebGPU] Optimized PagedAttention implementation (2/n) by hariharans29 · Pull Request #31727 · microsoft/onnxruntime

Hariharan Seshadri (hariharans29) · 2026-08-07T23:54:09Z

[WebGPU] PagedAttention: direct paged decode, fused paged prefill, Unpack/Repack skip (Phase 2 partial)

This PR is the Phase 2 follow-up to #31611. It replaces the "always gather + always Unpack/Repack" v1 fallback with two paged-aware FlashAttention programs that read the paged KV cache directly, and a fast path that lets FA consume the packed varlen Q buffer without materializing padded BSNH scratch. Net effect: ~2× faster decode, ~1.15× faster uniform prefill, ~1.3× faster varlen prefill on the shape matrix below, with no regressions.

The Phase 1 gather-then-flash fallback shipped in #31611 remains intact and still runs on adapters / configs where the paged-aware shaders can't safely dispatch (see Correctness invariants below).

What's shipped

Direct paged split-reduce decode. FlashAttentionPagedDecodeQKV + FlashAttentionPagedDecodeVxReduce index key_cache / value_cache directly through block_table. Selected when max_seqlen_q < 32 — mirrors the dense-FA split-reduce threshold. Eliminates the dense K/V scratch and its gather bandwidth for every decode step.
Fused paged prefill. FlashAttentionPagedPrefillProgram is a straight port of the dense-FA prefill shader's shared-memory path with page-table-aware K/V tile loads (bert/flash_attention_paged_prefill.wgsl.template). Supports fp16, BSNH Q, packed varlen Q (q_varlen template variant), and variable-Q-length causal masking via seqlen_k + seqlens_q. No attention_bias / head_sink / TurboQuant.
Unpack/Repack skip fast paths. When direct paged attention runs, we can hand FA a rank-4 view over the raw packed Q buffer instead of allocating padded BSNH scratch:
- Uniform mode (B * max_seqlen_q == token_count): view is [B, max_seqlen_q, N, H]. Covers decode, B==1 prefill, and equal-length batched prefill (the common continuous-batching case).
- Varlen mode: view is [token_count, 1, N, H] plus cumulative_seqlens_q; only the fused paged-prefill shader can index it (q_varlen).
Skipping Unpack+Repack removes 2 dispatches (~300–500 µs of CPU dispatch cost per Run on D3D12) plus a B * max_seqlen_q * hidden * 2 B scratch allocation (tens of MB at long prefill).

Dispatch-count reduction

Route (no rotary, non-packed)	#31611 (merged)	This PR (shm-path adapters)
Decode (`max_seqlen_q < 32`)	Scatter + Gather + UnpackQ + DecodeQKV + DecodeVxReduce + Repack = 6	Scatter + PagedDecodeQKV + PagedDecodeVxReduce = 3
Prefill (`max_seqlen_q ≥ 32`)	Scatter + Gather + UnpackQ + FlashAttention + Repack = 5	Scatter + FlashAttentionPagedPrefill = 2

Decode's FA is 2 kernels (split-K: QKV + VxReduce); prefill's FA is 1 kernel (FlashAttentionProgram). On configs where ShouldRunFusedPagedPrefill rejects (fp32, head_size > 256, block_size < max_k_step), the prefill row falls back to the 5-dispatch #31611 cascade; decode's direct paged split-reduce path has no such gate. Neither route is adapter-gated — the paged shaders use no subgroup intrinsics and run on every WebGPU adapter that meets the fp16 / shm-budget / alignment predicates.

Correctness invariants

Prefill selection consults one shared predicate:

bool ShouldRunFusedPagedPrefill(context, is_fp16, max_seqlen_q,
                                head_size, block_size);

It rejects (→ gather-then-flash fallback) when any of:

!is_fp16 — only fp16 variant is compiled today.
max_seqlen_q < 32 — decode uses the split-reduce programs instead.
head_size exceeds the workgroup shared-memory budget (fp16: head_size > 256).
block_size < max_k_step — the fused shader assumes one K/V tile lives in one paged block (one block_table lookup per tile). paged_attention_helper only enforces block_size >= 16 power-of-two; e.g. block_size=16 with fp16 head_size<=128 (max_k_step=32) would splice into a physically-adjacent block that isn't the next entry in the table.

The fused paged-prefill shader uses only workgroup shared memory (no subgroup intrinsics), so there is no adapter-class gate — subgroup adapters (Qualcomm / AMD / Intel with subgroups) take the paged shm kernel directly instead of falling back to gather + dense-FA-subgroup.

Because the same predicate gates the "skip RunGatherKV", "skip q_padded scratch", and "select fused shader" decisions, the three cannot drift.

Decode selection (max_seqlen_q < 32) is a pure shape check — no adapter, dtype, or block-size gate. The direct paged split-reduce kernels (FlashAttentionPagedDecodeQKV + FlashAttentionPagedDecodeVxReduce) are the sole decode path when the kernel dispatches at all (fp16 is enforced at kernel registration, so no fp32 fallback is possible). Unlike fused prefill, the decode kernels do one block_table lookup per K/V slot rather than per tile, so they have no block_size alignment requirement.

WGSL correctness gotcha handled in the fused prefill shader. cumulative_seqlens_q is array<i32> but row indices are u32. Both loadq and writeo explicitly cast (u32(cumulative_seqlens_q[b]) + q_idx); without the cast, tint surfaces the type-resolution failure as an opaque absl::…raw_hash_map<>::at at runtime.

Performance

Machine: dev-box discrete WebGPU adapter (D3D12), 24-core host. Google Benchmark harness at onnxruntime/test/onnx/microbenchmark/paged_attention.cc, --benchmark_min_time=0.3s, wall-clock timing via UseManualTime().

Earlier revisions of this PR included an ORT_WEBGPU_PAGED_ATTENTION_USE_FUSED env-var kill switch used for A/B measurement against the #31611 cascade. That toggle has been removed (the direct/fused paths are selected internally by shape and config; the numbers below are the reason). The A/B was performed by temporarily broadening the toggle locally to also force use_direct_paged_decode=false and skip_unpack_repack=false, so fused=0 exercised the exact gather-then-flash cascade shipped in #31611. All numbers below are with that broadened toggle; the broadening was reverted before final push.

Column meanings: nH = num query heads, nKV = num KV heads, H = head dim. Shape families:

MHA_H64 (nH=16, nKV=16, H=64), MHA_H128 (nH=16, nKV=16, H=128)
GQA_Qwen (nH=14, nKV=2, H=128), GQA_Llama (nH=32, nKV=4, H=128)

Decode (16 shapes)

Shape (B/nH/nKV/H/past)	this PR (µs)	#31611 (µs)	Speedup
1/16/16/128/2048	669	3612	5.40×
2/16/16/128/512	627	3248	5.18×
1/16/16/64/2048	861	3835	4.45×
2/16/16/64/2048	907	3617	3.99×
2/16/16/64/512	603	1369	2.27×
2/16/16/128/2048	2129	4816	2.26×
1/16/16/128/512	607	1150	1.89×
2/32/4/128/2048	1590	3010	1.89×
1/14/2/128/2048	979	1601	1.64×
2/32/4/128/512	656	990	1.51×
1/16/16/64/512	576	843	1.46×
2/14/2/128/512	694	984	1.42×
1/14/2/128/512	600	757	1.26×
2/14/2/128/2048	970	1194	1.23×
1/32/4/128/512	643	751	1.17×
1/32/4/128/2048	1434	1437	1.00×

Range 1.00×–5.40×, geomean ~2.0×. Biggest wins on long-past MHA (H=128, past=2048) where gather bandwidth dominated. The one 1.00× row is a small-K/V-cache GQA case where gather cost was already low.

Uniform prefill (24 shapes)

Range 1.01×–1.25×, geomean ~1.13×. Highlights (all wins):

Shape (B/nH/nKV/H/T)	this PR (µs)	#31611 (µs)	Speedup
1/32/4/128/128	1225	1530	1.25×
2/14/2/128/128	1111	1385	1.25×
2/16/16/64/128	766	948	1.24×
2/32/4/128/512	9723	12054	1.24×
1/16/16/128/128	856	1051	1.23×
2/16/16/128/1024	17296	21029	1.22×
1/14/2/128/128	718	877	1.22×
2/32/4/128/128	1640	1966	1.20×
1/16/16/64/128	750	891	1.19×
2/16/16/128/128	1294	1491	1.15×

(14 more rows 1.01×–1.15×; full log in tree.)

Short-T shapes gain most from Unpack/Repack skip; long-T shapes are dominated by FA compute time.

Varlen prefill (12 shapes, halving q_lens = `{max_T, max_T/2, max_T/4, …}`)

Range 1.14×–1.73×, geomean ~1.29×.

Shape (B/nH/nKV/H/maxT)	q_lens	this PR (µs)	#31611 (µs)	Speedup
4/16/16/128/512	`{512,256,128,64}`	5987	10375	1.73×
4/14/2/128/512	`{512,256,128,64}`	5312	7815	1.47×
4/32/4/128/512	`{512,256,128,64}`	11165	15427	1.38×
4/16/16/128/1024	`{1024,512,256,128}`	20127	27661	1.37×
2/32/4/128/512	`{512,256}`	8345	10732	1.29×
2/16/16/128/1024	`{1024,512}`	15121	19084	1.26×
4/14/2/128/1024	`{1024,512,256,128}`	17685	21562	1.22×
4/32/4/128/1024	`{1024,512,256,128}`	39571	47468	1.20×
2/14/2/128/1024	`{1024,512}`	13141	15430	1.17×
2/16/16/128/512	`{512,256}`	4528	5235	1.16×
2/32/4/128/1024	`{1024,512}`	28871	33383	1.16×
2/14/2/128/512	`{512,256}`	4067	4645	1.14×

Wins grow with batch size — bigger B means more of the padded-BSNH round-trip gets eliminated (B=4/maxT=512 packs only 46.9% of B·maxT tokens; the padded scratch #31611 allocates is >2× bigger than the actual data).

Tests

onnxruntime/test/contrib_ops/paged_attention_op_test.cc PagedAttention.EndToEnd_* — 12/12 non-CUDA tests pass. Covers MHA, GQA, single/multi-batch, variable past lengths, empty tokens, packed QKV, rotary, mixed prefill+decode, cache aliasing via IO-binding. New: EndToEnd_Prefill_MultiBatch_Varlen_Fused (B=2, token_count=48 with q_lens (32,16), head_size=128, MHA) — regression test for the fused varlen prefill path.
Micro-benchmark harness onnxruntime/test/onnx/microbenchmark/paged_attention.cc — 52 registered shapes (16 decode + 24 uniform prefill + 12 varlen prefill).

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>

## 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.

…9789 (#29826) Add regression test for #29789

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/.

Hariharan Seshadri (hariharans29) and others added 30 commits July 31, 2026 13:50

[WebGPU] Add PagedAttention kernel skeleton (Phase 0)

be19980

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.

[PagedAttention] Move paged_attention_helper.h from cuda/ to cpu/

58deec9

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/.

Merge remote-tracking branch 'origin' into hari/webgpu_paged_attention

3776422

PR feedback 1

221cf64

Lint changes

d15f781

PR feedback

4036b86

PR feedback

0f0203b

Merge remote-tracking branch 'origin/main' into hari/webgpu_paged_att…

0bb9324

…ention # Conflicts: # onnxruntime/test/python/transformers/test_paged_attention.py

WebGPU PagedAttention: adapt to merged schema surface

4b58b5c

Update onnxruntime/contrib_ops/webgpu/bert/paged_attention.h

610194d

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Move PagedAttention correctness tests to contrib ops

37b4d51

Merge remote-tracking branch 'origin/hari/webgpu_paged_attention' int…

9868a63

…o hari/webgpu_paged_attention

test: avoid CUDA aliasing failure in shared paged attention OpTester

354e04a

test: add IOBinding coverage for PagedAttention cache aliasing

789126b

Update onnxruntime/test/contrib_ops/paged_attention_op_test.cc

348b544

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update onnxruntime/test/contrib_ops/paged_attention_op_test.cc

bad5b81

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Update onnxruntime/test/contrib_ops/paged_attention_op_test.cc

f3c16b8

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Lint

e915692

Merge main

9e17b04

More tests

28dbdd3

Add regression test ConstantFoldingCopiesAliasedTensorBuffer for PR 2…

3964794

…9789 (#29826) Add regression test for #29789

[PagedAttention] Move paged_attention_helper.h from cuda/ to cpu/

d3f382b

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/.

Navigation Menu

[WebGPU] Optimized PagedAttention implementation (2/n) - #31727

[WebGPU] Optimized PagedAttention implementation (2/n)#31727
Hariharan Seshadri (hariharans29) merged 50 commits into
mainfrom
hari/webgpu_paged_attention_2

Hariharan Seshadri (hariharans29) commented Aug 7, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

Hariharan Seshadri (hariharans29) commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

[WebGPU] PagedAttention: direct paged decode, fused paged prefill, Unpack/Repack skip (Phase 2 partial)

What's shipped

Dispatch-count reduction

Correctness invariants

Performance

Decode (16 shapes)

Uniform prefill (24 shapes)

Varlen prefill (12 shapes, halving q_lens = {max_T, max_T/2, max_T/4, …})

Tests

Related

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Hariharan Seshadri (hariharans29) commented Aug 7, 2026 •

edited

Loading

Varlen prefill (12 shapes, halving q_lens = `{max_T, max_T/2, max_T/4, …}`)