[Copilot]

Description

Adds com.microsoft.DynamicSparseAttention v1 as a model-neutral CUDA executor for externally selected attention indices.

• Execution
  • selected_only over the contiguous main KV cache
  • local_plus_selected over local main KV, auxiliary KV, and an optional sink
  • One numerically stable FP32 softmax across all participating entries
• Attention preparation
  • Separate or packed QKV
  • GQA/MQA head mapping
  • QK RMSNorm and partial RoPE
• Caching
  • Contiguous main KV append
  • Past/present buffer aliasing
  • Read-only auxiliary KV with optional shared K/V storage
• Validation
  • Variable selected counts and -1 padding
  • Duplicate, bounds, and causal-index checks
  • Mode/source and tensor-layout constraints
• Integration
  • Contrib schema and shape inference
  • CUDA registration and symbolic shape inference
  • Operator documentation and deterministic CUDA coverage

Qwen4-Exp: selected_only + main
DeepSeek V4: local_plus_selected + auxiliary

Motivation and Context

Qwen4-Exp QSA and DeepSeek V4 CSA require a common attention executor that consumes selector-produced indices without owning model-specific scoring, TopK selection, or compression policy. This separates selection policy from contiguous-cache attention while preserving DeepSeek’s required joint local/selected/sink normalization.

[azure-pipelines]

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

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

🟡 Changes recommended

Logical-length validation and CUDA graph capture currently permit invalid metadata to produce incorrect attention results.

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

Pull request overview

Adds a CUDA-backed com.microsoft.DynamicSparseAttention operator for externally selected sparse attention over main and auxiliary KV caches.

Changes:

• Adds schema, validation, shape inference, and documentation.
• Implements CUDA preparation, cache append, validation, and attention kernels.
• Adds CUDA tests covering data types, modes, cache aliasing, and invalid inputs.

File summaries

File	Description
onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc	Adds CUDA operator tests.
onnxruntime/python/tools/symbolic_shape_infer.py	Adds symbolic shape inference.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers the schema.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Defines schema and shape inference.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers CUDA kernels.
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h	Declares the CUDA kernel.
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc	Implements operator dispatch and setup.
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h	Declares CUDA implementation interfaces.
onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu	Implements validation and attention kernels.
onnxruntime/contrib_ops/cpu/bert/dynamic_sparse_attention_helper.h	Adds shared input validation.
docs/OperatorKernels.md	Documents CUDA kernel support.
docs/ContribOperators.md	Documents the public operator schema.

Review details

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

---

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

[Copilot]

🟡 Changes recommended

Moderate unresolved contract, shape-inference, and test-coverage issues remain.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (8)

onnxruntime/contrib_ops/cpu/bert/dynamic_sparse_attention_helper.h:131

• The newly added no-past path is not exercised by the CUDA runtime tests: RunDynamicSparseAttentionCase requires past_key and past_value to have cache-sized data, and every helper-based case supplies them. The symbolic-shape test does not execute cache initialization, so regressions in the past_key == nullptr validation/capacity rules and zero-initialized cache append would go undetected; add a runtime case with both past inputs omitted and total_sequence_length == sequence_length.

  ORT_RETURN_IF_NOT(past_key != nullptr || total_length == sequence_length,
                    "DynamicSparseAttention: total_sequence_length must equal sequence length when no past cache "
                    "is provided.");
  ORT_RETURN_IF_NOT(past_key == nullptr || total_length <= past_cache_capacity,
                    "DynamicSparseAttention: total_sequence_length must not exceed the past cache capacity.");

onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc:82

• is_causal is exposed as a normal integer attribute with the description “Whether ...” and no schema restriction, but this constructor rejects every value except 1. A model exporter cannot use the advertised non-causal setting, and the schema/runtime contract is therefore misleading; either propagate is_causal=0 through validation and the kernel or make the attribute explicitly fixed to causal mode and document that limitation.

  const int64_t is_causal = info.GetAttrOrDefault<int64_t>("is_causal", 1);
  ORT_ENFORCE(is_causal == 1, "DynamicSparseAttention only supports is_causal=1.");

onnxruntime/core/graph/contrib_ops/bert_defs.cc:545

• The schema documentation says selected_only attends only to selected entries and documents the sink only under local_plus_selected, but the implementation also includes head_sink in selected_only (and the added SelectedOnlyExplicitSink test relies on that behavior). Please document the sink behavior for both modes, including in the regenerated operator docs, or reject sink inputs for selected_only so the published contract matches execution.

