[CUDA] [WebGPU] Optimize sparse attention indexers and add packed QSA state updates

docs/ContribOperators.md

@@ -4991,7 +4991,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>bias</tt> (optional) : T</dt>
 <dd>Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection</dd>
 <dt><tt>token_offset</tt> : M</dt>
-<dd>Offset of each token before packing, with shape (batch_size, sequence_length). Values must be in [0, batch_size * sequence_length).</dd>
+<dd>Offset of each token before packing, with shape (batch_size, sequence_length).</dd>
 <dt><tt>cumulative_sequence_length</tt> : M</dt>
 <dd>A tensor with shape (batch_size + 1). It specifies the cumulative sequence length.</dd>
 <dt><tt>attention_bias</tt> (optional) : T</dt>
@@ -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
 <dd>Scale applied to the per-head ReLU scores. Default is 1/sqrt(head_size).</dd>
 <dt><tt>state_capacity</tt> : int (required)</dt>
 <dd>Fixed capacity (number of entries) of past_key_state / present_key_state. Must be > 0.</dd>
+<dt><tt>state_update_capacity</tt> : int</dt>
+<dd>Only for policy_mode 'qsa': maximum number of leading token transitions captured per request. Must be in [0, 8]. Default is 0.</dd>
 <dt><tt>token_budget</tt> : int</dt>
 <dd>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'.</dd>
 </dl>
 
-#### Inputs
+#### Inputs (16 - 18)
 
 <dl>
 <dt><tt>query</tt> : T</dt>
-<dd>Packed indexer queries with shape (total_tokens, num_heads * head_size), before normalization, logical reshape, and rotary embedding.</dd>
-<dt><tt>key</tt> : T</dt>
-<dd>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.</dd>
+<dd>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).</dd>
+<dt><tt>key</tt> (optional) : T</dt>
+<dd>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.</dd>
 <dt><tt>query_norm_weight</tt> : T</dt>
 <dd>Effective RMSNorm multiplier of the queries, with shape (head_size).</dd>
 <dt><tt>key_norm_weight</tt> : T</dt>
@@ -5143,9 +5145,13 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>Only for policy_mode 'csa': buffered gate projections with the same shape as past_kv_buffer.</dd>
 <dt><tt>past_state_lengths</tt> : M</dt>
 <dd>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)).</dd>
+<dt><tt>state_update_capture_count</tt> (optional) : M</dt>
+<dd>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.</dd>
+<dt><tt>state_update_active</tt> (optional) : M</dt>
+<dd>Only for policy_mode 'qsa': optional capture gate with shape (1). A zero value disables capture.</dd>
 </dl>
 
-#### Outputs
+#### Outputs (6 - 7)
 
 <dl>
 <dt><tt>selected_indices</tt> : M</dt>
@@ -5160,6 +5166,8 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>Only for policy_mode 'csa': updated gate buffer with the same fixed shape as past_gate_buffer.</dd>
 <dt><tt>present_state_lengths</tt> : M</dt>
 <dd>Updated generic per-request state length, with the same fixed shape as past_state_lengths.</dd>
+<dt><tt>state_update</tt> (optional) : T</dt>
+<dd>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.</dd>
 </dl>
 
 #### Type Constraints
@@ -6943,7 +6951,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>input</tt> : T</dt>
 <dd>Input tensor with shape (total_tokens, hidden_size)</dd>
 <dt><tt>token_offset</tt> : M</dt>
-<dd>Offset of non-padding tokens and paddings. Its shape is (batch_size, sequence_length), and values must be in [0, batch_size * sequence_length).</dd>
+<dd>Offset of non-padding tokens and paddings. Its shape is (batch_size, sequence_length)</dd>
 </dl>
 
 #### Outputs
@@ -7616,11 +7624,11 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>key_norm_weight</tt> : T</dt>
 <dd>Effective RMSNorm multiplier of the compressed keys, with shape (head_size).</dd>
 <dt><tt>cos_cache</tt> : T</dt>
-<dd>Cosine rotary table indexed by absolute key position, with shape (batch_size, max_rotary_sequence_length, rotary_width).</dd>
+<dd>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).</dd>
 <dt><tt>sin_cache</tt> : T</dt>
 <dd>Sine rotary table with the same shape as cos_cache.</dd>
 <dt><tt>mask</tt> (optional) : TB</dt>
-<dd>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.</dd>
+<dd>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.</dd>
 <dt><tt>past_key</tt> : T</dt>
 <dd>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.</dd>
 <dt><tt>gate</tt> (optional) : T</dt>

docs/OperatorKernels.md

