[CUDA][WebGPU] Add packed sparse attention indexer for continuous batching

[Copilot]

Description

Adds com.microsoft.PackedSparseAttentionIndexer version 1 for packed, variable-length continuous batching, with CUDA and WebGPU implementations.

The operator supports both QSA and CSA policies over flattened token inputs and per-request metadata while preserving fixed-capacity state across prefill and decode calls. It is designed to produce selection metadata compatible with SparsePagedAttention without changing the existing dense SparseAttentionIndexer ABI or behavior.

Key changes:

• Add the shared operator schema, fixed generic QSA/CSA state slots, strict policy-specific validation, and fixed output shape inference.
• Add CUDA and WebGPU kernels and provider registrations for FP32, FP16, and BF16 where supported.
• Add device-side packed metadata validation without host synchronization in normal provider paths.
• Add deterministic TopK selection and complete per-request overflow handling that preserves prior state and returns empty selections.
• Support safe state aliasing where provider semantics permit it while avoiding unsafe CSA buffer aliasing.
• Share policy constants and math helpers with the existing sparse attention indexer where practical.
• Add symbolic shape inference support and focused Python coverage.
• Add C++ shape, QSA/CSA numerical, prefill/decode, zero-token request, malformed metadata, deterministic ordering, and overflow tests.
• Add operator and WebGPU implementation documentation and regenerate contrib operator/kernel documentation.

Validation performed:

• onnxruntime_provider_test CPU build completed successfully.
• All 12 PackedSparseAttentionIndexerShapeInferenceTest tests passed.
• Symbolic shape-inference tests passed.
• Repository linting, formatting checks, and secret scanning passed.
• CUDA/WebGPU numerical tests are included but could not be executed locally because those providers and compatible hardware/toolchains are unavailable in the environment.

Motivation and Context

The existing SparseAttentionIndexer contract uses dense [B, S, ...] inputs, dense masking, and growing contiguous state. OgaEngine-compatible paged-attention models instead use flattened packed tokens, cumulative_sequence_lengths, per-request past lengths, dynamic request scheduling, and fixed-address state.

This separate operator bridges that contract mismatch for sparse paged attention while keeping the existing dense operator stable. QSA selections reference request-local positions in the main paged cache; CSA selections reference fixed-capacity auxiliary compressed state suitable for SparsePagedAttention local-plus-selected execution.

Built on the sparse attention indexer work from #32528.

[azure-pipelines]

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

[Copilot]

🟡 Changes recommended

CUDA compile errors and unresolved shape-validation and WebGPU aliasing defects block approval.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds a packed, fixed-capacity sparse-attention indexer for continuous batching and paged-attention models.

Changes:

• Defines the generic QSA/CSA schema and symbolic shape inference.
• Implements and registers CUDA and WebGPU kernels.
• Adds documentation and provider-focused tests.

File summaries

File	Description
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Defines schema and shape inference.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers the schema.
onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h	Defines packed ABI constants.
onnxruntime/contrib_ops/cpu/sparse/sparse_attention_indexer_common.h	Makes shared helpers device-callable.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers CUDA kernels.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc	Implements CUDA host validation.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h	Declares the CUDA kernel.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu	Implements CUDA device kernels.
onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h	Declares CUDA launch interfaces.
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_device_math.cuh	Extracts shared CUDA math.
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu	Adopts shared CUDA helpers.
onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc	Implements WebGPU programs.
onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h	Declares WebGPU programs.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc	Registers the WebGPU kernel.
onnxruntime/python/tools/symbolic_shape_infer.py	Adds symbolic inference dispatch.
onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc	Adds schema and numeric tests.
onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py	Tests symbolic output shapes.
docs/ContribOperators.md	Adds generated schema documentation.
docs/OperatorKernels.md	Lists provider kernel support.
docs/contrib_ops/packed_sparse_attention_indexer.md	Documents the operator contract.
docs/contrib_ops/webgpu/packed_sparse_attention_indexer.md	Documents WebGPU support.

Review details

Suppressed comments (4)

onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu:571

• overflowed is not declared in the CSA kernel, so this is another CUDA compile failure. The guard should use the previously computed rejected flag.

    if (!overflowed) {

onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc:984

• The CSA aliasing path binds both past KV/gate buffers as read-only inputs and their aliased present buffers as read-write outputs. Dawn rejects this mixed usage for one storage buffer, so the advertised state aliasing cannot run. Add per-buffer alias detection and access aliased history only through the corresponding output binding.

                    {past_kv_buffer, ProgramTensorMetadataDependency::Type},
                    {past_gate_buffer, ProgramTensorMetadataDependency::Type}})
        .AddInput({past_state_lengths, ProgramTensorMetadataDependency::Type})
        .AddOutputs({{present_key_state, ProgramTensorMetadataDependency::Type},
                     {present_kv_buffer, ProgramTensorMetadataDependency::Type},

onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc:1089

• Like the QSA case, this “prefill followed by decode” test executes only one call with fabricated prior state. It cannot catch state handoff or aliasing defects between calls; execute prefill and decode sequentially using the first call's outputs.

TEST(PackedSparseAttentionIndexerTest, CsaPrefillThenDecodeIndependentState) {
  CsaPackedProblem problem;
  problem.cumulative_sequence_lengths = {0, 1, 4};
  problem.past_state_lengths = {1, 1, 0, 0};
  RunCsaPackedTest<float>(MakeCsaPackedProblem(std::move(problem)), 1.0e-5f);

onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc:1104

• CSA's advertised WebGPU FP16 path is also untested because this provider-specific test instantiates only float; the generic CsaFloat16 case runs on CUDA. Add a WebGPU FP16 case.

TEST(PackedSparseAttentionIndexerWebGpuTest, CsaFloat) {
  RunCsaPackedTest<float>(MakeCsaPackedProblem(), 1.0e-5f, ProviderKind::WebGpu);
}

• Files reviewed: 21/21 changed files
• Comments generated: 8
• Review effort level: Balanced

---

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