diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index e39641d06dc1b..3c1afcc34ed24 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -4991,7 +4991,7 @@ This version of the operator has been available since version 1 of the 'com.micr
bias (optional) : T
Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection
token_offset : M
-Offset of each token before packing, with shape (batch_size, sequence_length). Values must be in [0, batch_size * sequence_length).
+Offset of each token before packing, with shape (batch_size, sequence_length).
cumulative_sequence_length : M
A tensor with shape (batch_size + 1). It specifies the cumulative sequence length.
attention_bias (optional) : T
@@ -5046,8 +5046,8 @@ This version of the operator has been available since version 1 of the 'com.micr
sum_h ReLU(q_h . k), the token_budget / compress_ratio highest scoring blocks are kept, and
their token indices are emitted (request-local logical positions, i.e. the same numbering as
past_sequence_lengths + local offset) followed by the causally visible tokens of the trailing
- incomplete block. QSA positions are always the request-local logical cache positions derived
- from past_sequence_lengths and cumulative_sequence_lengths; position_ids must be omitted.
+ incomplete block. QSA positions are the request-local logical cache positions derived from
+ past_sequence_lengths and cumulative_sequence_lengths; position_ids must be omitted.
policy_mode = "csa" ("compressed sparse attention" block indexer)
Applies the same window-plan arithmetic as SparseAttentionIndexer (overlap/leftover/new window
@@ -5104,17 +5104,19 @@ This version of the operator has been available since version 1 of the 'com.micr
Scale applied to the per-head ReLU scores. Default is 1/sqrt(head_size).
state_capacity : int (required)
Fixed capacity (number of entries) of past_key_state / present_key_state. Must be > 0.
+state_update_capacity : int
+Only for policy_mode 'qsa': maximum number of leading token transitions captured per request. Must be in [0, 8]. Default is 0.
token_budget : int
Only for policy_mode 'qsa': maximum number of tokens selected from complete blocks. Must be > 0 and divisible by compress_ratio. Must be omitted when policy_mode is 'csa'.
-#### Inputs
+#### Inputs (16 - 18)
- query : T
-- Packed indexer queries with shape (total_tokens, num_heads * head_size), before normalization, logical reshape, and rotary embedding.
-- key : T
-- Packed indexer key projection of the new tokens. Shape is (total_tokens, head_size) for policy_mode 'qsa' and (total_tokens, 2 * head_size) for policy_mode 'csa', where the first head_size channels are the Ca series and the last head_size channels the Cb series.
+- Packed indexer queries with shape (total_tokens, num_heads * head_size), before normalization, logical reshape, and rotary embedding. For policy_mode 'qsa', key may be omitted and query then contains row-wise concatenated query and key projections with shape (total_tokens, (num_heads + 1) * head_size).
+- key (optional) : T
+- Packed indexer key projection of the new tokens. Shape is (total_tokens, head_size) for policy_mode 'qsa' and (total_tokens, 2 * head_size) for policy_mode 'csa', where the first head_size channels are the Ca series and the last head_size channels the Cb series. May be omitted for policy_mode 'qsa' when query contains the packed query/key projection.
- query_norm_weight : T
- Effective RMSNorm multiplier of the queries, with shape (head_size).
- key_norm_weight : T
@@ -5143,9 +5145,13 @@ This version of the operator has been available since version 1 of the 'com.micr
- Only for policy_mode 'csa': buffered gate projections with the same shape as past_kv_buffer.
- past_state_lengths : M
- Generic per-request state length with shape (batch_size, 2). Column 0 is the key_state entry count (policy_mode 'qsa': complete-block count; 'csa': compressed-entry count); column 1 is the pending-buffer length (policy_mode 'qsa': incomplete-block length in [0, compress_ratio); 'csa': buffer length in [0, 2 * compress_ratio)).
+- state_update_capture_count (optional) : M
+- Only for policy_mode 'qsa': number of leading token transitions to capture for each request, with shape (batch_size). Values are clamped to the request length and state_update_capacity. Required when state_update_capacity is positive.
+- state_update_active (optional) : M
+- Only for policy_mode 'qsa': optional capture gate with shape (1). A zero value disables capture.
-#### Outputs
+#### Outputs (6 - 7)
- selected_indices : M
@@ -5160,6 +5166,8 @@ This version of the operator has been available since version 1 of the 'com.micr
- Only for policy_mode 'csa': updated gate buffer with the same fixed shape as past_gate_buffer.
- present_state_lengths : M
- Updated generic per-request state length, with the same fixed shape as past_state_lengths.
+- state_update (optional) : T
+- Only for policy_mode 'qsa': compact transition payloads with shape (batch_size, state_update_capacity, head_size). A token that completes a compression block stores the prepared block representative; any other captured token stores its raw key. Inactive and unused slots are zero.
#### Type Constraints
@@ -6943,7 +6951,7 @@ This version of the operator has been available since version 1 of the 'com.micr
input : T
Input tensor with shape (total_tokens, hidden_size)
token_offset : M
-Offset of non-padding tokens and paddings. Its shape is (batch_size, sequence_length), and values must be in [0, batch_size * sequence_length).
+Offset of non-padding tokens and paddings. Its shape is (batch_size, sequence_length)
#### Outputs
@@ -7616,11 +7624,11 @@ This version of the operator has been available since version 1 of the 'com.micr
key_norm_weight : T
Effective RMSNorm multiplier of the compressed keys, with shape (head_size).
cos_cache : T
-Cosine rotary table indexed by absolute key position, with shape (batch_size, max_rotary_sequence_length, rotary_width).
+Cosine rotary table indexed by absolute key position, shared across the batch with shape (max_rotary_sequence_length, rotary_width) or request-specific with shape (batch_size, max_rotary_sequence_length, rotary_width).
sin_cache : T
Sine rotary table with the same shape as cos_cache.
mask (optional) : TB
-Only for policy_mode 'qsa': INT64 padding mask with shape (batch_size, total_sequence_length). Nonzero entries are visible subject to causal masking. total_sequence_length is past_sequence_length + sequence_length.
+Only for policy_mode 'qsa': INT64 padding mask with shape (batch_size, total_sequence_length). Nonzero entries are visible subject to causal masking. total_sequence_length is past_sequence_length + sequence_length. When omitted, every position through past_sequence_length plus the current query index is visible.
past_key : T
Cached indexer keys. For policy_mode 'qsa', these are raw keys; for 'csa', they are compressed keys. Shape is (batch_size, past_sequence_length, head_size), or (batch_size, max_cache_length, head_size) when a valid past_sequence_length is provided.
gate (optional) : T
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 10a1b5b96072b..96ac677d0966d 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -1131,7 +1131,7 @@ The **OpSet Version** column uses the following notation:
|NhwcConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|PackedAttention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* attention_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|PackedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* attention_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
-|PackedSparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* cumulative_sequence_lengths:**M**
*in* past_sequence_lengths:**M**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_key_state:**T**
*in* past_kv_buffer:**T**
*in* past_gate_buffer:**T**
*in* past_state_lengths:**M**
*out* selected_indices:**M**
*out* selected_counts:**M**
*out* present_key_state:**T**
*out* present_kv_buffer:**T**
*out* present_gate_buffer:**T**
*out* present_state_lengths:**M**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)|
+|PackedSparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* cumulative_sequence_lengths:**M**
*in* past_sequence_lengths:**M**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_key_state:**T**
*in* past_kv_buffer:**T**
*in* past_gate_buffer:**T**
*in* past_state_lengths:**M**
*in* state_update_capture_count:**M**
*in* state_update_active:**M**
*out* selected_indices:**M**
*out* selected_counts:**M**
*out* present_key_state:**T**
*out* present_kv_buffer:**T**
*out* present_gate_buffer:**T**
*out* present_state_lengths:**M**
*out* state_update:**T**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|PagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T_CACHE**
*in* value_cache:**T_CACHE**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* slot_mapping:**S**
*in* head_sink:**T**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* attention_metadata:**S**
*out* output:**T**
*out* key_cache_out:**T_CACHE**
*out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(int8), tensor(uint8)
**T_KV_SCALE** = tensor(float)|
|QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(int8)
**T2** = tensor(int8)
**T3** = tensor(float), tensor(float16)
**T4** = tensor(int32)|
|QMoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T1**
*in* fc1_scales:**T2**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T1**
*in* fc2_scales:**T2**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T1**
*in* fc3_scales:**T2**
*in* fc3_experts_bias:**T**
*in* fc1_zero_points:**T1**
*in* fc2_zero_points:**T1**
*in* fc3_zero_points:**T1**
*in* router_weights:**T**
*in* fc1_global_scale:**T4**
*in* fc2_global_scale:**T4**
*in* fc1_act_scale:**T4**
*in* fc2_act_scale:**T4**
*in* fc1_act_block_scale:**T2**
*in* fc2_act_block_scale:**T2**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(float8e8m0)
**T4** = tensor(float)|
diff --git a/docs/contrib_ops/cuda/sparse_attention_indexer.md b/docs/contrib_ops/cuda/sparse_attention_indexer.md
index acc9bd448a143..a1bfe8634d170 100644
--- a/docs/contrib_ops/cuda/sparse_attention_indexer.md
+++ b/docs/contrib_ops/cuda/sparse_attention_indexer.md
@@ -84,7 +84,7 @@ for a runtime:
| 1 | `key` | both | `T` | `(B, S, D)` for `qsa`, `(B, S, 2D)` for `csa` |
| 2 | `query_norm_weight` | both | `T` | `(D)` |
| 3 | `key_norm_weight` | both | `T` | `(D)` |
-| 4 | `cos_cache` | both | `T` | `(B, max_rotary_sequence_length, R)` |
+| 4 | `cos_cache` | both | `T` | `(max_rotary_sequence_length, R)` when shared across the batch, or `(B, max_rotary_sequence_length, R)` |
| 5 | `sin_cache` | both | `T` | same as `cos_cache` |
| 6 | `mask` | `qsa` | `TB` | INT64 padding mask `(B, T)` |
| 7 | `past_key` | both | `T` | `(B, P, D)`; raw keys for `qsa`, compressed keys for `csa` |
@@ -261,11 +261,11 @@ that assign one block to each work item clamp the grid to the CUDA `gridDim.x` l
| Stage | Kernel | Parallelism |
|---|---|---|
-| 1 | `CopyQsaPastKeyKernel` / `AppendQsaKeyKernel` | element |
+| 1 | `CopyQsaPastKeyKernel` / `AppendQsaKeyKernel` | element; the copy is skipped for aliased fixed-capacity caches |
| 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 |
+| 3 | `InitializePrefixVisibleCountKernel` or mask analysis/compaction | omitted masks initialize prefix counts directly; irregular masks use one block per `(b, s)` |
+| 4 | `QsaBlockScoreQwenKernel` / `QsaBlockScoreKernel` | one block per `(b, s, block)`; the Qwen `D=128, r=4, N=4` path assigns one warp per query head |
+| 5 | `QsaPartialTopKKernel` / `QsaSelectKernel` | one block per `(b, s)`; score-threshold radix selection and bounded winner sorting |
### `csa`
@@ -283,18 +283,19 @@ that assign one block to each work item clamp the grid to the CUDA `gridDim.x` l
| 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) |
+| `qsa` | int32 | optional `B*S*T` (arbitrary-mask visible indices) + `B*S` (visible counts) + up to `B*S*block_topk` (TopK indices) |
| `csa` | float | `B*S*N*D` (rotated query) + `B*S*present_compressed_length` (scores) + `B*W*D` (new compressed keys before persistent-cache rounding) |
Every size is derived from shapes and attributes only.
-### Selection without a visited bitmap
+### Selection
-`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.
+For QSA `block_topk` values from 33 through 512, `QsaPartialTopKKernel` finds the score threshold
+with four byte-wise radix-histogram passes, gathers exactly the winning candidates, resolves
+threshold ties by ascending index, and sorts only that bounded set. Selection storage and output
+traffic are therefore `O(block_topk)`, rather than `O(entries)`, while preserving deterministic
+"score descending, index ascending" ordering. Smaller QSA rows use the in-kernel block TopK;
+larger configured TopK values retain the repeated-scan fallback.
## 8. Numerics and Determinism
@@ -322,7 +323,8 @@ The CUDA kernel additionally validates concrete runtime dimensions and cross-inp
- `index_topk` absent or `<= 0` for `csa`;
- an output count other than 2 (`qsa`) or 3 (`csa`);
- a `past_proj_buffer` length outside `[0, 2 * compress_ratio)`;
-- rank mismatches and incompatible concrete dimensions between `query`, `key`, the caches and the buffers.
+- rank or dimension mismatches between `query`, `key`, the caches and the buffers. Rotary caches
+ may be shared rank-2 tables or request-specific rank-3 tables.
Schema and statically knowable errors fail at `Graph::Resolve()`; checks that require concrete runtime dimensions fail
before any CUDA kernel is launched.
@@ -341,7 +343,12 @@ contains:
- **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.
+ sees exactly the values the kernel reads. Shared rank-2 rotary caches are covered for both CUDA
+ and WebGPU. These tests skip when the requested EP is unavailable.
+- **Long-context parity tests** at 8K, 32K, 64K, 128K and 256K for
+ `compress_ratio=4, token_budget=2048`.
+- **Fixed-address CUDA graph capture and replay** in
+ `onnxruntime/test/contrib_ops/cuda_kernels/sparse_attention_indexer_kernel_test.cc`.
Run them with:
@@ -349,21 +356,28 @@ Run them with:
./build/Linux/Release/onnxruntime_provider_test --gtest_filter='SparseAttentionIndexer*'
```
+Run the CUDA decode sweep against a built Python package with:
+
+```bash
+python onnxruntime/python/tools/microbench/sparse_attention_indexer.py
+```
+
## 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 each batch's `(B, T)` padding
- mask for every `s`. Sharing work across query rows would reduce that traffic for long contexts.
-5. **`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.
+1. **Exact query-specific scoring remains `O(visible_tokens / compress_ratio)`.** Every decode query
+ can change, so an exact implementation must inspect every eligible block; persistent candidates
+ from an earlier query cannot preserve QSA semantics. The bounded radix path removes full-context
+ sorting and candidate buffers, but not this scoring lower bound.
+2. **Generic scoring is not tensor-core accelerated.** The common Qwen `D=128, r=4, N=4` path uses
+ vectorized loads and one warp per query head. Other QSA shapes and `CsaScoreKernel` retain the
+ generic shared-memory reductions.
+3. **Rotated queries use an FP32 workspace.** Query normalization and rotation are computed once
+ per `(b, s, h)` before block scoring instead of being recomputed for every candidate block.
+4. **Irregular-mask analysis remains `O(T)` per query row.** Omitting the mask declares
+ prefix-causal visibility and skips analysis and compaction. Padding masks with holes still
+ require exact compaction because block membership is defined by visible rank, not position.
+5. **Only the common Qwen dimensions are specialized.** Other `compress_ratio`, `head_size`, and
+ head-count combinations retain dynamic loop bounds in the generic scoring kernel.
6. **No CPU kernel.** The operator is CUDA-only today.
diff --git a/docs/contrib_ops/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md
index b6d4c0b3fefed..bd6877cf32189 100644
--- a/docs/contrib_ops/packed_sparse_attention_indexer.md
+++ b/docs/contrib_ops/packed_sparse_attention_indexer.md
@@ -10,7 +10,7 @@ Source:
(policy enum, selected-capacity formula and CSA window-plan arithmetic, shared unmodified with
`SparseAttentionIndexer`),
[packed_sparse_attention_indexer_common.h](../../onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h)
-(fixed 16-input / 6-output slot map),
+(fixed 18-input / 7-output slot map),
[packed_sparse_attention_indexer.cc](../../onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc) /
[packed_sparse_attention_indexer_impl.cu](../../onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu)
(CUDA),
@@ -51,16 +51,17 @@ Attributes:
| `epsilon` | default `1e-6` | RMSNorm epsilon |
| `scale` | default `1/sqrt(head_size)` | per-head score scale |
| `head_weight_scale` | `csa` only, default `1/sqrt(num_heads)` | head-weight score scale |
+| `state_update_capacity` | `qsa` only, default `0`, range `[0, 8]` | maximum captured transitions per request |
-Inputs are **fixed at 16 indices** for both policies (unlike the dense op, which uses a different
+Inputs are **fixed at 18 indices** for both policies (unlike the dense op, which uses a different
input/output count per policy). A slot not owned by the active policy is a *positional* optional:
its `NodeProto` input name is empty rather than the slot being removed from the list, so every
later slot keeps its fixed index.
| # | Name | Shape | Type | Policy |
|---|---|---|---|---|
-| 0 | `query` | `(total_tokens, num_heads*head_size)` | T | both |
-| 1 | `key` | `(total_tokens, head_size)` qsa / `(total_tokens, 2*head_size)` csa | T | both |
+| 0 | `query` | `(total_tokens, num_heads*head_size)`; when qsa omits `key`, `(total_tokens, (num_heads+1)*head_size)` containing row-wise `[query \| key]` | T | both |
+| 1 | `key` | `(total_tokens, head_size)` qsa / `(total_tokens, 2*head_size)` csa; optional for qsa when packed into `query` | T | both |
| 2 | `query_norm_weight` | `(head_size)` | T | both |
| 3 | `key_norm_weight` | `(head_size)` | T | both |
| 4 | `cos_cache` | `(max_position, rotary_width)` or `(batch_size, max_position, rotary_width)` | T | both |
@@ -75,9 +76,11 @@ later slot keeps its fixed index.
| 13 | `past_kv_buffer` | `(batch_size, 2*compress_ratio-1, width)` | T | both (generic) |
| 14 | `past_gate_buffer` | same shape as `past_kv_buffer` | T | csa only |
| 15 | `past_state_lengths` | `(batch_size, 2)` | int32, device-resident | both (generic) |
+| 16 | `state_update_capture_count` | `(batch_size)` | int32, device-resident | qsa capture only |
+| 17 | `state_update_active` | `(1)` | int32, device-resident | optional qsa capture gate |
-Outputs are **fixed at 6 indices** for both policies (`present_gate_buffer` is declared with an
-empty output name for `qsa`, the same positional-optional convention as above):
+Outputs have six fixed indices plus optional slot 6. `present_gate_buffer` is declared with an
+empty output name for `qsa`; `state_update` is declared only when QSA capture is enabled:
| # | Name | Shape | Type | Policy |
|---|---|---|---|---|
@@ -87,6 +90,7 @@ empty output name for `qsa`, the same positional-optional convention as above):
| 3 | `present_kv_buffer` | same shape as `past_kv_buffer` | T | both |
| 4 | `present_gate_buffer` | same shape as `past_gate_buffer` | T | csa only |
| 5 | `present_state_lengths` | same shape as `past_state_lengths` | int32 | both |
+| 6 | `state_update` | `(batch_size, state_update_capacity, head_size)` | T | qsa capture only |
`capacity` is `token_budget + compress_ratio - 1` for `qsa` and `index_topk` for `csa`, exactly the
same formula (`SelectedCapacity`) used by `SparseAttentionIndexer`.
@@ -118,16 +122,26 @@ Both policies read and write the *same four* state slots — there is no separat
`CsaWindowPlan`/`TryComputeCsaWindowPlan`).
State never grows. `present_*` always has exactly the same shape as `past_*`; only the *contents*
-change. Input/output aliasing is supported. CUDA avoids unsafe buffer aliases; WebGPU omits an
-aliased `past_*` read-only binding and reads the prior contents through the matching read-write
-`present_*` binding. Each request is handled by one invocation, and buffer compaction reads entries
-at or above the destination index before overwriting them.
+change. CUDA supports aliasing `past_key_state` with `present_key_state` and `past_state_lengths`
+with `present_state_lengths`; it skips the baseline key-state copy when those buffers alias.
+WebGPU supports aliasing every corresponding `past_*` / `present_*` pair and reads prior contents
+through the matching read-write `present_*` binding. Each request is handled by one invocation,
+and buffer compaction reads entries at or above the destination index before overwriting them.
**State overflow.** If a call would close more blocks/windows than
`state_capacity - old_entry_count` allows, that request's step is rejected as a deterministic
no-op: its state and state lengths remain unchanged, and its selection outputs stay empty. This
never reads or writes outside a tensor's fixed extent and never silently truncates state.
+**Transactional QSA capture.** When `state_update_capacity > 0`,
+`state_update_capture_count[b]` requests capture of that request's leading transitions, clamped to
+the request length and the configured capacity. `state_update_active == 0` disables capture for the
+call. Disabled, unused, and out-of-range slots are zero. A captured token that completes a
+compression block stores the prepared block representative; other tokens store their raw key.
+After restoring a pre-call state snapshot, a runtime can replay the accepted prefix in order:
+append raw entries to the pending buffer, and replace each block-completion transition with its
+prepared representative in `key_state`.
+
## 4. Packed metadata and device-side safety
`cumulative_sequence_lengths` and `past_sequence_lengths` are **device-resident** tensors, read
@@ -251,19 +265,22 @@ mechanics with no packed equivalent):
`onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc` covers:
- shape inference for `qsa` and `csa` (fixed `selected_indices`/`selected_counts` shapes, fixed
- state output shapes, the strict per-slot policy validation, and the always-6-outputs contract);
+ state output shapes, strict per-slot policy validation, and optional capture output);
- multi-request packed batches with unequal token counts, a zero-token request row, and prefill
followed by decode with independent per-request state (CUDA/WebGPU, skipped without the
respective execution provider);
- `qsa` state-capacity overflow safety;
+- active and disabled transactional QSA capture on CUDA and WebGPU;
- FP32/FP16 (and CUDA-only BF16) numeric coverage against an in-file reference that mirrors this
document's contract.
## 10. Known limitations and follow-ups
- The reference CUDA/WebGPU kernels prioritize correctness over throughput (see the top-of-file
- comments in the `.cu`/`.cc` implementations); they are not yet tuned for large `state_capacity`
- or long packed batches.
+ comments in the `.cu`/`.cc` implementations). CUDA QSA fuses query rotation into scoring, uses
+ the shared-memory TopK path for up to 32 selected blocks, bounded partial TopK for 33–512 blocks,
+ and a hierarchical TopK path for common long-context Qwen shapes. Values above 512 that do not
+ match the hierarchical specialization retain the correctness-first repeated-scan fallback.
- OgaEngine / Model Builder integration (declaring `past_key_state` etc. as Engine-managed,
per-request fixed-size state, analogous to a paged auxiliary cache) is out of scope for this
operator definition and is expected in a follow-up to `microsoft/onnxruntime-genai`.
diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h
index c457179e40a8c..a5b2be971c8aa 100644
--- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h
+++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h
@@ -51,7 +51,9 @@ enum InputIndex : int {
kPastKvBuffer = 13, // generic: [batch_size, 2 * compress_ratio - 1, width]
kPastGateBuffer = 14, // csa only: same shape as past_kv_buffer
kPastStateLengths = 15, // generic: [batch_size, 2], int32
- kInputCount = 16,
+ kStateUpdateCaptureCount = 16, // qsa only, optional: [batch_size], int32
+ kStateUpdateActive = 17, // qsa only, optional: [1], int32
+ kInputCount = 18,
};
// Fixed output slots. present_gate_buffer is declared (with an empty name) but not produced for
@@ -63,12 +65,14 @@ enum OutputIndex : int {
kPresentKvBuffer = 3, // same shape as past_kv_buffer
kPresentGateBuffer = 4, // csa only: same shape as past_gate_buffer
kPresentStateLengths = 5, // [batch_size, 2], int32
- kOutputCount = 6,
+ kStateUpdate = 6, // qsa only, optional: [batch_size, state_update_capacity, head_size]
+ kOutputCount = 7,
};
// Every PackedSparseAttentionIndexer node declares all 6 fixed outputs; present_gate_buffer is an
-// empty-name optional output for policy_mode="qsa".
-constexpr int kFixedOutputCount = kOutputCount;
+// empty-name optional output for policy_mode="qsa". State update capture is an optional appended
+// seventh output so existing nodes retain their output layout.
+constexpr int kFixedOutputCount = kStateUpdate;
// Column layout of past_state_lengths / present_state_lengths.
enum StateLengthColumn : int {
diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc
index 0e9e3e2d8325c..4e2bed76767d6 100644
--- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc
+++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc
@@ -19,22 +19,22 @@ namespace cuda {
using namespace onnxruntime::cuda;
namespace psai = onnxruntime::contrib::packed_sparse_attention_indexer;
-#define REGISTER_KERNEL_TYPED(T) \
- ONNX_OPERATOR_TYPED_KERNEL_EX( \
- PackedSparseAttentionIndexer, \
- kMSDomain, \
- 1, \
- T, \
- kCudaExecutionProvider, \
- (*KernelDefBuilder::Create()) \
- .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
- .TypeConstraint("I", DataTypeImpl::GetTensorType()) \
- .TypeConstraint("M", DataTypeImpl::GetTensorType()) \
- .MayInplace(psai::kPastKeyState, psai::kPresentKeyState) \
- .MayInplace(psai::kPastKvBuffer, psai::kPresentKvBuffer) \
- .MayInplace(psai::kPastGateBuffer, psai::kPresentGateBuffer) \
- .MayInplace(psai::kPastStateLengths, \
- psai::kPresentStateLengths), \
+#define REGISTER_KERNEL_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ PackedSparseAttentionIndexer, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCudaExecutionProvider, \
+ (*KernelDefBuilder::Create()) \
+ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("I", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("M", DataTypeImpl::GetTensorType()) \
+ .InputMemoryType(OrtMemTypeCPUInput, psai::kStateUpdateActive) \
+ .MayInplace(psai::kPastKeyState, psai::kPresentKeyState) \
+ .MayInplace(psai::kPastKvBuffer, psai::kPresentKvBuffer) \
+ .MayInplace(psai::kPastGateBuffer, psai::kPresentGateBuffer) \
+ .MayInplace(psai::kPastStateLengths, psai::kPresentStateLengths), \
PackedSparseAttentionIndexer);
REGISTER_KERNEL_TYPED(float)
@@ -110,6 +110,10 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo
"PackedSparseAttentionIndexer: state_capacity is required");
ORT_ENFORCE(state_capacity_ > 0 && state_capacity_ <= std::numeric_limits::max(),
"PackedSparseAttentionIndexer: state_capacity must be in (0, INT_MAX], got ", state_capacity_);
+ state_update_capacity_ = info.GetAttrOrDefault("state_update_capacity", 0);
+ ORT_ENFORCE(state_update_capacity_ >= 0 && state_update_capacity_ <= 8,
+ "PackedSparseAttentionIndexer: state_update_capacity must be in [0, 8], got ",
+ state_update_capacity_);
const bool has_token_budget = info.GetAttr("token_budget", &token_budget_).IsOK();
const bool has_index_topk = info.GetAttr("index_topk", &index_topk_).IsOK();
@@ -129,6 +133,8 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo
token_budget_, " compress_ratio=", compress_ratio_);
index_topk_ = 0;
} else {
+ ORT_ENFORCE(state_update_capacity_ == 0,
+ "PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'");
ORT_ENFORCE(has_index_topk, "PackedSparseAttentionIndexer: index_topk is required when policy_mode is 'csa'");
ORT_ENFORCE(!has_token_budget,
"PackedSparseAttentionIndexer: token_budget must not be set when policy_mode is 'csa'");
@@ -155,14 +161,22 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context
const bool position_ids_provided =
psai::kPositionIds < context->InputCount() && context->Input(psai::kPositionIds) != nullptr;
ORT_RETURN_IF(position_ids_provided == is_qsa,
- "PackedSparseAttentionIndexer: input ", psai::kPositionIds,
- is_qsa ? " (position_ids) must be omitted for policy_mode 'qsa'"
- : " (position_ids) is required for policy_mode 'csa'");
+ "PackedSparseAttentionIndexer: position_ids is required for "
+ "policy_mode 'csa' and must be omitted for policy_mode 'qsa'");
const bool gate_buffer_provided =
psai::kPastGateBuffer < context->InputCount() && context->Input(psai::kPastGateBuffer) != nullptr;
ORT_RETURN_IF(gate_buffer_provided != !is_qsa, "PackedSparseAttentionIndexer: input ", psai::kPastGateBuffer,
gate_buffer_provided ? " must be omitted for policy_mode 'qsa'"
: " is required for policy_mode 'csa'");
+ if (!is_qsa) {
+ for (int index : {psai::kStateUpdateCaptureCount, psai::kStateUpdateActive}) {
+ ORT_RETURN_IF(index < context->InputCount() && context->Input(index) != nullptr,
+ "PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'");
+ }
+ Tensor* state_update = context->Output(psai::kStateUpdate, TensorShape({0}));
+ ORT_RETURN_IF(state_update != nullptr,
+ "PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'");
+ }
return is_qsa ? ComputeQsa(context) : ComputeCsa(context);
}
@@ -179,9 +193,20 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con
const Tensor* sin_cache = context->Input(psai::kSinCache);
const Tensor* cumulative_sequence_lengths = context->Input(psai::kCumulativeSequenceLengths);
const Tensor* past_sequence_lengths = context->Input(psai::kPastSequenceLengths);
+ const Tensor* position_ids = context->Input(psai::kPositionIds);
const Tensor* past_key_state = context->Input(psai::kPastKeyState);
const Tensor* past_kv_buffer = context->Input(psai::kPastKvBuffer);
const Tensor* past_state_lengths = context->Input(psai::kPastStateLengths);
+ const Tensor* state_update_capture_count =
+ context->InputCount() > psai::kStateUpdateCaptureCount
+ ? context->Input(psai::kStateUpdateCaptureCount)
+ : nullptr;
+ const Tensor* state_update_active =
+ context->InputCount() > psai::kStateUpdateActive ? context->Input(psai::kStateUpdateActive) : nullptr;
+
+ ORT_RETURN_IF(state_update_capacity_ > 0 && state_update_capture_count == nullptr,
+ "PackedSparseAttentionIndexer: state_update_capture_count is required when "
+ "state_update_capacity is positive");
ORT_RETURN_IF(query == nullptr, "PackedSparseAttentionIndexer: query is required");
const auto& query_shape = query->Shape();
@@ -192,13 +217,16 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con
const int64_t total_tokens = query_shape[0];
ORT_RETURN_IF(query_norm_weight == nullptr, "PackedSparseAttentionIndexer: query_norm_weight is required");
const auto& query_norm_shape = query_norm_weight->Shape();
- ORT_RETURN_IF_NOT(query_norm_shape.NumDimensions() == 1,
+ ORT_RETURN_IF_NOT(query_norm_shape.NumDimensions() == 1 && query_norm_shape[0] > 0,
"PackedSparseAttentionIndexer: query_norm_weight must have shape (head_size), got ",
query_norm_shape.ToString());
const int64_t head_size = query_norm_shape[0];
- ORT_RETURN_IF(head_size <= 0 || query_shape[1] % head_size != 0,
- "PackedSparseAttentionIndexer: query width must be divisible by head_size");
- const int64_t num_heads = query_shape[1] / head_size;
+ const bool has_separate_key = key != nullptr;
+ const int64_t query_width = query_shape[1] - (has_separate_key ? 0 : head_size);
+ ORT_RETURN_IF(query_width <= 0 || query_width % head_size != 0,
+ "PackedSparseAttentionIndexer: query width must contain a positive multiple of head_size",
+ has_separate_key ? "" : " followed by one packed key head");
+ const int64_t num_heads = query_width / head_size;
ORT_RETURN_IF_ERROR(CheckIntDimension("total_tokens", total_tokens));
ORT_RETURN_IF_ERROR(CheckIntDimension("num_heads", num_heads, false));
ORT_RETURN_IF_ERROR(CheckIntDimension("head_size", head_size, false));
@@ -216,7 +244,12 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con
"PackedSparseAttentionIndexer: total_tokens must be 0 when batch_size is 0");
ORT_RETURN_IF_ERROR(CheckShape(past_sequence_lengths, "past_sequence_lengths", {batch_size}));
- ORT_RETURN_IF_ERROR(CheckShape(key, "key", {total_tokens, head_size}));
+ if (position_ids != nullptr) {
+ ORT_RETURN_IF_ERROR(CheckShape(position_ids, "position_ids", {total_tokens}));
+ }
+ if (has_separate_key) {
+ ORT_RETURN_IF_ERROR(CheckShape(key, "key", {total_tokens, head_size}));
+ }
ORT_RETURN_IF_ERROR(CheckShape(query_norm_weight, "query_norm_weight", {head_size}));
ORT_RETURN_IF_ERROR(CheckShape(key_norm_weight, "key_norm_weight", {head_size}));
@@ -243,6 +276,12 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con
ORT_RETURN_IF_ERROR(CheckShape(past_kv_buffer, "past_kv_buffer", {batch_size, buffer_capacity, head_size}));
ORT_RETURN_IF_ERROR(CheckShape(past_state_lengths, "past_state_lengths",
{batch_size, psai::kStateLengthColumns}));
+ if (state_update_capture_count != nullptr) {
+ ORT_RETURN_IF_ERROR(CheckShape(state_update_capture_count, "state_update_capture_count", {batch_size}));
+ }
+ if (state_update_active != nullptr) {
+ ORT_RETURN_IF_ERROR(CheckShape(state_update_active, "state_update_active", {1}));
+ }
const int64_t capacity = psai::SelectedCapacity(psai::Policy::kQsa, token_budget_, index_topk_, compress_ratio_);
@@ -257,19 +296,30 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con
context->Output(psai::kPresentKvBuffer, TensorShape({batch_size, buffer_capacity, head_size}));
Tensor* present_state_lengths =
context->Output(psai::kPresentStateLengths, TensorShape({batch_size, psai::kStateLengthColumns}));
+ Tensor* state_update = context->OutputCount() > psai::kStateUpdate
+ ? context->Output(
+ psai::kStateUpdate, TensorShape({batch_size, state_update_capacity_, head_size}))
+ : nullptr;
ORT_RETURN_IF(selected_indices == nullptr || selected_counts == nullptr || present_key_state == nullptr ||
present_kv_buffer == nullptr || present_state_lengths == nullptr,
"PackedSparseAttentionIndexer: policy_mode 'qsa' requires selected_indices, selected_counts, "
"present_key_state, present_kv_buffer and present_state_lengths outputs");
+ if (state_update != nullptr && state_update->SizeInBytes() > 0) {
+ CUDA_RETURN_IF_ERROR(cudaMemsetAsync(state_update->MutableDataRaw(), 0, state_update->SizeInBytes(),
+ Stream(context)));
+ }
PackedSparseAttentionIndexerParams params;
params.batch_size = static_cast(batch_size);
params.total_tokens = static_cast(total_tokens);
params.num_heads = static_cast(num_heads);
params.head_size = static_cast(head_size);
+ params.query_row_stride = static_cast(query_shape[1]);
+ params.key_row_stride = static_cast(has_separate_key ? head_size : query_shape[1]);
params.rotary_width = static_cast(rotary.rotary_width);
params.max_rotary_length = static_cast(rotary.max_rotary_length);
params.cos_cache_batched = rotary.batched;
+ params.has_position_ids = position_ids != nullptr;
params.compress_ratio = static_cast(compress_ratio_);
params.state_capacity = static_cast(state_capacity);
params.buffer_capacity = static_cast(buffer_capacity);
@@ -277,29 +327,38 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con
params.epsilon = epsilon_;
params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size));
params.block_topk = static_cast(token_budget_ / compress_ratio_);
+ params.state_update_capacity = static_cast(state_update_capacity_);
auto float_workspace = GetScratchBuffer(GetQsaPackedWorkspaceFloatCount(params), GetComputeStream(context));
auto overflow_flags =
GetScratchBuffer(static_cast(std::max(batch_size, 1)), GetComputeStream(context));
+ const CudaT* query_data = reinterpret_cast(query->Data());
+ const CudaT* key_data = has_separate_key
+ ? reinterpret_cast(key->Data())
+ : query_data + query_width;
return LaunchQsaPackedSparseAttentionIndexer(
Stream(context), params,
- reinterpret_cast(query->Data()),
- reinterpret_cast(key->Data()),
+ query_data,
+ key_data,
reinterpret_cast(query_norm_weight->Data()),
reinterpret_cast(key_norm_weight->Data()),
reinterpret_cast(cos_cache->Data()),
reinterpret_cast(sin_cache->Data()),
cumulative_sequence_lengths->Data(),
past_sequence_lengths->Data(),
+ position_ids != nullptr ? position_ids->Data() : nullptr,
reinterpret_cast(past_key_state->Data()),
reinterpret_cast(past_kv_buffer->Data()),
past_state_lengths->Data(),
+ state_update_capture_count != nullptr ? state_update_capture_count->Data() : nullptr,
+ state_update_active == nullptr || state_update_active->Data()[0] != 0,
selected_indices->MutableData(),
selected_counts->MutableData(),
reinterpret_cast(present_key_state->MutableData()),
reinterpret_cast(present_kv_buffer->MutableData()),
present_state_lengths->MutableData(),
+ state_update != nullptr ? reinterpret_cast(state_update->MutableData()) : nullptr,
float_workspace.get(),
overflow_flags.get());
}
@@ -416,9 +475,11 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context) con
params.total_tokens = static_cast(total_tokens);
params.num_heads = static_cast(num_heads);
params.head_size = static_cast(head_size);
+ params.query_row_stride = static_cast(query_shape[1]);
params.rotary_width = static_cast(rotary.rotary_width);
params.max_rotary_length = static_cast(rotary.max_rotary_length);
params.cos_cache_batched = rotary.batched;
+ params.has_position_ids = true;
params.compress_ratio = static_cast(compress_ratio_);
params.state_capacity = static_cast(state_capacity);
params.buffer_capacity = static_cast(buffer_capacity);
diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h
index c8f3ec19f2082..112be6e3f46af 100644
--- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h
+++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h
@@ -24,6 +24,7 @@ class PackedSparseAttentionIndexer final : public onnxruntime::cuda::CudaKernel
packed_sparse_attention_indexer::Policy policy_;
int64_t compress_ratio_;
int64_t state_capacity_;
+ int64_t state_update_capacity_;
int64_t token_budget_;
int64_t index_topk_;
float epsilon_;
diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu
index 9f882bf732e04..a5ed4610319e9 100644
--- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu
@@ -30,23 +30,37 @@
#include
#include
#include
+#include
#include
#include "contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h"
#include "contrib_ops/cuda/sparse/sparse_attention_indexer_device_math.cuh"
#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh"
+#include "core/providers/cuda/cu_inc/topk_warp_sort.cuh"
namespace onnxruntime {
namespace contrib {
namespace cuda {
namespace psai = onnxruntime::contrib::packed_sparse_attention_indexer;
+namespace topk = onnxruntime::cuda::topk;
namespace {
// The block reductions below halve the active thread count, so this must stay a power of two.
constexpr int kThreads = 128;
+constexpr int kWarpSize = 32;
+constexpr int kBoundedTopKMax = 512;
+constexpr int kBoundedTopKItemsPerThread = kBoundedTopKMax / kThreads;
+constexpr int kDistributedTopKMaxRows = 64;
+constexpr int kDistributedTopKMinBlocks = 2048;
+constexpr int kHierarchicalTileBlocks = 8;
+constexpr int kHierarchicalScoreThreads = 1024;
+constexpr int kHierarchicalMergeThreads = 256;
+constexpr int kQwenHeadSize = 128;
+constexpr int kQwenCompressRatio = 4;
+constexpr int kQwenNumHeads = 4;
template
Status ConfigureDynamicSharedMemory(Kernel kernel, size_t bytes, const char* kernel_name) {
@@ -60,6 +74,36 @@ Status ConfigureDynamicSharedMemory(Kernel kernel, size_t bytes, const char* ker
return CUDA_CALL(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_bytes));
}
+bool UseHierarchicalQsaTopK(const PackedSparseAttentionIndexerParams& params) {
+ return params.total_tokens > 0 && params.total_tokens <= kDistributedTopKMaxRows &&
+ params.state_capacity >= kDistributedTopKMinBlocks &&
+ params.head_size == kQwenHeadSize && params.compress_ratio == kQwenCompressRatio &&
+ params.num_heads == kQwenNumHeads && params.block_topk > kSaiFastTopKMax &&
+ params.block_topk <= kBoundedTopKMax;
+}
+
+bool UseRaggedQsaPrefill(const PackedSparseAttentionIndexerParams& params) {
+ return params.total_tokens > kDistributedTopKMaxRows &&
+ params.head_size == kQwenHeadSize && params.compress_ratio == kQwenCompressRatio &&
+ params.num_heads == kQwenNumHeads;
+}
+
+int GetHierarchicalTileCount(const PackedSparseAttentionIndexerParams& params) {
+ return (params.state_capacity + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks;
+}
+
+template
+struct alignas(sizeof(T) * 4) QsaVector4 {
+ T values[4];
+};
+
+template
+__device__ __forceinline__ float4 LoadQsaVector4(const T* source) {
+ const QsaVector4 packed = *reinterpret_cast*>(source);
+ return make_float4(to_float(packed.values[0]), to_float(packed.values[1]),
+ to_float(packed.values[2]), to_float(packed.values[3]));
+}
+
// Largest b such that cumulative_sequence_lengths[b] <= token, assuming the array is nondecreasing.
// If the data itself is malformed this may attribute a token to the wrong request, but the result
// is always an index in [0, batch_size), so it can never cause an out-of-bounds access.
@@ -98,9 +142,10 @@ template
__global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, const T* cos_cache,
const T* sin_cache, const int32_t* cumulative_sequence_lengths,
const int32_t* past_sequence_lengths, const T* past_kv_buffer,
- const int32_t* past_state_lengths,
+ const int32_t* past_state_lengths, const int32_t* state_update_capture_count,
+ bool state_update_active,
T* present_key_state, T* present_kv_buffer,
- int32_t* present_state_lengths, int32_t* overflow_flags,
+ int32_t* present_state_lengths, T* state_update, int32_t* overflow_flags,
PackedSparseAttentionIndexerParams params) {
extern __shared__ float shared[];
float* pooled = shared;
@@ -134,6 +179,12 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con
const bool rejected = invalid_metadata || full_new_block_count > capacity_left;
const int new_block_count = rejected ? 0 : full_new_block_count;
const int new_buf_len = rejected ? old_buf_len : (pending % params.compress_ratio);
+ const bool capture_active = !rejected && state_update != nullptr && state_update_capture_count != nullptr &&
+ state_update_active;
+ const int capture_count = capture_active
+ ? min(max(state_update_capture_count[b], 0),
+ min(req_len, params.state_update_capacity))
+ : 0;
// Barrier: every thread has now read past_state_lengths (identically) before any thread below
// writes present_state_lengths, which keeps this correct even if the two tensors alias.
@@ -145,6 +196,18 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con
overflow_flags[b] = rejected ? 1 : 0;
}
+ for (int t = static_cast(threadIdx.x); t < capture_count; t += static_cast(blockDim.x)) {
+ if ((old_buf_len + t + 1) % params.compress_ratio != 0) {
+ const int64_t capture_base =
+ (static_cast(b) * params.state_update_capacity + t) * params.head_size;
+ const int64_t key_base = (static_cast(req_start) + t) * params.key_row_stride;
+ for (int d = 0; d < params.head_size; ++d) {
+ state_update[capture_base + d] = key[key_base + d];
+ }
+ }
+ }
+ __syncthreads();
+
for (int k = 0; k < new_block_count; ++k) {
for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) {
float sum = 0.0f;
@@ -155,7 +218,7 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con
params.head_size +
d])
: to_float(key[(static_cast(req_start) + (virtual_pos - old_buf_len)) *
- params.head_size +
+ params.key_row_stride +
d]);
}
pooled[d] = sum / static_cast(params.compress_ratio);
@@ -186,8 +249,14 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con
__syncthreads();
const int64_t out_base = (static_cast(b) * params.state_capacity + entry) * params.head_size;
+ const int completion_token = (k + 1) * params.compress_ratio - old_buf_len - 1;
+ const int64_t capture_base =
+ (static_cast(b) * params.state_update_capacity + completion_token) * params.head_size;
for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) {
present_key_state[out_base + d] = from_float(rotated[d]);
+ if (completion_token < capture_count) {
+ state_update[capture_base + d] = from_float(rotated[d]);
+ }
}
__syncthreads();
}
@@ -207,7 +276,9 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con
params.head_size +
d])
: to_float(
- key[(static_cast(req_start) + (virtual_pos - old_buf_len)) * params.head_size + d]);
+ key[(static_cast(req_start) + (virtual_pos - old_buf_len)) *
+ params.key_row_stride +
+ d]);
present_kv_buffer[out_base + d] = from_float(value);
}
}
@@ -232,10 +303,11 @@ __global__ void PackedRotateQueryKernel(const T* query, const T* query_norm_weig
for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
const int token = static_cast(row / params.num_heads);
const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token);
- const int64_t base = row * params.head_size;
+ const int64_t input_base = static_cast(token) * params.query_row_stride +
+ static_cast(row % params.num_heads) * params.head_size;
for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) {
- shared[d] = to_float(query[base + d]);
+ shared[d] = to_float(query[input_base + d]);
}
__syncthreads();
@@ -251,9 +323,9 @@ __global__ void PackedRotateQueryKernel(const T* query, const T* query_norm_weig
__syncthreads();
const int64_t abs_position =
- kUseLeadingRope
- ? static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch])
- : position_ids[token];
+ params.has_position_ids
+ ? position_ids[token]
+ : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
const int position = SaiClampPosition(abs_position, params.max_rotary_length);
const int64_t cache_offset =
(static_cast(params.cos_cache_batched ? batch : 0) * params.max_rotary_length + position) *
@@ -262,10 +334,9 @@ __global__ void PackedRotateQueryKernel(const T* query, const T* query_norm_weig
const T* sin_row = sin_cache + cache_offset;
for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) {
- query_rotated[base + d] = kUseLeadingRope
- ? SaiLeadingRope(shared, params.rotary_width, cos_row, sin_row, d)
- : SaiTrailingRope(shared, params.head_size, params.rotary_width, cos_row,
- sin_row, d);
+ query_rotated[row * params.head_size + d] =
+ kUseLeadingRope ? SaiLeadingRope(shared, params.rotary_width, cos_row, sin_row, d)
+ : SaiTrailingRope(shared, params.head_size, params.rotary_width, cos_row, sin_row, d);
}
__syncthreads();
}
@@ -277,7 +348,7 @@ __global__ void PackedRotateQueryKernel(const T* query, const T* query_norm_weig
template
__global__ void QsaBlockScoreKernel(const T* present_key_state, const float* query_rotated,
const int32_t* cumulative_sequence_lengths,
- const int32_t* past_sequence_lengths,
+ const int32_t* past_sequence_lengths, const int64_t* position_ids,
const int32_t* present_state_lengths, float* block_scores,
PackedSparseAttentionIndexerParams params) {
extern __shared__ float reduction[];
@@ -296,7 +367,9 @@ __global__ void QsaBlockScoreKernel(const T* present_key_state, const float* que
}
const int64_t abs_position =
- static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
+ params.has_position_ids
+ ? position_ids[token]
+ : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
const int64_t causal_count = SaiCausalThreshold(abs_position, params.compress_ratio);
const int64_t visible_block_count = causal_count < key_len_after ? causal_count : key_len_after;
if (static_cast(block_index) >= visible_block_count) {
@@ -324,18 +397,248 @@ __global__ void QsaBlockScoreKernel(const T* present_key_state, const float* que
}
}
+// One CTA per packed query row. Eight state blocks are scored concurrently, with one warp per
+// (state block, query head). Unlike the generic rectangular launch, each CTA stops at its row's
+// live causal block count and never launches or writes work for unused fixed-capacity slots.
+template
+__global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const float* query_rotated,
+ const int32_t* cumulative_sequence_lengths,
+ const int32_t* past_sequence_lengths, const int64_t* position_ids,
+ const int32_t* present_state_lengths, const int32_t* overflow_flags,
+ float* block_scores, PackedSparseAttentionIndexerParams params) {
+ __shared__ float head_scores[kHierarchicalTileBlocks][kQwenNumHeads];
+
+ const int warp = static_cast(threadIdx.x) / kWarpSize;
+ const int lane = static_cast(threadIdx.x) % kWarpSize;
+ const int block_slot = warp / kQwenNumHeads;
+ const int head = warp % kQwenNumHeads;
+ for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) {
+ const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token);
+ const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength];
+ const int64_t abs_position =
+ params.has_position_ids
+ ? position_ids[token]
+ : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
+ const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio);
+ const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after);
+ const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0;
+
+ for (int tile_start = 0; tile_start < block_count; tile_start += kHierarchicalTileBlocks) {
+ const int block_index = tile_start + block_slot;
+ float score = 0.0f;
+ if (block_index < block_count) {
+ const int d = lane * 4;
+ const int64_t query_base =
+ (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize;
+ const int64_t key_base =
+ (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize;
+ const float4 query_value = LoadQsaVector4(query_rotated + query_base + d);
+ const float4 key_value = LoadQsaVector4(present_key_state + key_base + d);
+ score = query_value.x * key_value.x + query_value.y * key_value.y +
+ query_value.z * key_value.z + query_value.w * key_value.w;
+ score = topk::WarpReduceSum(score);
+ }
+ if (lane == 0) {
+ head_scores[block_slot][head] = fmaxf(score, 0.0f);
+ }
+ __syncthreads();
+ if (threadIdx.x < kHierarchicalTileBlocks && tile_start + threadIdx.x < block_count) {
+ float total_score = 0.0f;
+#pragma unroll
+ for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) {
+ total_score += head_scores[threadIdx.x][query_head];
+ }
+ const float scaled_score = total_score * params.scale;
+ block_scores[static_cast(token) * params.state_capacity + tile_start + threadIdx.x] =
+ scaled_score == 0.0f ? 0.0f : scaled_score;
+ }
+ __syncthreads();
+ }
+ }
+}
+
+template
+__global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* query_rotated,
+ const int32_t* cumulative_sequence_lengths,
+ const int32_t* past_sequence_lengths, const int64_t* position_ids,
+ const int32_t* present_state_lengths, const int32_t* overflow_flags,
+ uint64_t* tile_keys, int tile_count,
+ PackedSparseAttentionIndexerParams params) {
+ __shared__ uint64_t keys[kHierarchicalTileBlocks];
+ __shared__ float head_scores[kHierarchicalTileBlocks][kQwenNumHeads];
+
+ const int64_t total_tiles = static_cast(params.total_tokens) * tile_count;
+ const int warp = static_cast(threadIdx.x) / kWarpSize;
+ const int lane = static_cast(threadIdx.x) % kWarpSize;
+ const int block_slot = warp / kQwenNumHeads;
+ const int head = warp % kQwenNumHeads;
+ for (int64_t work = blockIdx.x; work < total_tiles; work += gridDim.x) {
+ const int tile = static_cast(work % tile_count);
+ const int token = static_cast(work / tile_count);
+ const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token);
+ const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength];
+ const int64_t abs_position =
+ params.has_position_ids
+ ? position_ids[token]
+ : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
+ const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio);
+ const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after);
+ const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0;
+
+ const int block_index = tile * kHierarchicalTileBlocks + block_slot;
+ float score = 0.0f;
+ if (block_index < block_count) {
+ const int d = lane * 4;
+ const int64_t query_base =
+ (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize;
+ const int64_t key_base =
+ (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize;
+ const float4 query_value = LoadQsaVector4(query_rotated + query_base + d);
+ const float4 key_value = LoadQsaVector4(present_key_state + key_base + d);
+ score = query_value.x * key_value.x + query_value.y * key_value.y +
+ query_value.z * key_value.z + query_value.w * key_value.w;
+ score = topk::WarpReduceSum(score);
+ }
+ if (lane == 0) {
+ head_scores[block_slot][head] = fmaxf(score, 0.0f);
+ }
+ __syncthreads();
+ if (threadIdx.x < kHierarchicalTileBlocks) {
+ const int tile_block = tile * kHierarchicalTileBlocks + static_cast(threadIdx.x);
+ uint64_t key = topk::kPaddingSortKey;
+ if (tile_block < block_count) {
+ float total_score = 0.0f;
+#pragma unroll
+ for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) {
+ total_score += head_scores[threadIdx.x][query_head];
+ }
+ const float scaled_score = total_score * params.scale;
+ key = topk::PackStableSortKey(scaled_score == 0.0f ? 0.0f : scaled_score, tile_block);
+ }
+ keys[threadIdx.x] = key;
+ }
+ __syncthreads();
+ if (threadIdx.x == 0) {
+ for (int index = 1; index < kHierarchicalTileBlocks; ++index) {
+ const uint64_t key = keys[index];
+ int insertion = index;
+ while (insertion > 0 && key > keys[insertion - 1]) {
+ keys[insertion] = keys[insertion - 1];
+ --insertion;
+ }
+ keys[insertion] = key;
+ }
+ }
+ __syncthreads();
+ const int64_t output_base = work * kHierarchicalTileBlocks;
+ if (threadIdx.x < kHierarchicalTileBlocks) {
+ tile_keys[output_base + threadIdx.x] = keys[threadIdx.x];
+ }
+ __syncthreads();
+ }
+}
+
+__device__ __forceinline__ uint64_t QsaMergeRank(const uint64_t* left, int left_count,
+ const uint64_t* right, int right_count, int rank) {
+ int low = max(0, rank - right_count);
+ int high = min(rank, left_count);
+ while (low <= high) {
+ const int left_rank = (low + high) / 2;
+ const int right_rank = rank - left_rank;
+ if (left_rank > 0 && right_rank < right_count && left[left_rank - 1] < right[right_rank]) {
+ high = left_rank - 1;
+ } else if (right_rank > 0 && left_rank < left_count && right[right_rank - 1] < left[left_rank]) {
+ low = left_rank + 1;
+ } else {
+ const uint64_t left_key = left_rank < left_count ? left[left_rank] : topk::kPaddingSortKey;
+ const uint64_t right_key = right_rank < right_count ? right[right_rank] : topk::kPaddingSortKey;
+ return max(left_key, right_key);
+ }
+ }
+ return topk::kPaddingSortKey;
+}
+
+__global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* output_keys,
+ int key_stride, int list_width, int rows) {
+ const int pairs_per_row = (key_stride + 2 * list_width - 1) / (2 * list_width);
+ const int64_t total_pairs = static_cast(rows) * pairs_per_row;
+ for (int64_t work = blockIdx.x; work < total_pairs; work += gridDim.x) {
+ const int pair = static_cast(work % pairs_per_row);
+ const int64_t row = work / pairs_per_row;
+ const int left_start = pair * 2 * list_width;
+ const int right_start = left_start + list_width;
+ const int left_count = min(min(list_width, kBoundedTopKMax), max(0, key_stride - left_start));
+ const int right_count = min(min(list_width, kBoundedTopKMax), max(0, key_stride - right_start));
+ const int output_count = min(left_count + right_count, kBoundedTopKMax);
+ const int64_t row_base = row * key_stride;
+ for (int rank = threadIdx.x; rank < output_count; rank += blockDim.x) {
+ output_keys[row_base + left_start + rank] =
+ QsaMergeRank(input_keys + row_base + left_start, left_count,
+ input_keys + row_base + right_start, right_count, rank);
+ }
+ }
+}
+
+__global__ void QsaEmitHierarchicalTopKKernel(const uint64_t* tile_keys,
+ const int32_t* cumulative_sequence_lengths,
+ const int32_t* past_sequence_lengths,
+ const int64_t* position_ids,
+ const int32_t* present_state_lengths,
+ const int32_t* overflow_flags, int key_stride,
+ int32_t* selected_indices, int32_t* selected_counts,
+ PackedSparseAttentionIndexerParams params) {
+ for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) {
+ int32_t* output = selected_indices + static_cast(token) * params.capacity;
+ const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token);
+ const int64_t abs_position =
+ params.has_position_ids
+ ? position_ids[token]
+ : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
+ const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength];
+ const int64_t causal_count = SaiCausalThreshold(abs_position, params.compress_ratio);
+ const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after);
+ const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0;
+ const int selected = min(params.block_topk, block_count);
+ for (int rank = threadIdx.x; rank < selected; rank += blockDim.x) {
+ const int selected_block = topk::UnpackStableSortIndex(tile_keys[static_cast(token) * key_stride + rank]);
+ for (int t = 0; t < params.compress_ratio; ++t) {
+ output[rank * params.compress_ratio + t] = selected_block * params.compress_ratio + t;
+ }
+ }
+ const int64_t block_start = static_cast(block_count) * params.compress_ratio;
+ const int64_t natural_tail = abs_position >= block_start ? abs_position - block_start + 1 : 0;
+ const int remaining_capacity = params.capacity - selected * params.compress_ratio;
+ const int64_t tail_count_64 = natural_tail < remaining_capacity ? natural_tail : remaining_capacity;
+ const int tail_count = overflow_flags[batch] == 0 ? static_cast(tail_count_64 < 0 ? 0 : tail_count_64) : 0;
+ for (int t = threadIdx.x; t < tail_count; t += blockDim.x) {
+ output[selected * params.compress_ratio + t] = static_cast(block_start + t);
+ }
+ for (int position = selected * params.compress_ratio + tail_count + threadIdx.x;
+ position < params.capacity; position += blockDim.x) {
+ output[position] = -1;
+ }
+ if (threadIdx.x == 0) {
+ selected_counts[token] = selected * params.compress_ratio + tail_count;
+ }
+ }
+}
+
// One block per query token. Emits the token indices of the highest scoring blocks followed by the
// causally visible tokens of the trailing incomplete block, and the exact active count. A request
// whose update step was rejected for exceeding state_capacity this call (overflow_flags[batch] set)
// always gets the safe empty result: indices stay -1 (already reset below) and count is 0.
-__global__ void QsaSelectKernel(const float* block_scores, const int32_t* cumulative_sequence_lengths,
- const int32_t* past_sequence_lengths,
+__global__ void QsaSelectKernel(const float* block_scores, const int32_t* topk_indices,
+ const int32_t* cumulative_sequence_lengths,
+ const int32_t* past_sequence_lengths, const int64_t* position_ids,
const int32_t* present_state_lengths, const int32_t* overflow_flags,
int32_t* selected_indices, int32_t* selected_counts,
PackedSparseAttentionIndexerParams params) {
extern __shared__ float shared[];
float* shared_value = shared;
- int* shared_index = reinterpret_cast(shared + blockDim.x);
+ const bool use_fast_topk = params.block_topk <= kSaiFastTopKMax;
+ const int shared_entries = use_fast_topk ? static_cast(blockDim.x) * params.block_topk
+ : static_cast(blockDim.x);
+ int* shared_index = reinterpret_cast(shared + shared_entries);
for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) {
int32_t* out_row = selected_indices + static_cast(token) * params.capacity;
@@ -355,7 +658,9 @@ __global__ void QsaSelectKernel(const float* block_scores, const int32_t* cumula
const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength];
const int64_t abs_position =
- static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
+ params.has_position_ids
+ ? position_ids[token]
+ : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
const int64_t causal_count = SaiCausalThreshold(abs_position, params.compress_ratio);
const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after);
const int visible_block_count = static_cast(visible_64 < 0 ? 0 : visible_64);
@@ -363,28 +668,52 @@ __global__ void QsaSelectKernel(const float* block_scores, const int32_t* cumula
const float* scores_row = block_scores + static_cast(token) * params.state_capacity;
- float previous_score = 0.0f;
- int previous_index = -1;
- int emitted_blocks = 0;
- for (int rank = 0; rank < selected; ++rank) {
- float best_value = 0.0f;
- int best_index = -1;
- SaiScanForNext(scores_row, visible_block_count, previous_score, previous_index, &best_value, &best_index);
- shared_value[threadIdx.x] = best_value;
- shared_index[threadIdx.x] = best_index;
- __syncthreads();
- SaiBlockArgMax(shared_value, shared_index);
- previous_index = shared_index[0];
- previous_score = shared_value[0];
- __syncthreads();
- if (previous_index < 0) {
- break;
+ int emitted_blocks = selected;
+ if (topk_indices != nullptr) {
+ const int32_t* topk_row = topk_indices + static_cast(token) * params.state_capacity;
+ for (int rank = 0; rank < selected; ++rank) {
+ const int selected_block = topk_row[rank];
+ for (int t = static_cast(threadIdx.x); t < params.compress_ratio;
+ t += static_cast(blockDim.x)) {
+ out_row[rank * params.compress_ratio + t] = selected_block * params.compress_ratio + t;
+ }
}
- for (int t = static_cast(threadIdx.x); t < params.compress_ratio; t += static_cast(blockDim.x)) {
- out_row[rank * params.compress_ratio + t] = previous_index * params.compress_ratio + t;
+ __syncthreads();
+ } else if (selected > 0 && use_fast_topk) {
+ SaiBlockTopK(scores_row, visible_block_count, selected, shared_value, shared_index);
+ for (int rank = 0; rank < selected; ++rank) {
+ const int selected_block = shared_index[rank];
+ for (int t = static_cast(threadIdx.x); t < params.compress_ratio;
+ t += static_cast(blockDim.x)) {
+ out_row[rank * params.compress_ratio + t] = selected_block * params.compress_ratio + t;
+ }
}
- emitted_blocks = rank + 1;
__syncthreads();
+ } else {
+ float previous_score = 0.0f;
+ int previous_index = -1;
+ emitted_blocks = 0;
+ for (int rank = 0; rank < selected; ++rank) {
+ float best_value = 0.0f;
+ int best_index = -1;
+ SaiScanForNext(scores_row, visible_block_count, previous_score, previous_index, &best_value, &best_index);
+ shared_value[threadIdx.x] = best_value;
+ shared_index[threadIdx.x] = best_index;
+ __syncthreads();
+ SaiBlockArgMax(shared_value, shared_index);
+ previous_index = shared_index[0];
+ previous_score = shared_value[0];
+ __syncthreads();
+ if (previous_index < 0) {
+ break;
+ }
+ for (int t = static_cast(threadIdx.x); t < params.compress_ratio;
+ t += static_cast(blockDim.x)) {
+ out_row[rank * params.compress_ratio + t] = previous_index * params.compress_ratio + t;
+ }
+ emitted_blocks = rank + 1;
+ __syncthreads();
+ }
}
// The trailing incomplete block is always causally visible in full up to this query's own
@@ -404,6 +733,148 @@ __global__ void QsaSelectKernel(const float* block_scores, const int32_t* cumula
}
}
+// One block per packed query token. Radix-selects a bounded candidate set, then sorts only that
+// set. The output aliases the corresponding score row after all scores for that row have been read.
+__global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulative_sequence_lengths,
+ const int32_t* past_sequence_lengths, const int64_t* position_ids,
+ const int32_t* present_state_lengths, const int32_t* overflow_flags,
+ PackedSparseAttentionIndexerParams params) {
+ using Sort = cub::BlockRadixSort;
+ __shared__ union {
+ uint32_t histogram[256];
+ typename Sort::TempStorage sort;
+ } temp;
+ __shared__ uint64_t selected_keys[kBoundedTopKMax];
+ __shared__ int32_t equal_warp_offsets[kThreads / kWarpSize + 1];
+ __shared__ uint32_t prefix;
+ __shared__ int remaining;
+ __shared__ int gathered;
+ __shared__ int equal_seen;
+
+ for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) {
+ const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token);
+ if (overflow_flags[batch] != 0) {
+ continue;
+ }
+ const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength];
+ const int64_t abs_position =
+ params.has_position_ids
+ ? position_ids[token]
+ : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]);
+ const int64_t causal_count = SaiCausalThreshold(abs_position, params.compress_ratio);
+ const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after);
+ const int block_count = static_cast(visible_64 < 0 ? 0 : visible_64);
+ const int selected = min(params.block_topk, block_count);
+ if (selected == 0) {
+ continue;
+ }
+ float* scores = block_scores + static_cast(token) * params.state_capacity;
+
+ if (threadIdx.x == 0) {
+ prefix = 0;
+ remaining = selected;
+ }
+ __syncthreads();
+
+ for (int shift = 24; shift >= 0; shift -= 8) {
+ for (int bucket = threadIdx.x; bucket < 256; bucket += blockDim.x) {
+ temp.histogram[bucket] = 0;
+ }
+ __syncthreads();
+ const uint32_t current_prefix = prefix;
+ for (int index = threadIdx.x; index < block_count; index += blockDim.x) {
+ const uint32_t score_key = static_cast(topk::PackStableSortKey(scores[index], 0) >> 32);
+ if (shift == 24 || (score_key >> (shift + 8)) == (current_prefix >> (shift + 8))) {
+ atomicAdd(&temp.histogram[(score_key >> shift) & 0xffu], 1u);
+ }
+ }
+ __syncthreads();
+ if (threadIdx.x == 0) {
+ int rank = remaining;
+ for (int bucket = 255; bucket >= 0; --bucket) {
+ const int count = static_cast(temp.histogram[bucket]);
+ if (rank > count) {
+ rank -= count;
+ } else {
+ prefix |= static_cast(bucket) << shift;
+ remaining = rank;
+ break;
+ }
+ }
+ }
+ __syncthreads();
+ }
+
+ if (threadIdx.x == 0) {
+ gathered = 0;
+ }
+ __syncthreads();
+ const uint32_t threshold = prefix;
+ for (int index = threadIdx.x; index < block_count; index += blockDim.x) {
+ const uint64_t key = topk::PackStableSortKey(scores[index], index);
+ if (static_cast(key >> 32) > threshold) {
+ const int slot = atomicAdd(&gathered, 1);
+ selected_keys[slot] = key;
+ }
+ }
+ __syncthreads();
+
+ if (threadIdx.x == 0) {
+ equal_seen = 0;
+ }
+ __syncthreads();
+ for (int base = 0; base < block_count && equal_seen < remaining; base += blockDim.x) {
+ const int index = base + static_cast(threadIdx.x);
+ const bool equal = index < block_count &&
+ static_cast(topk::PackStableSortKey(scores[index], 0) >> 32) == threshold;
+ const unsigned int warp_mask = __ballot_sync(0xffffffffu, equal);
+ const int lane = static_cast(threadIdx.x) % kWarpSize;
+ const int warp = static_cast(threadIdx.x) / kWarpSize;
+ const int lane_rank = __popc(warp_mask & (lane == 0 ? 0u : (1u << lane) - 1u));
+ if (lane == 0) {
+ equal_warp_offsets[warp] = __popc(warp_mask);
+ }
+ __syncthreads();
+ if (threadIdx.x == 0) {
+ int tile_total = 0;
+ for (int i = 0; i < kThreads / kWarpSize; ++i) {
+ const int warp_count = equal_warp_offsets[i];
+ equal_warp_offsets[i] = tile_total;
+ tile_total += warp_count;
+ }
+ equal_warp_offsets[kThreads / kWarpSize] = tile_total;
+ }
+ __syncthreads();
+ const int equal_rank = equal_seen + equal_warp_offsets[warp] + lane_rank;
+ if (equal && equal_rank < remaining) {
+ selected_keys[gathered + equal_rank] = topk::PackStableSortKey(scores[index], index);
+ }
+ __syncthreads();
+ if (threadIdx.x == 0) {
+ equal_seen += equal_warp_offsets[kThreads / kWarpSize];
+ }
+ __syncthreads();
+ }
+
+ uint64_t keys[kBoundedTopKItemsPerThread];
+#pragma unroll
+ for (int item = 0; item < kBoundedTopKItemsPerThread; ++item) {
+ const int rank = threadIdx.x * kBoundedTopKItemsPerThread + item;
+ keys[item] = rank < selected ? selected_keys[rank] : topk::kPaddingSortKey;
+ }
+ Sort(temp.sort).SortDescending(keys);
+ int32_t* topk_row = reinterpret_cast(scores);
+#pragma unroll
+ for (int item = 0; item < kBoundedTopKItemsPerThread; ++item) {
+ const int rank = threadIdx.x * kBoundedTopKItemsPerThread + item;
+ if (rank < selected) {
+ topk_row[rank] = topk::UnpackStableSortIndex(keys[item]);
+ }
+ }
+ __syncthreads();
+ }
+}
+
// ---------------------------------------------------------------------------------------------
// policy_mode = "csa"
// ---------------------------------------------------------------------------------------------
@@ -738,7 +1209,12 @@ __global__ void CsaSelectKernel(const float* scores, const int32_t* cumulative_s
size_t GetQsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& params) {
const size_t rows = static_cast(params.total_tokens);
- return rows * params.num_heads * params.head_size + rows * static_cast(std::max(params.state_capacity, 1));
+ const size_t query_count = rows * params.num_heads * params.head_size;
+ if (UseHierarchicalQsaTopK(params)) {
+ const size_t key_count = rows * GetHierarchicalTileCount(params) * kHierarchicalTileBlocks;
+ return query_count + 4 * key_count;
+ }
+ return query_count + rows * static_cast(std::max(params.state_capacity, 1));
}
size_t GetCsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& params) {
@@ -751,10 +1227,11 @@ Status LaunchQsaPackedSparseAttentionIndexer(
cudaStream_t stream, const PackedSparseAttentionIndexerParams& params, const T* query, const T* key,
const T* query_norm_weight, const T* key_norm_weight, const T* cos_cache, const T* sin_cache,
const int32_t* cumulative_sequence_lengths,
- const int32_t* past_sequence_lengths, const T* past_key_state,
- const T* past_kv_buffer, const int32_t* past_state_lengths, int32_t* selected_indices,
+ const int32_t* past_sequence_lengths, const int64_t* position_ids, const T* past_key_state,
+ const T* past_kv_buffer, const int32_t* past_state_lengths, const int32_t* state_update_capture_count,
+ bool state_update_active, int32_t* selected_indices,
int32_t* selected_counts, T* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths,
- float* float_workspace, int32_t* overflow_flags) {
+ T* state_update, float* float_workspace, int32_t* overflow_flags) {
if (params.batch_size > 0) {
const int64_t key_state_elems =
static_cast(params.batch_size) * params.state_capacity * params.head_size;
@@ -785,15 +1262,19 @@ Status LaunchQsaPackedSparseAttentionIndexer(
const int state_blocks = static_cast(std::min(params.batch_size, kSaiMaxGridDimX));
QsaUpdateStateKernel<<>>(
key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, past_kv_buffer,
- past_state_lengths, present_key_state, present_kv_buffer, present_state_lengths, overflow_flags, params);
+ past_state_lengths, state_update_capture_count, state_update_active, present_key_state, present_kv_buffer,
+ present_state_lengths, state_update, overflow_flags, params);
if (params.total_tokens == 0) {
return CUDA_CALL(cudaGetLastError());
}
float* query_rotated = float_workspace;
- float* block_scores =
- float_workspace + static_cast(params.total_tokens) * params.num_heads * params.head_size;
+ float* selection_workspace =
+ query_rotated + static_cast(params.total_tokens) * params.num_heads * params.head_size;
+ const bool use_hierarchical_topk = UseHierarchicalQsaTopK(params);
+ const bool use_ragged_prefill = UseRaggedQsaPrefill(params);
+ float* block_scores = use_hierarchical_topk ? nullptr : selection_workspace;
const int64_t rotate_rows = static_cast(params.total_tokens) * params.num_heads;
const int rotate_blocks = static_cast(std::min(rotate_rows, kSaiMaxGridDimX));
@@ -802,20 +1283,59 @@ Status LaunchQsaPackedSparseAttentionIndexer(
PackedRotateQueryKernel, rotate_shared_bytes, "PackedRotateQueryKernel"));
PackedRotateQueryKernel<<>>(
query, query_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths,
- nullptr, query_rotated, params);
+ position_ids, query_rotated, params);
- if (params.state_capacity > 0) {
+ const int token_blocks = static_cast(std::min(params.total_tokens, kSaiMaxGridDimX));
+ if (use_hierarchical_topk) {
+ const int tile_count = GetHierarchicalTileCount(params);
+ const int key_stride = tile_count * kHierarchicalTileBlocks;
+ const size_t key_count = static_cast(params.total_tokens) * key_stride;
+ uint64_t* merge_input = reinterpret_cast(selection_workspace);
+ uint64_t* merge_output = merge_input + key_count;
+ const int64_t tile_work = static_cast(params.total_tokens) * tile_count;
+ QsaScoreTileTopKKernel<<(std::min(tile_work, kSaiMaxGridDimX)),
+ kHierarchicalScoreThreads, 0, stream>>>(
+ present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids,
+ present_state_lengths, overflow_flags, merge_input, tile_count, params);
+ for (int list_width = kHierarchicalTileBlocks; list_width < key_stride; list_width *= 2) {
+ const int pairs_per_row = (key_stride + 2 * list_width - 1) / (2 * list_width);
+ const int64_t merge_work = static_cast(params.total_tokens) * pairs_per_row;
+ QsaMergeTileTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)),
+ kHierarchicalMergeThreads, 0, stream>>>(
+ merge_input, merge_output, key_stride, list_width, params.total_tokens);
+ std::swap(merge_input, merge_output);
+ }
+ QsaEmitHierarchicalTopKKernel<<>>(
+ merge_input, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths,
+ overflow_flags, key_stride, selected_indices, selected_counts, params);
+ } else if (use_ragged_prefill) {
+ QsaRaggedPrefillScoreKernel<<>>(
+ present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids,
+ present_state_lengths, overflow_flags, block_scores, params);
+ } else if (params.state_capacity > 0) {
const int64_t score_work = static_cast