[WebGPU] Add PagedAttention metadata and GPT-OSS support

[Tianlei Wu (tianleiwu)]

Description

Use the optional CPU attention_metadata input in WebGPU PagedAttention to obtain replay-wide query and KV bounds without downloading device sequence metadata. A new PagedAttentionPrepareMetadata WebGPU program derives exact per-request seqlen_k and seqlens_q arrays on device for masking.

This PR also adds the local-window and learned head-sink semantics needed by GPT-OSS. Unsupported optimized combinations route through the generic gather-then-flash path so output remains correct while direct paged kernel support can be added separately.

Models that do not declare attention_metadata retain the existing packed metadata readback path for compatibility.

Motivation

The previous WebGPU path performed a blocking GPU-to-CPU metadata copy in every PagedAttention node on every model run. In a 24-layer GPT-OSS model this creates 24 queue flushes per forward, serializes otherwise device-resident work, and prevents graph capture. GenAI already computes stable replay bounds on the host.

GPT-OSS also requires sliding-window attention and learned attention sinks. Supporting those semantics in the shared FlashAttention path enables correct WebGPU execution without waiting for every direct paged optimization.

Changes

• Keep input 16 in CPU memory and validate its two- or three-value bounds consistently with CUDA.
• Generate exact per-request query and KV lengths in a small GPU dispatch.
• Skip RunPackMetadata, the blocking device-to-host copy, and the two host-to-device length copies when metadata is present.
• Preserve the legacy metadata fallback for older model exports.
• Apply local-window masking and the learned head-sink term in generic FlashAttention.
• Keep head-sink-only short-query decode on direct paged split-reduce; route local-window cases and head-sink prefill through gather-then-flash.
• Reject is_causal=0 explicitly until WebGPU implements non-causal PagedAttention.
• Add WebGPU parity coverage for metadata bounds, observable ragged per-request lengths, empty requests, local-window masking, and head sinks.

Validation

Previously completed on the WebGPU branch:

• Regenerated WGSL template headers and compiled paged_attention.cc with WebGPU provider flags.
• Passed WebGPU PagedAttention parity with attention_metadata=[1, 64, 1] on Vulkan adapter 7.
• Passed legacy no-metadata fallback parity on the same provider.
• Passed current-head WebGPU CI before the review-fix commit.

Review-fix validation:

• py_compile passed for test_paged_attention.py.
• Ruff lint and format checks passed.
• git diff --check passed.
• The focused metadata test collects only with a built native ONNX Runtime package; this checkout does not have a WebGPU-enabled Python build, so the new [1, 2] numeric parity case is left to WebGPU CI.

A companion ONNX Runtime GenAI PR enables FP16 WebGPU paged exports and preserves the metadata input.

Companion: microsoft/onnxruntime-genai#2470

[Copilot]

Pull request overview

Adds a WebGPU metadata fast path that avoids blocking GPU-to-CPU readbacks in PagedAttention.

Changes:

• Uses CPU metadata bounds for allocation and dispatch.
• Derives exact sequence lengths with a GPU shader.
• Preserves legacy fallback and adds parity coverage.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
test_paged_attention.py	Adds metadata parity coverage.
paged_attention.h	Declares metadata preparation program.
paged_attention.cc	Implements metadata fast path and fallback.
paged_attention_prepare_metadata.wgsl.template	Computes per-request lengths on GPU.
webgpu_paged_attention.md	Documents the new path.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

Follow-up #32320 adds WebGPU local-window and attention-sink support on top of this metadata fast path.

[Jiajia Qin (qjia7)]

Review frame

• Problem/feature validity: Validated. Without input 16, every WebGPU PagedAttention node packs the device-resident cumulative query lengths and past lengths, copies them to CPU, and then copies derived per-request lengths back to GPU. For a multi-layer decoder this forces one queue synchronization per layer and prevents WebGPU graph capture. The merged attention_metadata schema and microsoft/onnxruntime-genai#2470 provide replay-wide CPU bounds specifically to avoid that synchronization.
• Risk/scope: Deep. Current head b2a94b06ea151a4e793228205308495cb03c1717 rebases the metadata and local-window/head-sink work from previously reviewed head e97807cd6b01157282200eddad931b39dd7e1bdb onto newer main, then adds a focused ragged metadata test. The behavioral surface remains dispatch selection, scratch sizing, device-derived causal masks, local-window masking, learned sink softmax, compatibility fallback, and graph replay.
• Direction gate: Pass. Host-owned replay bounds plus device-owned exact mask lengths remain the appropriate boundary. Applying local-window and sink semantics in generic FlashAttention while routing unsupported direct-paged combinations through gather-then-flash remains a sound correctness-first design. The rebase also explicitly rejects the new mainline is_causal=0 contract in WebGPU rather than silently returning causal output.

