[CUDA] Implement SparseAttentionIndexer operator - #32526
kunal-vaishnavi with Copilot wants to merge 17 commits into
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
There was a problem hiding this comment.
🟡 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_budgetandindex_topkalso remain uncheckedint64values even though capacity/top-k fields are later narrowed toint. For example,index_topk = 2^32passes 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[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
1c15246 to
f330287
Compare
kunal-vaishnavi
changed the title
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…_infer.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
c24495d to
ea75f95
Compare
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Unchecked shape arithmetic, precision mismatches, and unbounded CUDA shared-memory launches leave correctness failures unresolved.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (1)
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu:695
- The CSA path has the same unbounded shared-memory requirement:
CsaCompressKernelrequests4 * head_size + 512bytes, while runtime validation allowshead_sizeup toINT_MAX/2. For example,head_size=16384requests 66,048 bytes and can fail the launch although the input satisfies the public contract. Tile the kernel or validate/opt into the device's supported dynamic shared-memory size before dispatch.
CsaCompressKernel<T><<<compress_blocks, kThreads, value_bytes + kThreads * sizeof(float), stream>>>(
- Files reviewed: 15/15 changed files
- Comments generated: 8
- Review effort level: Balanced
| if (params.max_block_count > 0) { | ||
| const int64_t block_work = rows * params.max_block_count; | ||
| const int score_blocks = static_cast<int>(std::min<int64_t>(block_work, kMaxGridDimX)); | ||
| QsaBlockScoreKernel<T><<<score_blocks, kThreads, 2 * value_bytes + kThreads * sizeof(float), stream>>>( |
| const size_t expected_outputs = is_qsa ? sai::kQsaOutputCount : sai::kCsaOutputCount; | ||
| if (ctx.getNumOutputs() != expected_outputs) { | ||
| fail_shape_inference("SparseAttentionIndexer: policy_mode '", policy_mode, "' requires exactly ", | ||
| expected_outputs, " declared outputs, got ", ctx.getNumOutputs()); |
| plan.overlap_length = past_buffer_length >= compress_ratio ? compress_ratio : 0; | ||
| plan.leftover_length = past_buffer_length - plan.overlap_length; | ||
|
|
||
| const int64_t pending = plan.leftover_length + sequence_length; |
| const float* query_head = query_rotated + (row * params.num_heads + head) * params.head_size; | ||
| float dot = 0.0f; | ||
| for (int d = 0; d < params.head_size; ++d) { | ||
| dot += query_head[d] * to_float<T>(present_compressed_key[key_base + d]); |
| auto* total_dim = present_shape.add_dim(); | ||
| if (past_key_shape->dim(1).has_dim_value() && sequence_dim.has_dim_value()) { | ||
| total_dim->set_dim_value(past_key_shape->dim(1).dim_value() + sequence_dim.dim_value()); | ||
| } |
| } else if (plan_known && past_compressed_shape->dim(1).has_dim_value()) { | ||
| entry_dim->set_dim_value(past_compressed_shape->dim(1).dim_value() + plan.new_window_count); | ||
| } |
|
|
||
| | Stage | Kernel | Parallelism | | ||
| |---|---|---| | ||
| | 1 | `ConcatPastKeyKernel` | element | |
|
|
||
| ## 9. Validation Rules | ||
|
|
||
| Shape inference and the kernel both reject: |
| if (params.max_block_count > 0) { | ||
| const int64_t block_work = rows * params.max_block_count; | ||
| const int score_blocks = static_cast<int>(std::min<int64_t>(block_work, kMaxGridDimX)); | ||
| QsaBlockScoreKernel<T><<<score_blocks, kThreads, 3 * value_bytes + kThreads * sizeof(float), stream>>>( |
There was a problem hiding this comment.
This launch allocates 3 * head_size * sizeof(float) + kThreads * sizeof(float) bytes of dynamic shared memory without checking the device limit or opting the kernel into extended dynamic shared memory.
For example, head_size = 8192 requests 98,816 bytes, which exceeds the usual 48 KiB default launch limit even though this head size passes the current validation. The kernel can therefore fail with an invalid-configuration/resource error for an otherwise valid node.
Could this be tiled so shared-memory usage does not scale with the complete head dimension? Alternatively, please query the device limit, opt the kernel into extended shared memory when supported, and reject unsupported head sizes before launch.
| const int64_t out_base = | ||
| (static_cast<int64_t>(batch) * params.compressed_cache_capacity + entry) * params.head_size; | ||
| for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) { | ||
| present_compressed_key[out_base + d] = from_float<T>(TrailingRope<T>( |
There was a problem hiding this comment.
This rounds the newly generated compressed key to T when storing it in present_compressed_key. The scoring kernel then reloads the rounded FP16/BF16 value at line 587, whereas the documented/reference semantics retain the generated key in FP32 through scoring.
For near-tied candidates, this intermediate rounding can tie or reverse their scores and change selected_indices. Could newly generated keys be retained in an FP32 workspace for scoring during this invocation while still writing the rounded T value to the persistent cache? An adversarial FP16/BF16 near-tie selection test would also help cover this.
| const bool* mask_row = mask + row * params.total_sequence_length; | ||
| int32_t* out_row = visible_indices + row * params.total_sequence_length; | ||
| int32_t offset = 0; | ||
| for (int base = 0; base < params.total_sequence_length; base += blockDim.x) { |
There was a problem hiding this comment.
total_sequence_length is allowed to equal INT_MAX, but this loop uses a signed int counter incremented by blockDim.x. Near the upper bound, base += blockDim.x overflows to a negative value, after which the loop can continue and access mask_row or out_row with negative indices.
Could base and position use int64_t, narrowing only when the value is proven to be in range? The output-initialization loops at lines 372 and 604 have the same issue when capacity == INT_MAX.
| } | ||
| } | ||
| const size_t expected_outputs = is_qsa ? sai::kQsaOutputCount : sai::kCsaOutputCount; | ||
| if (ctx.getNumOutputs() != expected_outputs) { |
There was a problem hiding this comment.
Checking only getNumOutputs() does not ensure that the mandatory CSA outputs are actually present. An ONNX node can have three output slots while using an empty name for output 2, for example ["selected_indices", "present_key", ""]. That passes this check but fails later at runtime when present_proj_buffer is requested.
Could shape inference validate that output 2 is present rather than checking only the number of slots? Please also add a graph-resolution test with an empty third output slot.
| } else { | ||
| auto* total_dim = present_shape.add_dim(); | ||
| if (past_key_shape->dim(1).has_dim_value() && sequence_dim.has_dim_value()) { | ||
| total_dim->set_dim_value(past_key_shape->dim(1).dim_value() + sequence_dim.dim_value()); |
There was a problem hiding this comment.
This concrete-dimension addition is unchecked. Native shape inference accepts int64 ONNX dimensions and runs before the CUDA runtime's INT_MAX validation, so a large past dimension plus sequence_length can overflow int64_t during graph resolution.
Could this use a checked-add helper or enforce the CUDA-supported dimension range before calculating the output dimension? Similar unchecked additions occur in TryComputeCsaWindowPlan and when adding plan.new_window_count around line 2107.
# Conflicts: # docs/ContribOperators.md
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The symbolic inferrer uses an obsolete schema, and integer-overflow and stale-documentation issues remain.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 3
Open (15)
getNumOutputs()counts output slots, including an empty optional placeholder. A CSA node with…head_size, but this launch requests8 * head_size + 512…int64_t. Withcompress_ratio=1,…int64_t. For example, a valid metadata shape with…leftover_length + sequence_lengthis not range-checked. Shape inference can call this helper with…num_heads,…ConcatPastKeyKernel, but no such kernel exists in the implementation. QSA…
| __device__ __forceinline__ int64_t CausalThreshold(int64_t position, int compress_ratio) { | ||
| return position < 0 ? 0 : position / compress_ratio + (position % compress_ratio == compress_ratio - 1); | ||
| } |
| size_t GetQsaWorkspaceFloatCount(const SparseAttentionIndexerParams& params) { | ||
| const size_t rows = static_cast<size_t>(params.batch_size) * params.sequence_length; | ||
| return rows * params.num_heads * params.head_size + rows * std::max(params.max_block_count, 1); |
| return self._get_sympy_shape(node, index) | ||
|
|
||
| if policy_mode == "qsa": | ||
| past_key_shape = past_shape(6) |
| const int64_t threshold = | ||
| position < 0 ? 0 | ||
| : position / problem.compress_ratio + | ||
| (position % problem.compress_ratio == problem.compress_ratio - 1); |
| def test_sparse_attention_indexer_qsa(self): | ||
| node = helper.make_node( | ||
| "SparseAttentionIndexer", | ||
| ["query", "key", "key_norm_weight", "cos_cache", "sin_cache", "mask", "past_key"], |
| #### Inputs (7 - 13) | ||
|
|
||
| <dl> | ||
| <dt><tt>query</tt> : T</dt> | ||
| <dd>Indexer queries with shape (batch_size, sequence_length, num_heads, head_size), already normalized but not yet rotated.</dd> |
| |SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**<br> *out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| | ||
| |SkipSimplifiedLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**<br> *out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| | ||
| |SparseAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* block_row_indices:**M**<br> *in* block_col_indices:**M**<br> *in* total_sequence_length:**M**<br> *in* key_total_sequence_lengths:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float16)| | ||
| |SparseAttentionIndexer|*in* query:**T**<br> *in* key:**T**<br> *in* key_norm_weight:**T**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* mask:**TB**<br> *in* past_key:**T**<br> *in* gate:**T**<br> *in* position_bias:**T**<br> *in* head_weights:**T**<br> *in* position_ids:**I**<br> *in* past_sequence_length:**M**<br> *in* past_proj_buffer:**T**<br> *out* selected_indices:**M**<br> *out* present_key:**T**<br> *out* present_proj_buffer:**T**|1+|**I** = tensor(int64)<br/> **M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **TB** = tensor(bool)| |
Description
Adds
com.microsoft.SparseAttentionIndexerwith CUDA implementations for QSA token selection and CSA compressed-block selection.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.