From 3e309f4fd3255d79b600da42358fa6ef4999089e Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 10 Sep 2026 02:32:20 +0000
Subject: [PATCH 01/16] Implement SparseAttentionIndexer CUDA operator
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
---
docs/OperatorKernels.md | 1 +
.../cuda/sparse_attention_indexer.md | 364 +++++++
.../sparse/sparse_attention_indexer_common.h | 114 +++
.../contrib_ops/cuda/cuda_contrib_kernels.cc | 6 +
.../cuda/sparse/sparse_attention_indexer.cc | 343 +++++++
.../cuda/sparse/sparse_attention_indexer.h | 35 +
.../sparse/sparse_attention_indexer_impl.cu | 725 ++++++++++++++
.../sparse/sparse_attention_indexer_impl.h | 91 ++
.../core/graph/contrib_ops/bert_defs.cc | 395 ++++++++
onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 +
.../python/tools/symbolic_shape_infer.py | 74 ++
.../sparse_attention_indexer_op_test.cc | 927 ++++++++++++++++++
12 files changed, 3077 insertions(+)
create mode 100644 docs/contrib_ops/cuda/sparse_attention_indexer.md
create mode 100644 onnxruntime/contrib_ops/cpu/sparse/sparse_attention_indexer_common.h
create mode 100644 onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc
create mode 100644 onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.h
create mode 100644 onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu
create mode 100644 onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.h
create mode 100644 onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index de7882c5076f7..56368f439b128 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -1143,6 +1143,7 @@ The **OpSet Version** column uses the following notation:
|SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|SkipSimplifiedLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**
*out* input_skip_bias_sum:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|SparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* block_row_indices:**M**
*in* block_col_indices:**M**
*in* total_sequence_length:**M**
*in* key_total_sequence_lengths:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)|
+|SparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* mask:**TB**
*in* past_key:**T**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_compressed_key:**T**
*in* past_kv_buffer:**T**
*in* past_gate_buffer:**T**
*out* selected_indices:**M**
*out* present_key:**T**
*out* present_compressed_key:**T**
*out* present_kv_buffer:**T**
*out* present_gate_buffer:**T**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TB** = tensor(bool)|
|TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|UnfoldTensor|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
diff --git a/docs/contrib_ops/cuda/sparse_attention_indexer.md b/docs/contrib_ops/cuda/sparse_attention_indexer.md
new file mode 100644
index 0000000000000..c3c2b422957a6
--- /dev/null
+++ b/docs/contrib_ops/cuda/sparse_attention_indexer.md
@@ -0,0 +1,364 @@
+# SparseAttentionIndexer — Operator Documentation
+
+This document describes the `com.microsoft::SparseAttentionIndexer` contrib operator: the two
+indexer policies it implements, the schema and state contract, the CUDA kernel pipeline, and the
+known limitations and performance follow-ups.
+
+The operator answers a single question for every query token: *which* keys is the following
+attention operator allowed to read. It does not compute attention itself. Two policies are
+supported, selected by the `policy_mode` attribute:
+
+| `policy_mode` | Reference implementation | Selection granularity |
+|---|---|---|
+| `qsa` | `Qwen4ExpTextQSAIndexer` | token indices, grouped in blocks of `compress_ratio` |
+| `csa` | `DeepseekV4Indexer` / `DeepseekV4IndexerScorer` | compressed-entry indices |
+
+Source:
+[bert_defs.cc](../../../onnxruntime/core/graph/contrib_ops/bert_defs.cc) (schema),
+[sparse_attention_indexer_common.h](../../../onnxruntime/contrib_ops/cpu/sparse/sparse_attention_indexer_common.h)
+(slot map, capacity and window plan shared by the schema, the kernel and the tests),
+[sparse_attention_indexer.cc](../../../onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc),
+[sparse_attention_indexer_impl.cu](../../../onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu).
+
+---
+
+## Table of Contents
+
+1. [Overview](#1-overview)
+2. [Operator Schema](#2-operator-schema)
+3. [State Contract](#3-state-contract)
+4. [Policy `qsa`](#4-policy-qsa)
+5. [Policy `csa`](#5-policy-csa)
+6. [Rotary Embeddings](#6-rotary-embeddings)
+7. [CUDA Kernel Pipeline](#7-cuda-kernel-pipeline)
+8. [Numerics and Determinism](#8-numerics-and-determinism)
+9. [Validation Rules](#9-validation-rules)
+10. [Testing](#10-testing)
+11. [Known Limitations and Performance Follow-ups](#11-known-limitations-and-performance-follow-ups)
+
+---
+
+## 1. Overview
+
+Sparse-attention decoders run a small, cheap "indexer" attention next to the real attention. The
+indexer has its own low-rank query/key projections, its own RMSNorm and its own rotary embedding,
+and its only product is a set of indices. The real attention then reads just those keys.
+
+`SparseAttentionIndexer` implements that stage as one operator with three properties that matter
+for a runtime:
+
+- **Fixed output capacity.** `selected_indices` always has shape
+ `(batch_size, sequence_length, capacity)` where `capacity` is a pure function of the attributes
+ (`token_budget + compress_ratio - 1` for `qsa`, `index_topk` for `csa`). Entries that are not
+ used are `-1`. No output size depends on tensor *data*, so the kernel never copies a count back
+ to the host and never allocates from a device-computed size.
+- **Explicit, graph-visible state.** Everything that survives between calls — the concatenated
+ indexer key cache, the compressed-key cache and the partial-window buffers — is an input/output
+ pair. The operator caches nothing internally.
+- **A single versioned schema.** Both policies share one schema; the policy-specific inputs and
+ outputs are optional slots at fixed indices and are validated strictly (see
+ [§9](#9-validation-rules)).
+
+## 2. Operator Schema
+
+### Attributes
+
+| Attribute | Type | Required | Description |
+|---|---|---|---|
+| `policy_mode` | string | yes | Exactly `qsa` or `csa`. Any other value is rejected. |
+| `compress_ratio` | int | yes | Tokens folded into one block/entry. Must be `> 0`. |
+| `token_budget` | int | `qsa` only | Maximum number of tokens taken from complete blocks. Must be `> 0` and divisible by `compress_ratio`. Must be absent for `csa`. |
+| `index_topk` | int | `csa` only | Number of compressed entries selected per query. Must be `> 0`. Must be absent for `qsa`. |
+| `epsilon` | float | no | RMSNorm epsilon of the compressed key. Default `1e-6`. |
+| `scale` | float | no | Scale of the per-head ReLU score. Default `1/sqrt(head_size)`. |
+| `head_weight_scale` | float | `csa` only | Scale applied to `head_weights`. Default `1/sqrt(num_heads)`. Must be absent for `qsa`. |
+
+### Inputs
+
+`B` = batch size, `S` = sequence length, `N` = `num_heads`, `D` = `head_size`,
+`r` = `compress_ratio`, `P` = past sequence length, `T = P + S`, `R` = rotary width.
+
+| # | Name | Policy | Type | Shape |
+|---|---|---|---|---|
+| 0 | `query` | both | `T` | `(B, S, N, D)` |
+| 1 | `key` | both | `T` | `(B, S, D)` for `qsa`, `(B, S, 2D)` for `csa` |
+| 2 | `key_norm_weight` | both | `T` | `(D)` |
+| 3 | `cos_cache` | both | `T` | `(B, max_rotary_sequence_length, R)` |
+| 4 | `sin_cache` | both | `T` | same as `cos_cache` |
+| 5 | `mask` | `qsa` | `TB` | `(B, 1, S, T)` or `(B, S, T)` |
+| 6 | `past_key` | `qsa` | `T` | `(B, P, D)` |
+| 7 | `gate` | `csa` | `T` | `(B, S, 2D)` |
+| 8 | `position_bias` | `csa` | `T` | `(r, 2D)` |
+| 9 | `head_weights` | `csa` | `T` | `(B, S, N)` |
+| 10 | `position_ids` | `csa` | `I` | `(B, S)` |
+| 11 | `past_compressed_key` | `csa` | `T` | `(B, Pc, D)` |
+| 12 | `past_kv_buffer` | `csa` | `T` | `(B, Lb, 2D)`, `Lb` in `[0, 2r)` |
+| 13 | `past_gate_buffer` | `csa` | `T` | same as `past_kv_buffer` |
+
+### Outputs
+
+| # | Name | Policy | Type | Shape |
+|---|---|---|---|---|
+| 0 | `selected_indices` | both | `M` | `(B, S, capacity)` |
+| 1 | `present_key` | `qsa` | `T` | `(B, T, D)` |
+| 2 | `present_compressed_key` | `csa` | `T` | `(B, Pc + W, D)` |
+| 3 | `present_kv_buffer` | `csa` | `T` | `(B, Lb', 2D)` |
+| 4 | `present_gate_buffer` | `csa` | `T` | same as `present_kv_buffer` |
+
+`W` is the number of complete windows closed by this call and `Lb'` the new buffer length; both
+follow from `Lb`, `S` and `r` alone (see [§5](#5-policy-csa)).
+
+### Type constraints
+
+| Name | Allowed types |
+|---|---|
+| `T` | `tensor(float)`, `tensor(float16)`, `tensor(bfloat16)` |
+| `TB` | `tensor(bool)` |
+| `I` | `tensor(int64)` |
+| `M` | `tensor(int32)` |
+
+Only the CUDA execution provider registers a kernel. There is no CPU kernel; the header under
+`contrib_ops/cpu/sparse/` only holds the CUDA-free constants that the schema, the kernel and the
+tests must agree on.
+
+### Output slot discipline
+
+A `qsa` node declares exactly 2 outputs (`selected_indices`, `present_key`). A `csa` node declares
+exactly 5, leaving slot 1 as a missing optional so that the three `csa` state outputs keep their
+fixed indices. Shape inference verifies the declared output count *before* touching any output, so
+`getOutputType` is never called past the declared range.
+
+## 3. State Contract
+
+`SparseAttentionIndexer` is a pure function of its inputs. Every value it needs on the next call is
+returned as an output, so the caller (or the ORT session binding) owns the buffers:
+
+| Policy | State pair |
+|---|---|
+| `qsa` | `past_key` → `present_key` |
+| `csa` | `past_compressed_key` → `present_compressed_key` |
+| `csa` | `past_kv_buffer` → `present_kv_buffer` |
+| `csa` | `past_gate_buffer` → `present_gate_buffer` |
+
+`present_key` and `present_compressed_key` grow by a number of positions that is known from the
+input shapes and the attributes, so the graph can pre-allocate them. The two `csa` buffers stay
+bounded by `2r - 1` positions.
+
+## 4. Policy `qsa`
+
+For every `(b, s)`:
+
+1. **Visible set.** `visible = [t for t in range(T) if mask[b, s, t]]`, in ascending `t`.
+2. **Complete blocks.** `nblocks = len(visible) // r`. Block `j` covers
+ `visible[j*r : (j+1)*r]`.
+3. **Pooled key.** `k_j = mean` of the `r` raw `present_key` rows of block `j`, then
+ `RMSNorm(k_j) * key_norm_weight`, then rotary at absolute position `visible[j*r]`.
+4. **Score.** `score_j = scale * sum_h ReLU(q_h · k_j)` where `q_h` is the rotated query head
+ at absolute position `P + s`.
+5. **Selection.** `topk = min(token_budget // r, nblocks)` blocks, ordered by decreasing score.
+ Their tokens are emitted in that block order, `r` token indices per block.
+6. **Tail.** The `len(visible) % r` tokens of the trailing incomplete block,
+ `visible[nblocks*r:]`, are appended unconditionally.
+7. Remaining entries up to `capacity = token_budget + r - 1` are `-1`.
+
+The capacity is exactly `token_budget` selected tokens plus at most `r - 1` tail tokens.
+
+## 5. Policy `csa`
+
+### Window bookkeeping
+
+Let `ext = concat(past_kv_buffer, key)` along the sequence axis (and likewise for the gates).
+The buffer is split as
+
+```
+overlap = (Lb >= r) ? r : 0 # previous complete window, the "Ca" source
+leftover = Lb - overlap # tokens of the still-open window, always < r
+pending = leftover + S
+W = pending / r # complete windows closed by this call
+```
+
+Window `w` (`0 <= w < W`) covers `ext[overlap + w*r : overlap + (w+1)*r]`. Its `Ca` partner is the
+preceding `r` tokens, which exist iff `w >= 1 || overlap == r`. Because `leftover < r` always holds,
+a buffer length `>= r` unambiguously means "the previous complete window is present" — no extra
+state tensor is needed to disambiguate.
+
+The new buffer starts at `overlap + (W-1)*r` when `W > 0` (the last complete window followed by the
+new leftover) and at `0` otherwise, giving
+
+```
+Lb' = (W > 0) ? r + pending % r : Lb + S
+```
+
+`Lb'` is again in `[0, 2r)`.
+
+### Compression
+
+For window `w`, a `2r`-slot pooling window is built from `[B, 2r, D]` values and `[B, 2r, D]` gates:
+
+- slots `[r, 2r)` take the **`Cb`** half (channels `[D, 2D)`) of the window's own `r` tokens;
+- slots `[0, r)` take the **`Ca`** half (channels `[0, D)`) of the preceding `r` tokens, or are
+ masked out (value `0`, gate `-inf`) when that window does not exist.
+
+`position_bias[slot % r]` is added to the raw gate (the buffers store the **raw** projection, so
+the bias is re-applied at use time and never accumulates). A softmax over the `2r` slots is taken
+**per channel `d`**, the values are weighted and summed, the result is RMS-normalized with
+`key_norm_weight`, and finally rotated at absolute position `(Pc + w) * r`.
+
+### Scoring and selection
+
+```
+scores[b, s, e] = head_weight_scale * sum_h head_weights[b, s, h] * scale * ReLU(q_h · k_e)
+threshold[b, s] = (position_ids[b, s] + 1) / r # integer division
+scores[b, s, e] = -inf for e >= threshold[b, s]
+```
+
+The top `min(index_topk, Pc + W)` entries are emitted in decreasing score order; any emitted entry
+whose index is `>= threshold` (which only happens when fewer than `index_topk` entries are
+visible) is written as `-1`, and the remaining capacity is `-1`.
+
+## 6. Rotary Embeddings
+
+Both policies reuse the model's precomputed `cos_cache` / `sin_cache`, indexed by absolute
+position. This keeps the exact rotary variant of each reference model without re-deriving
+frequencies inside the kernel:
+
+- **`qsa` (leading, split-half).** `R = cos_cache.shape[2]` channels are rotated, and the rotation
+ is the split-half form
+ `out[i] = x[i]*cos[i] - x[i + R/2]*sin[i]`, `out[i + R/2] = x[i + R/2]*cos[i + R/2] + x[i]*sin[i + R/2]`,
+ applied to channels `[0, R)`. Channels `[R, D)` pass through. Feeding an MRoPE-expanded
+ `cos_cache`/`sin_cache` therefore reproduces the model's MRoPE exactly, because the operator
+ never recomputes the position→angle mapping.
+- **`csa` (trailing, interleaved).** `2R` channels are rotated and they are the **last** `2R`
+ channels of the head. `cos`/`sin` are half-width, so entry `j >> 1` covers the channel pair
+ `(j, j+1)`, and the rotation is the interleaved form
+ `out[2i] = x[2i]*cos[i] - x[2i+1]*sin[i]`, `out[2i+1] = x[2i+1]*cos[i] + x[2i]*sin[i]`.
+
+Positions are clamped into `[0, max_rotary_sequence_length - 1]` inside the kernel, so an
+out-of-range `position_ids` value cannot read out of bounds.
+
+### `key_norm_weight` and zero-centered gamma
+
+`key_norm_weight` is the **effective** multiplier: the kernel computes
+`normalized * key_norm_weight`. Qwen's `Qwen4ExpTextRMSNorm` multiplies by `1 + gamma`. Exporters
+targeting that model must fold the addition into the initializer (`key_norm_weight = 1 + gamma`).
+DeepSeek's `DeepseekV4RMSNorm` uses a plain `weight *`, so its tensor is passed through unchanged.
+
+## 7. CUDA Kernel Pipeline
+
+All kernels use a fixed block of 128 threads (a power of two, required by the shared-memory block
+reductions). Grid sizes are clamped to 65535 blocks and the elementwise kernels use grid-stride
+loops, so no launch configuration depends on tensor data.
+
+### `qsa`
+
+| Stage | Kernel | Parallelism |
+|---|---|---|
+| 1 | `ConcatPastKeyKernel` | element |
+| 2 | `RotateQueryKernel` | one block per `(b, s, h)` |
+| 3 | `CompactVisibleKernel` | one block per `(b, s)`; Hillis–Steele scan in shared memory |
+| 4 | `QsaBlockScoreKernel` | one block per `(b, s, block)`; mean-pool → RMSNorm → rotary → score |
+| 5 | `QsaSelectKernel` | one block per `(b, s)`; iterated block arg-max |
+
+### `csa`
+
+| Stage | Kernel | Parallelism |
+|---|---|---|
+| 1 | `CsaCopyPastCompressedKernel` | element |
+| 2 | `CsaCompressKernel` | one block per `(b, window)`; per-channel softmax over `2r` slots |
+| 3 | `CsaCopyBufferKernel` | element |
+| 4 | `RotateQueryKernel` | one block per `(b, s, h)` |
+| 5 | `CsaScoreKernel` | one block per `(b, s, entry)` |
+| 6 | `CsaSelectKernel` | one block per `(b, s)`; iterated block arg-max |
+
+### Workspaces
+
+| Policy | Buffer | Elements |
+|---|---|---|
+| `qsa` | float | `B*S*N*D` (rotated query) + `B*S*max_block_count` (block scores) |
+| `qsa` | int32 | `B*S*T` (visible indices) + `B*S` (visible counts) |
+| `csa` | float | `B*S*N*D` (rotated query) + `B*S*present_compressed_length` (scores) |
+
+Every size is derived from shapes and attributes only.
+
+### Selection without a visited bitmap
+
+`QsaSelectKernel` and `CsaSelectKernel` repeatedly scan for the next element in the total order
+"score descending, index ascending", using the previously emitted `(score, index)` pair as the
+cursor. That avoids an `O(entries)` bitmap in shared memory, keeps the selection deterministic and
+makes ties resolve to the smaller index. The cost is `O(topk * entries)` per query row, which is
+the main performance follow-up below.
+
+## 8. Numerics and Determinism
+
+- Pooling, softmax, RMS normalization, rotary and scoring are all done in `float32`; only the final
+ store is rounded to the tensor element type. This is a deliberate deviation from the reference
+ implementations, which for `qsa` run in the model dtype — the operator is strictly more accurate,
+ never less.
+- The `2r`-slot softmax is a two-pass (max-subtracted) formulation, so a fully masked slot column
+ cannot produce `NaN`.
+- Selection order is a total order, so the output is bitwise reproducible for a given input.
+- `float16` and `bfloat16` are supported for all `T` tensors. `bfloat16` uses the conversion
+ helpers in `cu_inc/cuda_type_helper.cuh`, which are emulated on pre-`sm_80` devices, so no
+ architecture guard is required.
+
+## 9. Validation Rules
+
+Shape inference and the kernel both reject:
+
+- a `policy_mode` other than `qsa` / `csa`;
+- `compress_ratio <= 0`;
+- a `qsa` node that provides any `csa`-only input (slots 7–13) or attribute, and vice versa;
+- a missing required input for the active policy;
+- `token_budget` absent, `<= 0`, or not divisible by `compress_ratio` for `qsa`;
+- `index_topk` absent or `<= 0` for `csa`;
+- an output count other than 2 (`qsa`) or 5 (`csa`);
+- `past_kv_buffer` / `past_gate_buffer` lengths outside `[0, 2 * compress_ratio)` or differing from
+ each other;
+- rank or dimension mismatches between `query`, `key`, the caches and the buffers.
+
+Because the checks live in shape inference, most misuse fails at `Graph::Resolve()` with a clear
+message rather than at kernel launch.
+
+## 10. Testing
+
+[`onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc`](../../../onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc)
+contains:
+
+- **Shape-inference tests** that build a real `Model`, run `Graph::Resolve()` and assert the
+ inferred element type and shape of every output for both policies, including the buffer-only
+ `csa` case (`W == 0`).
+- **Negative tests** that assert the exact validation message for each rule in
+ [§9](#9-validation-rules), including a `csa` node that declares too few outputs — the case that
+ would otherwise write past the declared output range.
+- **Numeric tests** (`float`, `float16`, `bfloat16`) that run the CUDA kernel against a float
+ reference implementation of both policies written directly from the reference semantics. Inputs
+ are round-tripped through the tested element type before the reference runs, so the reference
+ sees exactly the values the kernel reads. These tests skip when no CUDA EP is available.
+
+Run them with:
+
+```bash
+./build/Linux/Release/onnxruntime_provider_test --gtest_filter='SparseAttentionIndexer*'
+```
+
+## 11. Known Limitations and Performance Follow-ups
+
+The implementation is correctness-first. The following are known and deliberate:
+
+1. **Selection is `O(topk * entries)` per query row.** Each emitted index costs a full block-wide
+ scan. A radix-select or a per-row bitonic top-k would reduce this to roughly one pass, and is the
+ single biggest win for large `index_topk` / `token_budget`.
+2. **Scoring is not tensor-core accelerated.** `QsaBlockScoreKernel` and `CsaScoreKernel` compute
+ `q · k` with a shared-memory block reduction, one dot product per block. A tiled GEMM (or a
+ fused `ReLU`+reduce epilogue) would be far better once shapes grow.
+3. **The rotated query is materialized in `float32`.** That costs `B*S*N*D` floats of workspace.
+ Fusing the rotation into the scoring kernels removes the traffic at the cost of recomputing the
+ rotation per block.
+4. **`CompactVisibleKernel` is `O(T)` per query row** and re-reads the mask for every `s`. For long
+ contexts a batched exclusive scan over the whole `(B, S, T)` mask would be cheaper.
+5. **`present_key` / `present_compressed_key` are copied every call.** An in-place cache with a
+ `past_sequence_length` input (as `GroupQueryAttention` does) would avoid the copy, at the cost of
+ a less explicit state contract.
+6. **`compress_ratio` and `head_size` are not specialized.** Templating the hot kernels on a small
+ set of common values would remove the dynamic loop bounds.
+7. **No CPU kernel.** The operator is CUDA-only today.
diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_indexer_common.h
new file mode 100644
index 0000000000000..d9eb3d1671cca
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_indexer_common.h
@@ -0,0 +1,114 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include
+#include
+
+namespace onnxruntime {
+namespace contrib {
+namespace sparse_attention_indexer {
+
+// Values accepted by the policy_mode attribute of com.microsoft.SparseAttentionIndexer.
+constexpr const char* kPolicyModeQsa = "qsa";
+constexpr const char* kPolicyModeCsa = "csa";
+
+enum class Policy {
+ kQsa,
+ kCsa,
+};
+
+inline bool TryParsePolicy(const std::string& policy_mode, Policy& policy) {
+ if (policy_mode == kPolicyModeQsa) {
+ policy = Policy::kQsa;
+ return true;
+ }
+ if (policy_mode == kPolicyModeCsa) {
+ policy = Policy::kCsa;
+ return true;
+ }
+ return false;
+}
+
+// Input slots. Slots 5-6 belong to policy_mode="qsa" and slots 7-13 to policy_mode="csa";
+// a slot that does not belong to the active policy must be omitted from the node.
+enum InputIndex : int {
+ kQuery = 0,
+ kKey = 1,
+ kKeyNormWeight = 2,
+ kCosCache = 3,
+ kSinCache = 4,
+ kMask = 5,
+ kPastKey = 6,
+ kGate = 7,
+ kPositionBias = 8,
+ kHeadWeights = 9,
+ kPositionIds = 10,
+ kPastCompressedKey = 11,
+ kPastKvBuffer = 12,
+ kPastGateBuffer = 13,
+ kInputCount = 14,
+};
+
+// Output slots. Slot 1 belongs to policy_mode="qsa" and slots 2-4 to policy_mode="csa".
+enum OutputIndex : int {
+ kSelectedIndices = 0,
+ kPresentKey = 1,
+ kPresentCompressedKey = 2,
+ kPresentKvBuffer = 3,
+ kPresentGateBuffer = 4,
+ kOutputCount = 5,
+};
+
+// A "qsa" node declares selected_indices + present_key; a "csa" node declares every slot so that
+// the three csa state outputs keep their fixed indices (slot 1 is left as a missing optional).
+constexpr int kQsaOutputCount = 2;
+constexpr int kCsaOutputCount = 5;
+
+// Number of selected entries emitted per query. The capacity only depends on attributes, so it is
+// a compile-time constant of the graph rather than a function of the data.
+inline int64_t SelectedCapacity(Policy policy, int64_t token_budget, int64_t index_topk,
+ int64_t compress_ratio) {
+ return policy == Policy::kQsa ? token_budget + compress_ratio - 1 : index_topk;
+}
+
+// How the "csa" token buffer is split and how many new compressed entries this call emits.
+//
+// past_kv_buffer / past_gate_buffer carry `overlap_length` tokens of the previous complete window
+// (the Ca operand of the next window) followed by `leftover_length` tokens of an incomplete window.
+// Since `leftover_length` is always < compress_ratio, a buffer length >= compress_ratio uniquely
+// means "the previous complete window is present", so no extra state tensor is needed.
+struct CsaWindowPlan {
+ int64_t overlap_length = 0; // compress_ratio, or 0 before the first complete window
+ int64_t leftover_length = 0; // tokens of the current incomplete window
+ int64_t new_window_count = 0; // complete windows closed by this call
+ int64_t present_buffer_length = 0; // length of present_kv_buffer / present_gate_buffer
+ int64_t present_buffer_start = 0; // offset of that buffer inside [past buffer | new tokens]
+};
+
+inline bool TryComputeCsaWindowPlan(int64_t past_buffer_length, int64_t sequence_length,
+ int64_t compress_ratio, CsaWindowPlan& plan) {
+ if (compress_ratio <= 0 || sequence_length < 0 || past_buffer_length < 0 ||
+ past_buffer_length >= 2 * compress_ratio) {
+ return false;
+ }
+
+ 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;
+ plan.new_window_count = pending / compress_ratio;
+ if (plan.new_window_count > 0) {
+ plan.present_buffer_length = compress_ratio + pending % compress_ratio;
+ plan.present_buffer_start = plan.overlap_length + (plan.new_window_count - 1) * compress_ratio;
+ } else {
+ plan.present_buffer_length = past_buffer_length + sequence_length;
+ plan.present_buffer_start = 0;
+ }
+ return true;
+}
+
+} // namespace sparse_attention_indexer
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
index 95d81e53a3d49..03529ea5eb417 100644
--- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
@@ -249,6 +249,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16_MLFloat16, DecoderMaskedMultiHead
class CUDA_MS_OP_CLASS_NAME(1, GemmFloat8);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, SparseAttention);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, SparseAttention);
+class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, SparseAttentionIndexer);
+class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, SparseAttentionIndexer);
+class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, SparseAttentionIndexer);
class CUDA_MS_OP_THREE_TYPED_CLASS_NAME(1, uint8_t, float, int32_t, GatherBlockQuantized);
class CUDA_MS_OP_THREE_TYPED_CLASS_NAME(1, uint8_t, MLFloat16, int32_t, GatherBlockQuantized);
@@ -551,6 +554,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc
new file mode 100644
index 0000000000000..5a3d84eacb7a6
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc
@@ -0,0 +1,343 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cuda/sparse/sparse_attention_indexer.h"
+
+#include
+#include
+
+#include "contrib_ops/cuda/sparse/sparse_attention_indexer_impl.h"
+#include "core/providers/cuda/cuda_common.h"
+#include "core/providers/cuda/cuda_type_conversion.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+namespace sai = onnxruntime::contrib::sparse_attention_indexer;
+
+#define REGISTER_KERNEL_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ SparseAttentionIndexer, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCudaExecutionProvider, \
+ (*KernelDefBuilder::Create()) \
+ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("TB", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("I", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("M", DataTypeImpl::GetTensorType()), \
+ SparseAttentionIndexer);
+
+REGISTER_KERNEL_TYPED(float)
+REGISTER_KERNEL_TYPED(MLFloat16)
+REGISTER_KERNEL_TYPED(BFloat16)
+
+#undef REGISTER_KERNEL_TYPED
+
+namespace {
+
+Status CheckShape(const Tensor* tensor, const char* name, std::initializer_list expected) {
+ ORT_RETURN_IF(tensor == nullptr, "SparseAttentionIndexer: ", name, " is required");
+ const TensorShape expected_shape(expected);
+ ORT_RETURN_IF_NOT(tensor->Shape() == expected_shape, "SparseAttentionIndexer: ", name, " must have shape ",
+ expected_shape.ToString(), ", got ", tensor->Shape().ToString());
+ return Status::OK();
+}
+
+} // namespace
+
+template
+SparseAttentionIndexer::SparseAttentionIndexer(const OpKernelInfo& info) : CudaKernel(info) {
+ std::string policy_mode;
+ ORT_ENFORCE(info.GetAttr("policy_mode", &policy_mode).IsOK(),
+ "SparseAttentionIndexer: policy_mode is required");
+ ORT_ENFORCE(sai::TryParsePolicy(policy_mode, policy_), "SparseAttentionIndexer: policy_mode must be '",
+ sai::kPolicyModeQsa, "' or '", sai::kPolicyModeCsa, "', got '", policy_mode, "'");
+
+ ORT_ENFORCE(info.GetAttr("compress_ratio", &compress_ratio_).IsOK(),
+ "SparseAttentionIndexer: compress_ratio is required");
+ ORT_ENFORCE(compress_ratio_ > 0, "SparseAttentionIndexer: compress_ratio must be > 0, got ", compress_ratio_);
+
+ const bool has_token_budget = info.GetAttr("token_budget", &token_budget_).IsOK();
+ const bool has_index_topk = info.GetAttr("index_topk", &index_topk_).IsOK();
+ float head_weight_scale = 0.0f;
+ const bool has_head_weight_scale = info.GetAttr("head_weight_scale", &head_weight_scale).IsOK();
+
+ if (policy_ == sai::Policy::kQsa) {
+ ORT_ENFORCE(has_token_budget, "SparseAttentionIndexer: token_budget is required when policy_mode is 'qsa'");
+ ORT_ENFORCE(!has_index_topk && !has_head_weight_scale,
+ "SparseAttentionIndexer: index_topk and head_weight_scale must not be set when policy_mode is 'qsa'");
+ ORT_ENFORCE(token_budget_ > 0 && token_budget_ % compress_ratio_ == 0,
+ "SparseAttentionIndexer: token_budget must be > 0 and divisible by compress_ratio, got token_budget=",
+ token_budget_, " compress_ratio=", compress_ratio_);
+ index_topk_ = 0;
+ } else {
+ ORT_ENFORCE(has_index_topk, "SparseAttentionIndexer: index_topk is required when policy_mode is 'csa'");
+ ORT_ENFORCE(!has_token_budget, "SparseAttentionIndexer: token_budget must not be set when policy_mode is 'csa'");
+ ORT_ENFORCE(index_topk_ > 0, "SparseAttentionIndexer: index_topk must be > 0, got ", index_topk_);
+ token_budget_ = 0;
+ }
+
+ epsilon_ = info.GetAttrOrDefault("epsilon", 1.0e-6f);
+ ORT_ENFORCE(epsilon_ >= 0.0f, "SparseAttentionIndexer: epsilon must be >= 0, got ", epsilon_);
+ scale_ = info.GetAttrOrDefault("scale", 0.0f);
+ head_weight_scale_ = has_head_weight_scale ? head_weight_scale : 0.0f;
+}
+
+template
+Status SparseAttentionIndexer::ComputeInternal(OpKernelContext* context) const {
+ const bool is_qsa = policy_ == sai::Policy::kQsa;
+ for (int index = sai::kMask; index < sai::kInputCount; ++index) {
+ const bool policy_owns_slot = is_qsa ? (index <= sai::kPastKey) : (index >= sai::kGate);
+ const bool provided = index < context->InputCount() && context->Input(index) != nullptr;
+ ORT_RETURN_IF(provided != policy_owns_slot, "SparseAttentionIndexer: input ", index,
+ provided ? " must be omitted for policy_mode '" : " is required for policy_mode '",
+ is_qsa ? sai::kPolicyModeQsa : sai::kPolicyModeCsa, "'");
+ }
+
+ return is_qsa ? ComputeQsa(context) : ComputeCsa(context);
+}
+
+template
+Status SparseAttentionIndexer::ComputeQsa(OpKernelContext* context) const {
+ using CudaT = typename OrtToCudaType::type;
+
+ const Tensor* query = context->Input(sai::kQuery);
+ const Tensor* key = context->Input(sai::kKey);
+ const Tensor* key_norm_weight = context->Input(sai::kKeyNormWeight);
+ const Tensor* cos_cache = context->Input(sai::kCosCache);
+ const Tensor* sin_cache = context->Input(sai::kSinCache);
+ const Tensor* mask = context->Input(sai::kMask);
+ const Tensor* past_key = context->Input(sai::kPastKey);
+
+ const auto& query_shape = query->Shape();
+ ORT_RETURN_IF_NOT(query_shape.NumDimensions() == 4,
+ "SparseAttentionIndexer: query must have shape (batch_size, sequence_length, num_heads, head_size)"
+ ", got ",
+ query_shape.ToString());
+ const int64_t batch_size = query_shape[0];
+ const int64_t sequence_length = query_shape[1];
+ const int64_t num_heads = query_shape[2];
+ const int64_t head_size = query_shape[3];
+
+ const auto& cos_shape = cos_cache->Shape();
+ ORT_RETURN_IF_NOT(cos_shape.NumDimensions() == 3 && cos_shape[0] == batch_size && cos_shape[1] > 0,
+ "SparseAttentionIndexer: cos_cache must have shape "
+ "(batch_size, max_rotary_sequence_length, rotary_width), got ",
+ cos_shape.ToString());
+ const int64_t max_rotary_length = cos_shape[1];
+ const int64_t rotary_width = cos_shape[2];
+ ORT_RETURN_IF_NOT(sin_cache->Shape() == cos_shape,
+ "SparseAttentionIndexer: sin_cache must have the same shape as "
+ "cos_cache");
+ ORT_RETURN_IF_NOT(rotary_width > 0 && rotary_width % 2 == 0 && rotary_width <= head_size,
+ "SparseAttentionIndexer: policy_mode 'qsa' requires an even rotary_width in (0, head_size], got ",
+ rotary_width);
+
+ const auto& past_shape = past_key->Shape();
+ ORT_RETURN_IF_NOT(past_shape.NumDimensions() == 3 && past_shape[0] == batch_size && past_shape[2] == head_size,
+ "SparseAttentionIndexer: past_key must have shape (batch_size, past_sequence_length, head_size)"
+ ", got ",
+ past_shape.ToString());
+ const int64_t past_sequence_length = past_shape[1];
+ const int64_t total_sequence_length = past_sequence_length + sequence_length;
+
+ ORT_RETURN_IF_ERROR(CheckShape(key, "key", {batch_size, sequence_length, head_size}));
+ ORT_RETURN_IF_ERROR(CheckShape(key_norm_weight, "key_norm_weight", {head_size}));
+
+ const auto& mask_shape = mask->Shape();
+ const bool mask_is_4d = mask_shape.NumDimensions() == 4;
+ ORT_RETURN_IF_NOT(
+ (mask_is_4d && mask_shape[0] == batch_size && mask_shape[1] == 1 && mask_shape[2] == sequence_length &&
+ mask_shape[3] == total_sequence_length) ||
+ (mask_shape.NumDimensions() == 3 && mask_shape[0] == batch_size && mask_shape[1] == sequence_length &&
+ mask_shape[2] == total_sequence_length),
+ "SparseAttentionIndexer: mask must have shape (batch_size, 1, sequence_length, total_sequence_length) or "
+ "(batch_size, sequence_length, total_sequence_length) with total_sequence_length=",
+ total_sequence_length, ", got ", mask_shape.ToString());
+
+ SparseAttentionIndexerParams params;
+ params.batch_size = static_cast(batch_size);
+ params.sequence_length = static_cast(sequence_length);
+ params.num_heads = static_cast(num_heads);
+ params.head_size = static_cast(head_size);
+ params.rotary_width = static_cast(rotary_width);
+ params.max_rotary_length = static_cast(max_rotary_length);
+ params.compress_ratio = static_cast(compress_ratio_);
+ params.capacity = static_cast(
+ sai::SelectedCapacity(sai::Policy::kQsa, token_budget_, index_topk_, compress_ratio_));
+ params.epsilon = epsilon_;
+ params.scale = scale_ != 0.0f ? scale_ : 1.0f / std::sqrt(static_cast(head_size));
+ params.past_sequence_length = static_cast(past_sequence_length);
+ params.total_sequence_length = static_cast(total_sequence_length);
+ params.max_block_count = static_cast(total_sequence_length / compress_ratio_);
+ params.block_topk = static_cast(token_budget_ / compress_ratio_);
+
+ Tensor* selected_indices = context->Output(sai::kSelectedIndices,
+ TensorShape({batch_size, sequence_length, params.capacity}));
+ Tensor* present_key = context->Output(sai::kPresentKey, TensorShape({batch_size, total_sequence_length, head_size}));
+ ORT_RETURN_IF(selected_indices == nullptr || present_key == nullptr,
+ "SparseAttentionIndexer: policy_mode 'qsa' requires both selected_indices and present_key outputs");
+
+ auto float_workspace = GetScratchBuffer(GetQsaWorkspaceFloatCount(params), context->GetComputeStream());
+ auto int_workspace = GetScratchBuffer(GetQsaWorkspaceIntCount(params), context->GetComputeStream());
+
+ return LaunchQsaSparseAttentionIndexer(
+ Stream(context), params,
+ reinterpret_cast(query->Data()),
+ reinterpret_cast(key->Data()),
+ reinterpret_cast(key_norm_weight->Data()),
+ reinterpret_cast(cos_cache->Data()),
+ reinterpret_cast(sin_cache->Data()),
+ mask->Data(),
+ reinterpret_cast(past_key->Data()),
+ selected_indices->MutableData(),
+ reinterpret_cast(present_key->MutableData()),
+ float_workspace.get(),
+ int_workspace.get());
+}
+
+template
+Status SparseAttentionIndexer::ComputeCsa(OpKernelContext* context) const {
+ using CudaT = typename OrtToCudaType::type;
+
+ const Tensor* query = context->Input(sai::kQuery);
+ const Tensor* key = context->Input(sai::kKey);
+ const Tensor* key_norm_weight = context->Input(sai::kKeyNormWeight);
+ const Tensor* cos_cache = context->Input(sai::kCosCache);
+ const Tensor* sin_cache = context->Input(sai::kSinCache);
+ const Tensor* gate = context->Input(sai::kGate);
+ const Tensor* position_bias = context->Input(sai::kPositionBias);
+ const Tensor* head_weights = context->Input(sai::kHeadWeights);
+ const Tensor* position_ids = context->Input(sai::kPositionIds);
+ const Tensor* past_compressed_key = context->Input(sai::kPastCompressedKey);
+ const Tensor* past_kv_buffer = context->Input(sai::kPastKvBuffer);
+ const Tensor* past_gate_buffer = context->Input(sai::kPastGateBuffer);
+
+ const auto& query_shape = query->Shape();
+ ORT_RETURN_IF_NOT(query_shape.NumDimensions() == 4,
+ "SparseAttentionIndexer: query must have shape (batch_size, sequence_length, num_heads, head_size)"
+ ", got ",
+ query_shape.ToString());
+ const int64_t batch_size = query_shape[0];
+ const int64_t sequence_length = query_shape[1];
+ const int64_t num_heads = query_shape[2];
+ const int64_t head_size = query_shape[3];
+ const int64_t width = 2 * head_size;
+
+ const auto& cos_shape = cos_cache->Shape();
+ ORT_RETURN_IF_NOT(cos_shape.NumDimensions() == 3 && cos_shape[0] == batch_size && cos_shape[1] > 0,
+ "SparseAttentionIndexer: cos_cache must have shape "
+ "(batch_size, max_rotary_sequence_length, rotary_width), got ",
+ cos_shape.ToString());
+ const int64_t max_rotary_length = cos_shape[1];
+ const int64_t rotary_width = cos_shape[2];
+ ORT_RETURN_IF_NOT(sin_cache->Shape() == cos_shape,
+ "SparseAttentionIndexer: sin_cache must have the same shape as "
+ "cos_cache");
+ ORT_RETURN_IF_NOT(rotary_width > 0 && 2 * rotary_width <= head_size,
+ "SparseAttentionIndexer: policy_mode 'csa' requires 0 < 2 * rotary_width <= head_size, got "
+ "rotary_width=",
+ rotary_width, " head_size=", head_size);
+
+ const auto& past_compressed_shape = past_compressed_key->Shape();
+ ORT_RETURN_IF_NOT(past_compressed_shape.NumDimensions() == 3 && past_compressed_shape[0] == batch_size &&
+ past_compressed_shape[2] == head_size,
+ "SparseAttentionIndexer: past_compressed_key must have shape "
+ "(batch_size, past_compressed_length, head_size), got ",
+ past_compressed_shape.ToString());
+ const int64_t past_compressed_length = past_compressed_shape[1];
+
+ const auto& past_buffer_shape = past_kv_buffer->Shape();
+ ORT_RETURN_IF_NOT(past_buffer_shape.NumDimensions() == 3 && past_buffer_shape[0] == batch_size &&
+ past_buffer_shape[2] == width,
+ "SparseAttentionIndexer: past_kv_buffer must have shape "
+ "(batch_size, buffer_length, 2 * head_size), got ",
+ past_buffer_shape.ToString());
+ const int64_t past_buffer_length = past_buffer_shape[1];
+
+ ORT_RETURN_IF_ERROR(CheckShape(key, "key", {batch_size, sequence_length, width}));
+ ORT_RETURN_IF_ERROR(CheckShape(key_norm_weight, "key_norm_weight", {head_size}));
+ ORT_RETURN_IF_ERROR(CheckShape(gate, "gate", {batch_size, sequence_length, width}));
+ ORT_RETURN_IF_ERROR(CheckShape(position_bias, "position_bias", {compress_ratio_, width}));
+ ORT_RETURN_IF_ERROR(CheckShape(head_weights, "head_weights", {batch_size, sequence_length, num_heads}));
+ ORT_RETURN_IF_ERROR(CheckShape(position_ids, "position_ids", {batch_size, sequence_length}));
+ ORT_RETURN_IF_ERROR(CheckShape(past_gate_buffer, "past_gate_buffer",
+ {batch_size, past_buffer_length, width}));
+
+ sai::CsaWindowPlan plan;
+ ORT_RETURN_IF_NOT(sai::TryComputeCsaWindowPlan(past_buffer_length, sequence_length, compress_ratio_, plan),
+ "SparseAttentionIndexer: past_kv_buffer sequence length must be in [0, 2 * compress_ratio), got ",
+ past_buffer_length);
+ const int64_t present_compressed_length = past_compressed_length + plan.new_window_count;
+
+ SparseAttentionIndexerParams params;
+ params.batch_size = static_cast(batch_size);
+ params.sequence_length = static_cast(sequence_length);
+ params.num_heads = static_cast(num_heads);
+ params.head_size = static_cast(head_size);
+ params.rotary_width = static_cast(rotary_width);
+ params.max_rotary_length = static_cast(max_rotary_length);
+ params.compress_ratio = static_cast(compress_ratio_);
+ params.capacity = static_cast(
+ sai::SelectedCapacity(sai::Policy::kCsa, token_budget_, index_topk_, compress_ratio_));
+ params.epsilon = epsilon_;
+ params.scale = scale_ != 0.0f ? scale_ : 1.0f / std::sqrt(static_cast(head_size));
+ params.past_compressed_length = static_cast(past_compressed_length);
+ params.present_compressed_length = static_cast(present_compressed_length);
+ params.past_buffer_length = static_cast(past_buffer_length);
+ params.overlap_length = static_cast(plan.overlap_length);
+ params.new_window_count = static_cast(plan.new_window_count);
+ params.present_buffer_length = static_cast(plan.present_buffer_length);
+ params.present_buffer_start = static_cast(plan.present_buffer_start);
+ params.index_topk = static_cast(index_topk_);
+ params.head_weight_scale =
+ head_weight_scale_ != 0.0f ? head_weight_scale_ : 1.0f / std::sqrt(static_cast(num_heads));
+
+ Tensor* selected_indices = context->Output(sai::kSelectedIndices,
+ TensorShape({batch_size, sequence_length, params.capacity}));
+ Tensor* present_compressed_key = context->Output(
+ sai::kPresentCompressedKey, TensorShape({batch_size, present_compressed_length, head_size}));
+ Tensor* present_kv_buffer =
+ context->Output(sai::kPresentKvBuffer, TensorShape({batch_size, plan.present_buffer_length, width}));
+ Tensor* present_gate_buffer =
+ context->Output(sai::kPresentGateBuffer, TensorShape({batch_size, plan.present_buffer_length, width}));
+ ORT_RETURN_IF(selected_indices == nullptr || present_compressed_key == nullptr || present_kv_buffer == nullptr ||
+ present_gate_buffer == nullptr,
+ "SparseAttentionIndexer: policy_mode 'csa' requires selected_indices, present_compressed_key, "
+ "present_kv_buffer and present_gate_buffer outputs");
+
+ auto float_workspace = GetScratchBuffer(GetCsaWorkspaceFloatCount(params), context->GetComputeStream());
+
+ const CudaT* empty = nullptr;
+ return LaunchCsaSparseAttentionIndexer(
+ Stream(context), params,
+ reinterpret_cast(query->Data()),
+ reinterpret_cast(key->Data()),
+ reinterpret_cast(key_norm_weight->Data()),
+ reinterpret_cast(cos_cache->Data()),
+ reinterpret_cast(sin_cache->Data()),
+ reinterpret_cast(gate->Data()),
+ reinterpret_cast(position_bias->Data()),
+ reinterpret_cast(head_weights->Data()),
+ position_ids->Data(),
+ past_compressed_length > 0 ? reinterpret_cast(past_compressed_key->Data()) : empty,
+ past_buffer_length > 0 ? reinterpret_cast(past_kv_buffer->Data()) : empty,
+ past_buffer_length > 0 ? reinterpret_cast(past_gate_buffer->Data()) : empty,
+ selected_indices->MutableData(),
+ present_compressed_length > 0 ? reinterpret_cast(present_compressed_key->MutableData()) : nullptr,
+ plan.present_buffer_length > 0 ? reinterpret_cast(present_kv_buffer->MutableData()) : nullptr,
+ plan.present_buffer_length > 0 ? reinterpret_cast(present_gate_buffer->MutableData()) : nullptr,
+ float_workspace.get());
+}
+
+template class SparseAttentionIndexer;
+template class SparseAttentionIndexer;
+template class SparseAttentionIndexer;
+
+} // namespace cuda
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.h b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.h
new file mode 100644
index 0000000000000..0abbb6064a3a4
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.h
@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "contrib_ops/cpu/sparse/sparse_attention_indexer_common.h"
+#include "core/common/common.h"
+#include "core/providers/cuda/cuda_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+template
+class SparseAttentionIndexer final : public onnxruntime::cuda::CudaKernel {
+ public:
+ explicit SparseAttentionIndexer(const OpKernelInfo& info);
+ Status ComputeInternal(OpKernelContext* context) const override;
+
+ private:
+ Status ComputeQsa(OpKernelContext* context) const;
+ Status ComputeCsa(OpKernelContext* context) const;
+
+ sparse_attention_indexer::Policy policy_;
+ int64_t compress_ratio_;
+ int64_t token_budget_;
+ int64_t index_topk_;
+ float epsilon_;
+ float scale_; // 0 means "derive from head_size"
+ float head_weight_scale_; // 0 means "derive from num_heads"
+};
+
+} // namespace cuda
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu
new file mode 100644
index 0000000000000..ea05afcd7886e
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu
@@ -0,0 +1,725 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+//
+// Correctness-first implementation of com.microsoft.SparseAttentionIndexer. Every stage is a
+// straightforward kernel that mirrors the reference semantics; see
+// docs/contrib_ops/cuda/sparse_attention_indexer.md for the performance follow-ups.
+
+#include "contrib_ops/cuda/sparse/sparse_attention_indexer_impl.h"
+
+#include
+#include
+#include
+#include
+
+#include
+
+#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+namespace {
+
+// The block reductions below halve the active thread count, so this must stay a power of two.
+constexpr int kThreads = 128;
+constexpr int64_t kMaxGridDimX = 2147483647;
+
+__device__ __forceinline__ float NegativeInfinity() { return -CUDART_INF_F; }
+
+int GridForElements(int64_t count) {
+ const int64_t blocks = (count + kThreads - 1) / kThreads;
+ return static_cast(std::clamp(blocks, 1, 65535));
+}
+
+// ---------------------------------------------------------------------------------------------
+// Shared device helpers
+// ---------------------------------------------------------------------------------------------
+
+__device__ __forceinline__ float BlockSum(float value, float* shared) {
+ shared[threadIdx.x] = value;
+ __syncthreads();
+ for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+ if (threadIdx.x < stride) {
+ shared[threadIdx.x] += shared[threadIdx.x + stride];
+ }
+ __syncthreads();
+ }
+ const float total = shared[0];
+ __syncthreads();
+ return total;
+}
+
+// Reduces (value, index) pairs to the largest value, breaking ties towards the smaller index.
+// A negative index marks an empty slot. shared_value/shared_index must already be filled and synced.
+__device__ __forceinline__ void BlockArgMax(float* shared_value, int* shared_index) {
+ for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+ if (threadIdx.x < stride) {
+ const int other_index = shared_index[threadIdx.x + stride];
+ if (other_index >= 0) {
+ const int this_index = shared_index[threadIdx.x];
+ const float other_value = shared_value[threadIdx.x + stride];
+ const float this_value = shared_value[threadIdx.x];
+ if (this_index < 0 || other_value > this_value ||
+ (other_value == this_value && other_index < this_index)) {
+ shared_value[threadIdx.x] = other_value;
+ shared_index[threadIdx.x] = other_index;
+ }
+ }
+ }
+ __syncthreads();
+ }
+}
+
+// Per-thread scan for the best entry that comes strictly after (previous_score, previous_index) in
+// the total order "score descending, then index ascending". Entries already emitted are therefore
+// skipped without needing a visited bitmap.
+__device__ __forceinline__ void ScanForNext(const float* scores, int count, float previous_score,
+ int previous_index, float* best_value, int* best_index) {
+ *best_index = -1;
+ *best_value = 0.0f;
+ for (int candidate = static_cast(threadIdx.x); candidate < count;
+ candidate += static_cast(blockDim.x)) {
+ const float value = scores[candidate];
+ if (previous_index >= 0 &&
+ !(value < previous_score || (value == previous_score && candidate > previous_index))) {
+ continue;
+ }
+ if (*best_index < 0 || value > *best_value ||
+ (value == *best_value && candidate < *best_index)) {
+ *best_value = value;
+ *best_index = candidate;
+ }
+ }
+}
+
+// Split-half rotary over the leading `rotary_width` channels (the convention used by the qsa
+// reference). Channels beyond `rotary_width` pass through unchanged.
+template
+__device__ __forceinline__ float LeadingRope(const float* value, int rotary_width, const T* cos_row,
+ const T* sin_row, int d) {
+ if (d >= rotary_width) {
+ return value[d];
+ }
+ const int half = rotary_width / 2;
+ const float paired = (d < half) ? -value[d + half] : value[d - half];
+ return value[d] * to_float(cos_row[d]) + paired * to_float(sin_row[d]);
+}
+
+// Interleaved rotary over the trailing 2 * rotary_width channels (the convention used by the csa
+// reference). Each cos/sin entry covers one channel pair, matching repeat_interleave(2).
+template
+__device__ __forceinline__ float TrailingRope(const float* value, int head_size, int rotary_width,
+ const T* cos_row, const T* sin_row, int d) {
+ const int base = head_size - 2 * rotary_width;
+ if (d < base) {
+ return value[d];
+ }
+ const int offset = d - base;
+ const float paired = ((offset & 1) == 0) ? -value[d + 1] : value[d - 1];
+ return value[d] * to_float(cos_row[offset >> 1]) + paired * to_float(sin_row[offset >> 1]);
+}
+
+// Highest compressed entry a query at `position` may attend to, matching (position + 1) // ratio.
+__device__ __forceinline__ int64_t CausalThreshold(int64_t position, int compress_ratio) {
+ return position < 0 ? 0 : (position + 1) / compress_ratio;
+}
+
+__device__ __forceinline__ int ClampPosition(int64_t position, int max_rotary_length) {
+ if (position < 0) {
+ return 0;
+ }
+ const int64_t limit = max_rotary_length - 1;
+ return static_cast(position < limit ? position : limit);
+}
+
+// ---------------------------------------------------------------------------------------------
+// Shared kernels
+// ---------------------------------------------------------------------------------------------
+
+// One block per (batch, token, head). kUseLeadingRope selects the qsa convention; otherwise the
+// csa convention with positions taken from position_ids.
+template
+__global__ void RotateQueryKernel(const T* query, const T* cos_cache, const T* sin_cache,
+ const int64_t* position_ids, float* query_rotated,
+ SparseAttentionIndexerParams params) {
+ extern __shared__ float shared[];
+ const int64_t rows = static_cast(params.batch_size) * params.sequence_length * params.num_heads;
+ for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+ const int token = static_cast((row / params.num_heads) % params.sequence_length);
+ const int batch = static_cast(row / (static_cast(params.num_heads) * params.sequence_length));
+ const int64_t base = row * params.head_size;
+
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ shared[d] = to_float(query[base + d]);
+ }
+ __syncthreads();
+
+ const int64_t raw_position = kUseLeadingRope
+ ? static_cast(params.past_sequence_length) + token
+ : position_ids[static_cast(batch) * params.sequence_length + token];
+ const int position = ClampPosition(raw_position, params.max_rotary_length);
+ const int64_t cache_offset =
+ (static_cast(batch) * params.max_rotary_length + position) * params.rotary_width;
+ const T* cos_row = cos_cache + cache_offset;
+ const T* sin_row = sin_cache + cache_offset;
+
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ query_rotated[base + d] =
+ kUseLeadingRope ? LeadingRope(shared, params.rotary_width, cos_row, sin_row, d)
+ : TrailingRope(shared, params.head_size, params.rotary_width, cos_row, sin_row, d);
+ }
+ __syncthreads();
+ }
+}
+
+// ---------------------------------------------------------------------------------------------
+// policy_mode = "qsa"
+// ---------------------------------------------------------------------------------------------
+
+template
+__global__ void ConcatPastKeyKernel(const T* past_key, const T* key, T* present_key,
+ SparseAttentionIndexerParams params) {
+ const int64_t total = static_cast(params.batch_size) * params.total_sequence_length * params.head_size;
+ for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total;
+ index += static_cast(gridDim.x) * blockDim.x) {
+ const int d = static_cast(index % params.head_size);
+ const int64_t token_row = index / params.head_size;
+ const int token = static_cast(token_row % params.total_sequence_length);
+ const int batch = static_cast(token_row / params.total_sequence_length);
+ present_key[index] =
+ token < params.past_sequence_length
+ ? past_key[(static_cast(batch) * params.past_sequence_length + token) * params.head_size + d]
+ : key[(static_cast(batch) * params.sequence_length + (token - params.past_sequence_length)) *
+ params.head_size +
+ d];
+ }
+}
+
+// One block per query row; compacts the visible key positions of that row into visible_indices.
+__global__ void CompactVisibleKernel(const bool* mask, int32_t* visible_indices, int32_t* visible_count,
+ SparseAttentionIndexerParams params) {
+ extern __shared__ int32_t shared_scan[];
+ const int64_t rows = static_cast(params.batch_size) * params.sequence_length;
+ for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+ 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) {
+ const int position = base + static_cast(threadIdx.x);
+ const int32_t flag = (position < params.total_sequence_length && mask_row[position]) ? 1 : 0;
+ shared_scan[threadIdx.x] = flag;
+ __syncthreads();
+ for (int stride = 1; stride < blockDim.x; stride <<= 1) {
+ const int32_t addend = (threadIdx.x >= static_cast(stride)) ? shared_scan[threadIdx.x - stride] : 0;
+ __syncthreads();
+ shared_scan[threadIdx.x] += addend;
+ __syncthreads();
+ }
+ if (flag != 0) {
+ out_row[offset + shared_scan[threadIdx.x] - 1] = position;
+ }
+ const int32_t tile_total = shared_scan[blockDim.x - 1];
+ __syncthreads();
+ offset += tile_total;
+ }
+ if (threadIdx.x == 0) {
+ visible_count[row] = offset;
+ }
+ __syncthreads();
+ }
+}
+
+// One block per (query row, block index). Pools compress_ratio visible keys, normalizes, rotates
+// and scores the result against every query head.
+template
+__global__ void QsaBlockScoreKernel(const T* present_key, const T* key_norm_weight, const T* cos_cache,
+ const T* sin_cache, const float* query_rotated,
+ const int32_t* visible_indices, const int32_t* visible_count,
+ float* block_scores, SparseAttentionIndexerParams params) {
+ extern __shared__ float shared[];
+ float* pooled = shared;
+ float* rotated = shared + params.head_size;
+ float* reduction = shared + 2 * params.head_size;
+
+ const int64_t total = static_cast(params.batch_size) * params.sequence_length * params.max_block_count;
+ for (int64_t work = blockIdx.x; work < total; work += gridDim.x) {
+ const int block_index = static_cast(work % params.max_block_count);
+ const int64_t row = work / params.max_block_count;
+ const int batch = static_cast(row / params.sequence_length);
+ const int block_count = visible_count[row] / params.compress_ratio;
+
+ if (block_index >= block_count) {
+ if (threadIdx.x == 0) {
+ block_scores[row * params.max_block_count + block_index] = NegativeInfinity();
+ }
+ continue;
+ }
+
+ const int32_t* index_row = visible_indices + row * params.total_sequence_length;
+ const int32_t* group = index_row + static_cast(block_index) * params.compress_ratio;
+ const int64_t key_base = static_cast(batch) * params.total_sequence_length * params.head_size;
+
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ float sum = 0.0f;
+ for (int t = 0; t < params.compress_ratio; ++t) {
+ sum += to_float(present_key[key_base + static_cast(group[t]) * params.head_size + d]);
+ }
+ pooled[d] = sum / static_cast(params.compress_ratio);
+ }
+ __syncthreads();
+
+ float sum_squares = 0.0f;
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ sum_squares += pooled[d] * pooled[d];
+ }
+ sum_squares = BlockSum(sum_squares, reduction);
+ const float inverse_rms = rsqrtf(sum_squares / static_cast(params.head_size) + params.epsilon);
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ pooled[d] = pooled[d] * inverse_rms * to_float(key_norm_weight[d]);
+ }
+ __syncthreads();
+
+ const int position = ClampPosition(group[0], params.max_rotary_length);
+ const int64_t cache_offset =
+ (static_cast(batch) * params.max_rotary_length + position) * params.rotary_width;
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ rotated[d] = LeadingRope(pooled, params.rotary_width, cos_cache + cache_offset,
+ sin_cache + cache_offset, d);
+ }
+ __syncthreads();
+
+ float score = 0.0f;
+ for (int head = 0; head < params.num_heads; ++head) {
+ const float* query_head = query_rotated + (row * params.num_heads + head) * params.head_size;
+ float partial = 0.0f;
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ partial += query_head[d] * rotated[d];
+ }
+ score += fmaxf(BlockSum(partial, reduction), 0.0f);
+ }
+
+ if (threadIdx.x == 0) {
+ block_scores[row * params.max_block_count + block_index] = score * params.scale;
+ }
+ __syncthreads();
+ }
+}
+
+// One block per query row. Emits the token indices of the highest scoring blocks followed by the
+// visible tokens of the trailing incomplete block.
+__global__ void QsaSelectKernel(const float* block_scores, const int32_t* visible_indices,
+ const int32_t* visible_count, int32_t* selected_indices,
+ SparseAttentionIndexerParams params) {
+ extern __shared__ float shared[];
+ float* shared_value = shared;
+ int* shared_index = reinterpret_cast(shared + blockDim.x);
+
+ const int64_t rows = static_cast(params.batch_size) * params.sequence_length;
+ for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+ int32_t* out_row = selected_indices + row * params.capacity;
+ for (int position = threadIdx.x; position < params.capacity; position += blockDim.x) {
+ out_row[position] = -1;
+ }
+ __syncthreads();
+
+ const int visible = visible_count[row];
+ const int block_count = visible / params.compress_ratio;
+ const int selected = min(params.block_topk, block_count);
+ const float* scores_row = block_scores + row * params.max_block_count;
+ const int32_t* index_row = visible_indices + row * params.total_sequence_length;
+
+ float previous_score = 0.0f;
+ int previous_index = -1;
+ int emitted = 0;
+ for (int rank = 0; rank < selected; ++rank) {
+ float best_value = 0.0f;
+ int best_index = -1;
+ ScanForNext(scores_row, block_count, previous_score, previous_index, &best_value, &best_index);
+ shared_value[threadIdx.x] = best_value;
+ shared_index[threadIdx.x] = best_index;
+ __syncthreads();
+ BlockArgMax(shared_value, shared_index);
+ previous_index = shared_index[0];
+ previous_score = shared_value[0];
+ __syncthreads();
+ if (previous_index < 0) {
+ break;
+ }
+ for (int t = threadIdx.x; t < params.compress_ratio; t += blockDim.x) {
+ out_row[rank * params.compress_ratio + t] = index_row[previous_index * params.compress_ratio + t];
+ }
+ emitted = rank + 1;
+ __syncthreads();
+ }
+
+ const int tail_start = block_count * params.compress_ratio;
+ for (int t = threadIdx.x; t < visible - tail_start; t += blockDim.x) {
+ out_row[emitted * params.compress_ratio + t] = index_row[tail_start + t];
+ }
+ __syncthreads();
+ }
+}
+
+// ---------------------------------------------------------------------------------------------
+// policy_mode = "csa"
+// ---------------------------------------------------------------------------------------------
+
+// Reads channel `channel` of token `position` of the virtual sequence [past buffer | new tokens].
+template
+__device__ __forceinline__ float ExtendedValue(const T* past_buffer, const T* current, int batch,
+ int position, int channel,
+ const SparseAttentionIndexerParams& params) {
+ const int width = 2 * params.head_size;
+ if (position < params.past_buffer_length) {
+ return to_float(
+ past_buffer[(static_cast(batch) * params.past_buffer_length + position) * width + channel]);
+ }
+ return to_float(
+ current[(static_cast(batch) * params.sequence_length + (position - params.past_buffer_length)) * width +
+ channel]);
+}
+
+// One block per (batch, new window). Softmax-pools the 2 * compress_ratio window slots, normalizes
+// and rotates the result into present_compressed_key.
+template
+__global__ void CsaCompressKernel(const T* key, const T* gate, const T* past_kv_buffer,
+ const T* past_gate_buffer, const T* position_bias,
+ const T* key_norm_weight, const T* cos_cache, const T* sin_cache,
+ T* present_compressed_key, SparseAttentionIndexerParams params) {
+ extern __shared__ float shared[];
+ float* pooled = shared;
+ float* reduction = shared + params.head_size;
+
+ const int width = 2 * params.head_size;
+ const int64_t total = static_cast(params.batch_size) * params.new_window_count;
+ for (int64_t work = blockIdx.x; work < total; work += gridDim.x) {
+ const int window = static_cast(work % params.new_window_count);
+ const int batch = static_cast(work / params.new_window_count);
+ const bool has_previous = window >= 1 || params.overlap_length >= params.compress_ratio;
+ const int previous_base = params.overlap_length + (window - 1) * params.compress_ratio;
+ const int current_base = params.overlap_length + window * params.compress_ratio;
+
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ float max_gate = NegativeInfinity();
+ if (has_previous) {
+ for (int slot = 0; slot < params.compress_ratio; ++slot) {
+ const float value =
+ ExtendedValue(past_gate_buffer, gate, batch, previous_base + slot, d, params) +
+ to_float(position_bias[static_cast(slot) * width + d]);
+ max_gate = fmaxf(max_gate, value);
+ }
+ }
+ for (int slot = 0; slot < params.compress_ratio; ++slot) {
+ const float value =
+ ExtendedValue(past_gate_buffer, gate, batch, current_base + slot, params.head_size + d, params) +
+ to_float(position_bias[static_cast(slot) * width + params.head_size + d]);
+ max_gate = fmaxf(max_gate, value);
+ }
+
+ float denominator = 0.0f;
+ float accumulator = 0.0f;
+ if (has_previous) {
+ for (int slot = 0; slot < params.compress_ratio; ++slot) {
+ const float logit =
+ ExtendedValue(past_gate_buffer, gate, batch, previous_base + slot, d, params) +
+ to_float(position_bias[static_cast(slot) * width + d]);
+ const float weight = __expf(logit - max_gate);
+ denominator += weight;
+ accumulator += weight * ExtendedValue(past_kv_buffer, key, batch, previous_base + slot, d, params);
+ }
+ }
+ for (int slot = 0; slot < params.compress_ratio; ++slot) {
+ const float logit =
+ ExtendedValue(past_gate_buffer, gate, batch, current_base + slot, params.head_size + d, params) +
+ to_float(position_bias[static_cast(slot) * width + params.head_size + d]);
+ const float weight = __expf(logit - max_gate);
+ denominator += weight;
+ accumulator +=
+ weight * ExtendedValue(past_kv_buffer, key, batch, current_base + slot, params.head_size + d, params);
+ }
+ pooled[d] = denominator > 0.0f ? accumulator / denominator : 0.0f;
+ }
+ __syncthreads();
+
+ float sum_squares = 0.0f;
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ sum_squares += pooled[d] * pooled[d];
+ }
+ sum_squares = BlockSum(sum_squares, reduction);
+ const float inverse_rms = rsqrtf(sum_squares / static_cast(params.head_size) + params.epsilon);
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ pooled[d] = pooled[d] * inverse_rms * to_float(key_norm_weight[d]);
+ }
+ __syncthreads();
+
+ const int64_t entry = static_cast(params.past_compressed_length) + window;
+ const int position = ClampPosition(entry * params.compress_ratio, params.max_rotary_length);
+ const int64_t cache_offset =
+ (static_cast(batch) * params.max_rotary_length + position) * params.rotary_width;
+ const int64_t out_base =
+ (static_cast(batch) * params.present_compressed_length + entry) * params.head_size;
+ for (int d = threadIdx.x; d < params.head_size; d += blockDim.x) {
+ present_compressed_key[out_base + d] = from_float(TrailingRope(
+ pooled, params.head_size, params.rotary_width, cos_cache + cache_offset, sin_cache + cache_offset, d));
+ }
+ __syncthreads();
+ }
+}
+
+template
+__global__ void CsaCopyPastCompressedKernel(const T* past_compressed_key, T* present_compressed_key,
+ SparseAttentionIndexerParams params) {
+ const int64_t total = static_cast(params.batch_size) * params.past_compressed_length * params.head_size;
+ for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total;
+ index += static_cast(gridDim.x) * blockDim.x) {
+ const int64_t entry_row = index / params.head_size;
+ const int entry = static_cast(entry_row % params.past_compressed_length);
+ const int batch = static_cast(entry_row / params.past_compressed_length);
+ const int d = static_cast(index % params.head_size);
+ present_compressed_key[(static_cast(batch) * params.present_compressed_length + entry) * params.head_size +
+ d] = past_compressed_key[index];
+ }
+}
+
+template
+__global__ void CsaCopyBufferKernel(const T* key, const T* gate, const T* past_kv_buffer,
+ const T* past_gate_buffer, T* present_kv_buffer, T* present_gate_buffer,
+ SparseAttentionIndexerParams params) {
+ const int width = 2 * params.head_size;
+ const int64_t total = static_cast(params.batch_size) * params.present_buffer_length * width;
+ for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total;
+ index += static_cast(gridDim.x) * blockDim.x) {
+ const int channel = static_cast(index % width);
+ const int64_t token_row = index / width;
+ const int token = static_cast(token_row % params.present_buffer_length);
+ const int batch = static_cast(token_row / params.present_buffer_length);
+ const int source = params.present_buffer_start + token;
+ present_kv_buffer[index] = from_float(ExtendedValue(past_kv_buffer, key, batch, source, channel, params));
+ present_gate_buffer[index] =
+ from_float(ExtendedValue(past_gate_buffer, gate, batch, source, channel, params));
+ }
+}
+
+// One thread per (query row, compressed entry). Also applies the causal mask so that the selection
+// kernel only has to read scores.
+template
+__global__ void CsaScoreKernel(const float* query_rotated, const T* present_compressed_key,
+ const T* head_weights, const int64_t* position_ids, float* scores,
+ SparseAttentionIndexerParams params) {
+ const int64_t total = static_cast(params.batch_size) * params.sequence_length *
+ params.present_compressed_length;
+ for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total;
+ index += static_cast(gridDim.x) * blockDim.x) {
+ const int entry = static_cast(index % params.present_compressed_length);
+ const int64_t row = index / params.present_compressed_length;
+ const int batch = static_cast(row / params.sequence_length);
+
+ const int64_t threshold = CausalThreshold(position_ids[row], params.compress_ratio);
+ if (static_cast