Confirmed findings

T1 (partially resolved from the prior review): Make the ragged metadata test observe both derived query lengths

onnxruntime/test/python/transformers/test_paged_attention.py:1656-1680

The new case correctly enables attention_metadata, supplies new_seqlens=[2, 0], and reaches PagedAttentionPrepareMetadataProgram, including the seqlen_k=-1 sentinel for the empty request. It therefore adds useful zero-request safety coverage.

It still does not discriminate the second seqlens_q value. A request with q_len=0 owns no packed query or output rows. Although FlashAttention runs over padded rows for that batch, PagedAttentionRepackOutputProgram copies no row from it into the observable output. Consequently, a faulty metadata shader that wrote q_len=2 for both requests could still produce exactly the expected packed output.

Please add an observable ragged case with the shorter request first, for example new_seqlens=[1, 2] with the existing [2, 2, 0] metadata bounds. If the shader accidentally reused batch 0's length for every request, batch 1 would get seqlen_k=0 instead of 1; its second query row would then omit the second key and fail parity. The existing [2,0] case can remain for the empty-request sentinel.

Q1 (unresolved from the prior review): Update the design constraints to match the enabled features

docs/design/webgpu_paged_attention.md:36,50

The document describes local-window and head-sink support as complete at lines 380-388, but its non-goals still group head-sink with deferred QK-Norm and its locked-constraint table still says every local_window_size != -1 is rejected. Those statements contradict PagedAttention::ComputeInternal, which accepts both features and routes them through the implemented paths. Please update the non-goals and constraint table so the provider support contract is unambiguous.

Q2 (non-blocking follow-up): Support local-window masking in split-reduce

onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc:816-834, onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc:1005-1012

Both routing layers require !has_local_window for the short-query fast path. Consequently, setting local_window_size > 0 disables direct paged decode and split-reduce even for one-token decoding. The generic FlashAttention fallback is correct, but gathering dense K/V and processing the full history may not be performant for decoding.

As a non-blocking follow-up, please add local-window masking to the split-reduce path so local-window decoding can retain the optimized paged decode path.

Clarifications

None.

Test coverage

The legacy no-metadata matrix covers the compatibility readback. The uniform metadata case covers CPU input placement and the shape-(3) contract. The newly added [2,0] case covers shader dispatch and safe integration of the empty-request sentinel, but T1 remains because its only nonempty request is batch 0, so reusing batch 0's length for every request would still pass.

The local-window/head-sink tests cover combined prefill/decode, short-history window clamping, sink-only prefill/decode, generic gather-then-flash, and direct paged split-reduce routing. They still do not compose those features with attention_metadata, but the metadata and feature control flows are independently covered and the implementation trace found no new correctness defect from the rebase. Q2 records non-blocking follow-up work for an optimized local-window decode path.

Current-head CI is green, including Linux WebGPU builds, macOS WebGPU build-and-test jobs, WGSL template tests, Python formatting, and optional C++ lint. The checked macOS logs do not show the Python TestPagedAttentionWebGpu case running, so CI success is build/backend coverage rather than evidence that the new Python ragged case executed numerically.

Verdict

The metadata-readback problem and the absorbed GPT-OSS feature requirement remain valid, and the design direction remains appropriate after the rebase. T1 is the remaining merge-blocking test gap because the new metadata shader still lacks an observable per-batch ragged case. There are no clarification requests. Q1 is non-blocking documentation cleanup, and Q2 is a non-blocking follow-up to support local windows in split-reduce for performant decoding. Updating the stale PR title/body to describe the local-window/head-sink scope is also cleanup only.

[Tianlei Wu (tianleiwu)]

Addressed the review feedback in fa8ae86:

• T1: kept the [2, 0] empty-request sentinel case and added an observable [1, 2] ragged case, so both device-derived seqlens_q values affect output.
• Q1: updated the design non-goals and constraint table to describe the implemented head-sink support and the local-window gather-then-flash fallback.
• Q2: leaving local-window masking in split-reduce as a non-blocking follow-up. The current fallback is correct; extending the optimized path requires split-reduce WGSL changes plus path-specific correctness and performance coverage.
• Updated the PR title and description to include the GPT-OSS local-window/head-sink scope.

Local checks passed for Python compilation, Ruff lint/format, and git diff --check. Numeric execution of the new WebGPU parity case requires the WebGPU CI environment.

