[Copilot]

Description

Adds com.microsoft.SparseAttentionIndexer with CUDA implementations for QSA token selection and CSA compressed-block selection.

• Operator contract
  • Adds schema, shape inference, symbolic shape inference, and CUDA kernel registration.
  • Supports FP32, FP16, and BF16 inputs.
• CUDA execution
  • Implements rotary embedding, key compression, scoring, deterministic top-k selection, and cache updates.
  • Handles batched inputs, long sequences, and strided channel processing.
• Validation and coverage
  • Rejects non-positive attention-head counts during graph resolution and execution.
  • Adds reference-based coverage for QSA, CSA, cache-only steps, batch strides, and multi-tile scans.
• Documentation
  • Documents policy semantics, tensor contracts, cache behavior, and kernel pipelines.

Motivation and Context

Recent sparse-attention decoders require either token-level QSA indexing or stateful CSA compressed-block indexing. This operator provides both policies through one validated CUDA interface while preserving the cache state required by subsequent decoding steps.

[azure-pipelines]

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

[Copilot]

🟡 Changes recommended

Attribute narrowing, zero-value semantics, CUDA edge cases, missing symbolic tests, and incomplete generated documentation remain unresolved.

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

Pull request overview

Adds the CUDA-only com.microsoft.SparseAttentionIndexer contrib operator for QSA token and CSA compressed-block selection.

Changes:

• Defines schema, validation, shape inference, and symbolic inference.
• Implements CUDA kernels and state-cache handling for FP32, FP16, and BF16.
• Adds reference-based tests and operator documentation.

File summaries

File	Description
onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc	Adds shape and CUDA numerical tests.
onnxruntime/python/tools/symbolic_shape_infer.py	Adds symbolic output inference.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers the schema.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Defines the operator contract and shape inference.
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.h	Declares the CUDA kernel.
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc	Validates inputs and dispatches policies.
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.h	Defines launch interfaces and parameters.
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu	Implements CUDA pipelines.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers typed CUDA kernels.
onnxruntime/contrib_ops/cpu/sparse/sparse_attention_indexer_common.h	Shares policy and window-planning utilities.
docs/OperatorKernels.md	Lists kernel availability and types.
docs/contrib_ops/cuda/sparse_attention_indexer.md	Documents semantics and implementation.

Review details

Suppressed comments (1)

onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc:65

• token_budget and index_topk also remain unchecked int64 values even though capacity/top-k fields are later narrowed to int. For example, index_topk = 2^32 passes validation but becomes zero, producing an output shape and selection behavior that disagree with the schema-inferred capacity. Add representability checks before accepting these attributes and mirror them in shape inference.

  const bool has_token_budget = info.GetAttr<int64_t>("token_budget", &token_budget_).IsOK();
  const bool has_index_topk = info.GetAttr<int64_t>("index_topk", &index_topk_).IsOK();

• Files reviewed: 12/12 changed files
• Comments generated: 6
• Review effort level: Balanced

---

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

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Copilot review overview

🟡 Changes recommended

CUDA indexing overflow and shared-memory configuration races remain unresolved.

Review effort: Balanced
Findings: 3

Open (3)

• Reject query widths exceeding INT_MAX · New
• Avoid race from per-call CUDA shared-memory configuration · New
• Prevent signed overflow in selection scan counter · New

Resolved since last review (15)

• Saturate threshold for INT64_MAX with compression ratio one
• getNumOutputs() counts output slots, including an empty optional placeholder. A CSA node with…
• The schema permits arbitrary positive head_size, but this launch requests 8 * head_size + 512…
• Update fixture to the current 14-input schema
• Fix reference threshold overflow regression case
• Update inferrer indices and query head-size source
• Remove unused rotated-query workspace allocation
• The inferred compressed-cache length can overflow signed int64_t. With compress_ratio=1,…
• These shape dimensions can overflow signed int64_t. For example, a valid metadata shape with…
• For newly compressed entries, this load no longer observes the float32 value used by the reference:…
• leftover_length + sequence_length is not range-checked. Shape inference can call this helper with…
• Regenerate stale kernel documentation entry
• Regenerate API documentation for the current schema
• This statement overstates graph-time validation. The inference function checks ranks, num_heads,…
• This pipeline names ConcatPastKeyKernel, but no such kernel exists in the implementation. QSA…

[Copilot]

Copilot review overview

🔵 Needs a closer look

Aliased cache updates and top-k scans beyond 128 candidates need regression coverage.

Review effort: Balanced
Findings: None

Resolved since last review (3)

• Prevent signed overflow in selection scan counter
• Avoid race from per-call CUDA shared-memory configuration
• Reject query widths exceeding INT_MAX

Previously missed (2)

In code that hasn't changed since last review

Add aliased shared-cache tests for copy skipping and cache updates

onnxruntime/​contrib_ops/​cuda/​sparse/​sparse_attention_indexer.cc:33

The shared-cache tests check capacity, but never alias past_key and present_key. They supply past_key as a graph input and leave the output unbound; graph inputs are not reused by the allocation planner (allocation_planner.cc:643–659). Consequently, neither policy tests the copy-skipping branch or updates to the caller's cache. Add CUDA-backed QSA and CSA tests that bind both cache tensors to the same device buffer, assert pointer equality, and compare successive decoding steps with the reference, including preservation of untouched cache rows.

Test candidate selection beyond 128 entries and tie ordering

onnxruntime/​test/​contrib_ops/​sparse_attention_indexer_op_test.cc:1112

The largest QSA numeric case has only 100 complete visible blocks, and CSA cases have at most four entries. With 128 selection threads, none reaches a second iteration of ScanForNext, so later-candidate selection and tie ordering are untested for both policies. Add reference-based cases with more than 128 visible candidates: make entries beyond index 127 win, and use equal scores with enough selected entries to cross that boundary. The suggested fixtures exercise both scenarios while retaining QSA's strided-channel coverage.

[github-actions]

You can commit the suggested changes from lintrunner.