`attention_mode="selected_only"` attends only to selected entries. `attention_mode="local_plus_selected"` jointly
normalizes causally valid main-cache entries in `local_window_size`, selected auxiliary entries, and an optional
per-query-head sink. The sink contributes to the softmax denominator but has no value vector. A row with no entries and
no sink produces zero output.

onnxruntime/core/graph/contrib_ops/bert_defs.cc:491

• These guards are specifically needed because present_key and present_value are optional, but every new shape-inference test declares all three outputs; the manual Graph::Resolve test also provides only element types, not input shapes. Add a model-level regression with shaped inputs and one/two outputs to exercise the optional-output path and protect this bounds check.

  for (int output_index = 1; output_index <= 2; ++output_index) {
    const int input_index = output_index + 2;
    if (ctx.getNumOutputs() <= static_cast<size_t>(output_index) || !ctx.hasOutput(output_index)) {
      continue;

onnxruntime/core/graph/contrib_ops/bert_defs.cc:508

• The separate-QKV path floor-divides the query width by num_heads without checking divisibility. A model with D=10 and num_heads=3 therefore gets a cache shape with head_size=3 during Graph::Resolve, even though CheckInputs rejects the same model at execution; add the analogous divisibility check before deriving head_size.

        const bool is_packed = ctx.getInputType(2) == nullptr;
        const int64_t head_size = query_dims[2].dim_value() /
                                  (is_packed ? num_heads + 2 * kv_num_heads : num_heads);

onnxruntime/python/tools/symbolic_shape_infer.py:2637

• This handler also uses floor division for packed QKV widths, so SymbolicShapeInference rewrites an invalid input width to a plausible but incorrect output width instead of checking the divisor. The cache derivation below has the same issue for separate QKV; guard both derivations, following the adjacent GQA handler's non-truncating behavior.

            if node.input[1] == "" and node.input[2] == "" and isinstance(output_shape[2], int):
                num_heads = get_attribute(node, "num_heads")
                kv_num_heads = get_attribute(node, "kv_num_heads")
                head_size = output_shape[2] // (num_heads + 2 * kv_num_heads)
                output_shape[2] = num_heads * head_size

onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc:303

• The only successful auxiliary-attention case sets auxiliary_kv_shared=1 and omits auxiliary_value, so it never verifies the ordinary path that reads a distinct auxiliary value tensor. Add a case with different auxiliary key/value contents and auxiliary_kv_shared=0 to catch accidentally using the key as the value.

  DynamicSparseAttentionCase c;
  c.attention_mode = "local_plus_selected";
  c.selected_kv_source = "auxiliary";
  c.local_window_size = 1;
  c.auxiliary_kv_shared = 1;
  c.auxiliary_sequence_length = 2;

onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc:396

• The only rotary test explicitly sets rotary_interleaved=1, leaving the half-split branch in ApplyRotary untested. Add a case (or parameterize this one) with rotary_interleaved=0 and expected values that distinguish the two layouts.

  c.scale = 0.1f;
  c.do_rotary = 1;
  c.rotary_interleaved = 1;
  c.rotary_offset = 8;

• Files reviewed: 14/14 changed files
• Comments generated: 1
• Review effort level: Lite (auto)

Note

Copilot is running an experiment and ran this review at Lite.

[Copilot]

Copilot review overview

🟡 Changes recommended

Explicit zero scale is misinterpreted, and unchecked schema arithmetic can overflow during graph resolution.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1 · 3

Open (5)

• Validate Head Counts Before Packed-Head Arithmetic · New
• Preserve Explicit Zero Scale Instead of Applying Default · New
• Test Smooth Softmax Without an Explicit Head Sink · New
• Document Rotary Offset Multiple-of-8 Requirement · New
• Document Packed QKV Row-Wise Layout · New

Resolved since last review (1)

• The schema describes is_causal as a switch that can enable or disable causal visibility, but the…

[Copilot]

Copilot review overview

🔵 Needs a closer look

Unresolved correctness issues and the new CUDA cache-aliasing path require fixes and maintainer GPU verification.

Review effort: Balanced
Findings: 5

Open (5)

• Pairwise duplicate validation scales quadratically · New
• CUDA tests lack batch and KV-head coverage for GQA · New
• Fallback incorrectly infers shape for supplied unshaped past cache · New
• Fallback overwrites capacity for present unshaped past caches · New
• Invalid CUDA attributes abort under ORT_NO_EXCEPTIONS · New

Resolved since last review (5)

• Validate Head Counts Before Packed-Head Arithmetic
• Preserve Explicit Zero Scale Instead of Applying Default
• Document Packed QKV Row-Wise Layout
• Document Rotary Offset Multiple-of-8 Requirement
• Test Smooth Softmax Without an Explicit Head Sink