[WebGPU] Add SparsePagedAttention

[Copilot]

Description

Adds WebGPU Execution Provider support for com.microsoft.SparsePagedAttention using the shared CPU/CUDA schema.

• Attention
  • Supports selected main-cache reads.
  • Supports local paged attention plus selected contiguous auxiliary K/V.
  • Computes one FP32 softmax across local, selected, and sink logits.
• Cache and metadata
  • Reuses PagedAttention QKV splitting, RoPE, and cache-update paths.
  • Keeps selection, request metadata, and block-table resolution device-resident.
  • Supports slot mapping, suppressed writes, cache aliasing, and shared auxiliary K/V.
• Safety
  • Validates signed indices before unsigned conversion.
  • Bounds logical positions and shader iteration counts by cache capacity.
  • Rejects unsupported types, layouts, dimensions, quantization, and resource limits without truncation.
• Coverage
  • Adds deterministic tests for selected-main, local-plus-auxiliary, joint softmax, grouped-query batching, cache writes, invalid inputs, and extreme lengths.
  • Documents the supported WebGPU feature matrix and limits.

Motivation and Context

SparsePagedAttention enables QSA-style selected reads and CSA-style local-plus-compressed attention over paged caches. WebGPU support allows these models to preserve the common operator ABI and device-resident scheduling semantics without provider-specific graph rewrites.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🔵 Needs a closer look

It introduces a new attention operator plus substantial new GPU kernels (WebGPU WGSL pipeline and CUDA implementation) that warrant final human review for long-tail correctness/perf/device-limit edge cases.

Pull request overview

This PR introduces com.microsoft.SparsePagedAttention (schema + docs) and adds Execution Provider implementations, with a staged WebGPU pipeline that reuses parts of the existing WebGPU PagedAttention prologue (packed-QKV split, RoPE, and block-table scatter), plus a functional CUDA implementation and unit tests to validate key behaviors (cache writes, selection semantics, and joint softmax).

Changes:

• Added the SparsePagedAttention contrib op schema (Microsoft opset v1) and registered kernels for WebGPU and CUDA.
• Implemented a staged WebGPU kernel (scatter/token-meta/main+aux partials/finalize) with explicit feature guarding and device-limit validation.
• Added deterministic unit tests and updated generated contrib-op/operator-kernel documentation, plus WebGPU/CUDA operator docs.

File summaries

File	Description
onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc	New deterministic tests covering WebGPU (and a couple CUDA reference cases) for selection, cache updates, and sanitization.
onnxruntime/python/tools/symbolic_shape_infer.py	Routes SparsePagedAttention to the existing PagedAttention symbolic inference logic.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Adds SparsePagedAttention schema class to Microsoft opset v1 list.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Defines the SparsePagedAttention v1 schema and reuses PagedAttention type/shape inference.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc	Registers the WebGPU SparsePagedAttention kernel and includes its header.
onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention.h	Declares the staged WebGPU programs and SparsePagedAttention kernel class.
onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention.cc	Implements WebGPU kernel dispatch, validation, stage binding-budget checks, and cache alias/copy behavior.
onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_token_meta.wgsl.template	WGSL stage to compute sanitized per-token metadata on device.
onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_scatter_kv.wgsl.template	WGSL stage to scatter current K/V into cache using slot_mapping.
onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_main.wgsl.template	WGSL main-cache partial attention stage producing FP32 softmax state.
onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_auxiliary.wgsl.template	WGSL auxiliary-cache partial attention stage producing FP32 softmax state.
onnxruntime/contrib_ops/webgpu/bert/sparse_paged_attention_finalize.wgsl.template	WGSL finalize stage merging partial softmax states (and optional sink) into output.
onnxruntime/contrib_ops/webgpu/bert/paged_attention.h	Exposes prologue dispatch helpers for reuse by SparsePagedAttention.
onnxruntime/contrib_ops/webgpu/bert/paged_attention.cc	Makes prologue dispatch helpers externally callable and updates internal call sites.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers CUDA SparsePagedAttention typed kernels.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.h	Declares the CUDA SparsePagedAttention kernel class template.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc	Implements CUDA kernel plumbing/validation and invokes the CUDA implementation.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.h	Declares the CUDA implementation entry point and enums.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu	Implements the functional CUDA kernel and launch code for sparse attention.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h	Adds a reusable PreparePagedAttentionQueryAndCache API used by sparse attention.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu	Implements/instantiates PreparePagedAttentionQueryAndCache.
docs/OperatorKernels.md	Updates generated kernel documentation to include SparsePagedAttention.
docs/ContribOperators.md	Updates generated contrib-op documentation to include SparsePagedAttention schema details.
docs/contrib_ops/webgpu/sparse_paged_attention.md	New WebGPU-specific doc describing staged execution and supported feature matrix/limits.
docs/contrib_ops/cuda/sparse_paged_attention.md	New CUDA-specific doc describing modes and capability matrix.

Review details

• Files reviewed: 25/25 changed files
• Comments generated: 0
• Review effort level: Lite

---

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