@@ -1131,7 +1131,7 @@ The **OpSet Version** column uses the following notation:
 |NhwcConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
 |PackedAttention|*in* input:**T**<br> *in* weights:**T**<br> *in* bias:**T**<br> *in* token_offset:**M**<br> *in* cumulative_sequence_length:**M**<br> *in* attention_bias:**T**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
 |PackedMultiHeadAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* bias:**T**<br> *in* token_offset:**M**<br> *in* cumulative_sequence_length:**M**<br> *in* attention_bias:**T**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
-|PackedSparseAttentionIndexer|*in* query:**T**<br> *in* key:**T**<br> *in* query_norm_weight:**T**<br> *in* key_norm_weight:**T**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* cumulative_sequence_lengths:**M**<br> *in* past_sequence_lengths:**M**<br> *in* gate:**T**<br> *in* position_bias:**T**<br> *in* head_weights:**T**<br> *in* position_ids:**I**<br> *in* past_key_state:**T**<br> *in* past_kv_buffer:**T**<br> *in* past_gate_buffer:**T**<br> *in* past_state_lengths:**M**<br> *out* selected_indices:**M**<br> *out* selected_counts:**M**<br> *out* present_key_state:**T**<br> *out* present_kv_buffer:**T**<br> *out* present_gate_buffer:**T**<br> *out* present_state_lengths:**M**|1+|**I** = tensor(int64)<br/> **M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)|
+|PackedSparseAttentionIndexer|*in* query:**T**<br> *in* key:**T**<br> *in* query_norm_weight:**T**<br> *in* key_norm_weight:**T**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* cumulative_sequence_lengths:**M**<br> *in* past_sequence_lengths:**M**<br> *in* gate:**T**<br> *in* position_bias:**T**<br> *in* head_weights:**T**<br> *in* position_ids:**I**<br> *in* past_key_state:**T**<br> *in* past_kv_buffer:**T**<br> *in* past_gate_buffer:**T**<br> *in* past_state_lengths:**M**<br> *in* state_update_capture_count:**M**<br> *in* state_update_active:**M**<br> *out* selected_indices:**M**<br> *out* selected_counts:**M**<br> *out* present_key_state:**T**<br> *out* present_kv_buffer:**T**<br> *out* present_gate_buffer:**T**<br> *out* present_state_lengths:**M**<br> *out* state_update:**T**|1+|**I** = tensor(int64)<br/> **M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)|
 |PagedAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* key_cache:**T_CACHE**<br> *in* value_cache:**T_CACHE**<br> *in* cumulative_sequence_length:**S**<br> *in* past_seqlens:**S**<br> *in* block_table:**S**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* slot_mapping:**S**<br> *in* head_sink:**T**<br> *in* q_norm_weight:**T**<br> *in* k_norm_weight:**T**<br> *in* k_scale:**T_KV_SCALE**<br> *in* v_scale:**T_KV_SCALE**<br> *in* attention_metadata:**S**<br> *out* output:**T**<br> *out* key_cache_out:**T_CACHE**<br> *out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float16)<br/> **T_CACHE** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(int8), tensor(uint8)<br/> **T_KV_SCALE** = tensor(float)|
 |QAttention|*in* input:**T1**<br> *in* weight:**T2**<br> *in* bias:**T3**<br> *in* input_scale:**T3**<br> *in* weight_scale:**T3**<br> *in* mask_index:**T4**<br> *in* input_zero_point:**T1**<br> *in* weight_zero_point:**T2**<br> *in* past:**T3**<br> *out* output:**T3**<br> *out* present:**T3**|1+|**T1** = tensor(int8)<br/> **T2** = tensor(int8)<br/> **T3** = tensor(float), tensor(float16)<br/> **T4** = tensor(int32)|
 |QMoE|*in* input:**T**<br> *in* router_probs:**T**<br> *in* fc1_experts_weights:**T1**<br> *in* fc1_scales:**T2**<br> *in* fc1_experts_bias:**T**<br> *in* fc2_experts_weights:**T1**<br> *in* fc2_scales:**T2**<br> *in* fc2_experts_bias:**T**<br> *in* fc3_experts_weights:**T1**<br> *in* fc3_scales:**T2**<br> *in* fc3_experts_bias:**T**<br> *in* fc1_zero_points:**T1**<br> *in* fc2_zero_points:**T1**<br> *in* fc3_zero_points:**T1**<br> *in* router_weights:**T**<br> *in* fc1_global_scale:**T4**<br> *in* fc2_global_scale:**T4**<br> *in* fc1_act_scale:**T4**<br> *in* fc2_act_scale:**T4**<br> *in* fc1_act_block_scale:**T2**<br> *in* fc2_act_block_scale:**T2**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)<br/> **T1** = tensor(float8e4m3fn), tensor(uint8)<br/> **T2** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(float8e8m0)<br/> **T4** = tensor(float)|

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<T, /*leading=*/true>` | 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,29 +343,41 @@ 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:
 
 ```bash
 ./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.