[Copilot]

Description

Adds com.microsoft.SparsePagedAttention, a model-neutral CUDA executor for externally selected attention indices over ORT paged KV caches.

• Attention modes
  • Selected main-cache attention for Qwen4-Exp QSA.
  • Joint local-main and selected auxiliary attention for DeepSeek V4 CSA.
  • Single online softmax across local, selected, and sink logits.
• Cache integration
  • Reuses PagedAttention writes, block tables, slot mapping, cache aliasing, QK-Norm, RoPE, and INT8 quantization.
  • Supports contiguous FP16/BF16 auxiliary caches and shared auxiliary K=V.
• CUDA execution
  • Resolves request-local indices directly on device.
  • Uses fixed-capacity launches and device-resident counts for CUDA Graph compatibility.
  • Guards invalid, unmapped, out-of-range, and non-causal indices.
• Integration
  • Adds schema and symbolic shape inference.
  • Adds CUDA correctness tests, operator documentation, and support-matrix entries.

QSA: attention_mode="selected_only", selected_kv_source="main"
CSA: attention_mode="local_plus_selected", selected_kv_source="auxiliary"

Motivation and Context

Qwen4-Exp QSA and DeepSeek V4 CSA require sparse attention over different KV lifetimes and logical index spaces. This operator separates externally produced selection metadata from paged-cache writes and attention execution without introducing model-specific policy modes.

[azure-pipelines]

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

[Copilot]

🟡 Changes recommended

A validation failure can occur after mutating the KV cache, and central sparse, quantized, and joint-softmax behaviors lack effective coverage.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds CUDA SparsePagedAttention for externally selected indices over paged or auxiliary KV caches.

Changes:

• Defines and registers the new contrib operator and symbolic shape inference.
• Implements CUDA cache preparation, sparse attention, quantization, and auxiliary-cache support.
• Adds CUDA tests and operator documentation.

File summaries

File	Description
onnxruntime/test/contrib_ops/sparse_paged_attention_op_test.cc	Adds basic CUDA correctness tests.
onnxruntime/python/tools/symbolic_shape_infer.py	Registers symbolic inference.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers the schema.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Defines the operator contract.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers CUDA kernels.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.h	Declares the CUDA kernel class.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention.cc	Validates inputs and dispatches execution.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.h	Declares sparse implementation APIs.
onnxruntime/contrib_ops/cuda/bert/sparse_paged_attention_impl.cu	Implements sparse attention.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h	Exposes the shared cache prologue.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu	Implements and instantiates that prologue.
docs/OperatorKernels.md	Documents CUDA kernel support.
docs/ContribOperators.md	Documents the generated schema.
docs/contrib_ops/cuda/sparse_paged_attention.md	Adds usage and support details.

Review details

• Files reviewed: 14/14 changed files
• Comments generated: 5
• Review effort level: Balanced

---

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

[azure-pipelines]

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

[Copilot]

Copilot review overview

🟡 Changes recommended

Cache-read correctness and launch/input validation defects can produce incorrect attention, mutate caches before failure, or crash the host.

Review effort: Balanced
Findings: 3

Open (3)

• Reject zero head size before workspace and launch calculations · New
• Handle suppressed writes as cache reads and prevent capacity overflow · New
• Prevent oversized CUDA grid Y dimension during sparse prefill · New

Resolved since last review (4)

• physical_block * block_size is evaluated as int before the result is used in the 64-bit cache…
• This block-table offset is also formed in 32-bit arithmetic. A valid table with more than INT_MAX…
• The flattened selected-index offset is computed as int. A valid large tensor whose `token_id *…
• CheckInputs currently derives head_size with truncating integer division and does not verify…