[Jiajia Qin (qjia7)]

Review frame

• Problem/feature validity: Validated. Without input 16, every WebGPU PagedAttention node downloads device-resident sequence metadata and uploads the derived lengths again, forcing a queue synchronization per layer and preventing graph capture. GPT-OSS also requires local-window masking and learned attention sinks.
• Risk/scope: Deep for the full PR because it changes WebGPU attention dispatch, masking, scratch sizing, compatibility fallback, and replay behavior. The incremental range from previously reviewed head b2a94b06ea151a4e793228205308495cb03c1717 to current head fa8ae86e74542d91cb06c0e95a2ca7a1f78b2098 is narrowly limited to one additional metadata parity case and corrections to the design document.
• Direction gate: Pass. Host replay-wide bounds select dispatch and size work, while exact per-request lengths remain device-derived for masking. Unsupported optimized local-window combinations correctly fall back to gather-then-flash, and the performance debt is explicitly documented.

Prior review disposition

• T1 resolved: The added [1, 2] case uses metadata bounds [2, 2, 0] and gives both requests observable output. It specifically distinguishes per-batch metadata derivation: if batch 0's query length (1) were reused for batch 1, batch 1 would get seqlen_k=0 instead of 1, so its second query row would omit the second key and fail parity. The existing [2, 0] case continues to cover the empty-request sentinel.
• Q1 resolved: The non-goals now defer only QK-Norm, identify both implemented head-sink routes, and describe local_window_size > 0 as supported through gather-then-flash while direct paged paths remain disabled.
• Q2 remains non-blocking follow-up: Local-window masking is still not implemented in split-reduce, so local-window decode uses the correct but potentially slower gather-then-flash path. This optimization was explicitly deferred and the current commit does not regress it.

Confirmed findings

Q3 (non-blocking performance): Start fused FlashAttention at the earliest active window tile

onnxruntime/contrib_ops/webgpu/bert/flash_attention.wgsl.template:380,439

Both fused FlashAttention implementations still start k_start at zero when has_local_window is enabled. Keys before local_window_start are masked correctly, but the shader first loads their K/V data and computes their QK scores. Consequently, long-context local-window attention retains computation proportional to the full KV history even though only the active window can affect output.

Could we start each key loop at the aligned earliest window boundary needed by the query workgroup? For causal attention, the workgroup-uniform bound is based on its first query row:

query_tile_start = query_tile_index * workgroup_size
earliest_causal_end = past_sequence_length + query_tile_start + 1
earliest_window_start = max(0, earliest_causal_end - local_window_size)

Round earliest_window_start down by the loop's actual step (max_k_step for the shared-memory path and capped_sg_size for the subgroup path), and retain the existing per-query is_key_visible() checks for the partially overlapping first tile and later rows' shifted windows. The lower bound must remain workgroup-uniform because both loops contain workgroupBarrier(). For the non-causal specialization, derive the start from total_sequence_length instead.

This should reduce fused QK/PV traversal to approximately the window size plus one query tile. The full-history PagedAttention gather and split-reduce local-window support can remain separate follow-ups.

Clarifications

None.

Test coverage

The metadata path now has uniform, observable ragged [1, 2], and empty-request [2, 0] coverage. Static control-flow tracing confirms all three cases provide input 16 and therefore execute PagedAttentionPrepareMetadataProgram; the [1, 2] case then reaches short-query paged split-reduce with independently observable batch outputs.

The changed Python file compiles and passes targeted lint. Current WebGPU-relevant CI passes, including Linux WebGPU build, all three macOS WebGPU build-and-test jobs, and both WGSL template jobs. The macOS logs do not show this Python parity test executing numerically, consistent with the PR's stated local limitation.

Three unrelated checks currently fail: Optional Lint cannot build reviewdog's Debian image because the repository Release file is expired; React Native iOS fails in Xcode SDK headers on an unknown dev_t; and CUDA Plugin EP tests segfault in CudaPluginUserStreamGraphTest.CaptureAndReplayOnUserStream. None intersects the two changed files or WebGPU PagedAttention path.

Verdict

The feature remains valid and the design direction remains appropriate. T1 and Q1 are resolved; Q2 remains separate non-blocking split-reduce follow-up work; and Q3 suggests a non-blocking optimization to avoid traversing fully masked KV tiles in fused FlashAttention. There is no correctness blocker or clarification request at current head fa8ae86e74542d91cb06c0e95a2ca7a1f78b2098; the three failing CI jobs appear unrelated and require rerun or owner-side CI handling rather than a PR code change.

[github-actions]

You can commit the suggested changes from lintrunner.