Skip to content

[CUDA] Extend DynamicSparseAttention INT8 KV cache and PackedSparseAttentionIndexer QSA #33472

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Akshay Sonawane (apsonawane) merged 26 commits into from
Oct 28, 2026
+3,413 −608
Merged

[CUDA] Extend DynamicSparseAttention INT8 KV cache and PackedSparseAttentionIndexer QSA #33472

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
52c73b5
Extend sparse attention caches and optimize packed indexing
apsonawane Oct 19, 2026
7be8d68
Align QSA position overrides and CUDA-only outputs across providers
apsonawane Oct 19, 2026
b538d1f
Update generated QSA position override documentation
apsonawane Oct 19, 2026
538bd98
Fix sparse attention workspace and optional interface handling
apsonawane Oct 19, 2026
9356555
Fix WebGPU binding-limit test and clarify mixed key-state support
apsonawane Oct 19, 2026
d21946e
Add compact QSA selection and specialized INT8 attention dispatch
apsonawane Oct 21, 2026
b42576e
Merge main into asonawane/main-ops-sparse-attention-indexer
apsonawane Oct 22, 2026
b71498b
Apply lintrunner formatting to PR changes
apsonawane Oct 22, 2026
c0c4e7d
Enable default-on sparse attention optimizations, S<=8 grouped decode…
tianleiwu Oct 22, 2026
6548c21
[CUDA] Fold #33546 and S=4 prefix-state outputs (mid_kv_buffer2/3, mi…
tianleiwu Oct 23, 2026
58b6da6
Merge remote-tracking branch 'origin/main' into tlwu/pr33472-updates
tianleiwu Oct 23, 2026
3b52694
Apply lintrunner formatting
tianleiwu Oct 23, 2026
5ed039c
[CUDA] Revert unused QSA position_ids and WebGPU shader changes in Pa…
tianleiwu Oct 24, 2026
6511e8a
fix(cuda): fix shared-provider include order, Windows macro clash, te…
tianleiwu Oct 24, 2026
a9afd69
Merge remote-tracking branch 'origin/main' into pr33472-work
tianleiwu Oct 24, 2026
a523af4
Optimize GroupedDecodeDynamicSparseAttentionKernel for 12 query heads…
tianleiwu Oct 24, 2026
61c5d87
Remove legacy mid-state snapshot outputs from PackedSparseAttentionIn…
tianleiwu Oct 25, 2026
aa78776
perf(cuda): optimize INT8 KV DynamicSparseAttention prefill and Packe…
tianleiwu Oct 26, 2026
18a13c7
fix(cuda): fix FP16 ratio-12 grouped decode MMA and expand prefill te…
tianleiwu Oct 26, 2026
fea46be
fix(cuda): avoid narrowing in sparse attention test
tianleiwu Oct 26, 2026
9116cb9
[CUDA] Support per-batch committed_length in DynamicSparseAttention (…
tianleiwu Oct 26, 2026
a4151cc
[CUDA] DSA: fail strict validation on inconsistent committed cursor
tianleiwu Oct 26, 2026
977e981
test(cuda): run alternate DSA/indexer switches in fresh processes
tianleiwu Oct 26, 2026
0639efa
docs: document DSA/indexer dispatch switches, INT8 encoding and statu…
tianleiwu Oct 26, 2026
cd4e2e4
[CUDA] Remove packed sparse indexer radix top-K path
tianleiwu Oct 26, 2026
c8c7b77
fix(cuda): guard DSA wide loads on buffer alignment and fix iOS test …
tianleiwu Oct 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 35 additions & 13 deletions docs/ContribOperators.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1789,6 +1789,14 @@ This version of the operator has been available since version 1 of the 'com.micr
Auxiliary K/V use BNSH layout and are read-only. When `auxiliary_kv_shared=1`, auxiliary_value may be omitted and
auxiliary_key is used as both K and V. DeepSeek-style post-attention inverse/conjugate RoPE and output projection remain
the exporter's responsibility.

Main caches may instead use INT8 with positive finite float32 key_scale and value_scale vectors, one scale per
KV-head channel. The cache is symmetric signed int8 without a zero point, like the 8-bit KV cache of GroupQueryAttention
and PagedAttention; tensor(uint8) is reserved for packed 4-bit caches in those operators and is not accepted here.
Newly appended keys (after optional normalization and RoPE) and values are divided by their scales,
rounded to nearest with ties to even, and saturated to [-128, 127]. Attention reads the quantized representation,
including newly appended rows. Auxiliary caches remain floating point. Invalid scale values produce validation_status=10
when validation is requested; CUDA strict validation also reports them as an error.

#### Version

Expand Down Expand Up @@ -1825,7 +1833,7 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Add a zero-valued sink logit when no explicit head_sink is supplied.</dd>
</dl>

#### Inputs (11 - 17)
#### Inputs (11 - 20)

<dl>
<dt><tt>query</tt> : T</dt>
Expand All @@ -1834,10 +1842,10 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Current main key [batch, sequence, kv_num_heads * head_size].</dd>
<dt><tt>value</tt> (optional) : T</dt>
<dd>Current main value [batch, sequence, kv_num_heads * head_size].</dd>
<dt><tt>past_key</tt> (optional) : T</dt>
<dd>Main key cache in BNSH layout.</dd>
<dt><tt>past_value</tt> (optional) : T</dt>
<dd>Main value cache in BNSH layout.</dd>
<dt><tt>past_key</tt> (optional) : TCache</dt>
<dd>Main key cache in BNSH layout, same type as query or INT8.</dd>
<dt><tt>past_value</tt> (optional) : TCache</dt>
<dd>Main value cache in BNSH layout, same type as past_key.</dd>
<dt><tt>auxiliary_key</tt> (optional) : T</dt>
<dd>Read-only auxiliary key sequence in BNSH layout.</dd>
<dt><tt>auxiliary_value</tt> (optional) : T</dt>
Expand All @@ -1862,24 +1870,36 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Optional K RMSNorm weight [head_size].</dd>
<dt><tt>head_sink</tt> (optional) : T</dt>
<dd>Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.</dd>
<dt><tt>committed_length</tt> (optional) : M</dt>
<dd>Optional per-request committed token count ([1] or [batch]) for cursor validation. When supplied, validation_status is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent, and CUDA strict validation fails with a cursor-specific error. CUDA only.</dd>
<dt><tt>key_scale</tt> (optional) : S</dt>
<dd>Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].</dd>
<dt><tt>value_scale</tt> (optional) : S</dt>
<dd>Positive finite per-channel INT8 value-cache scales [kv_num_heads * head_size].</dd>
</dl>

#### Outputs (1 - 3)
#### Outputs (1 - 4)

<dl>
<dt><tt>output</tt> : T</dt>
<dd>Attention output [batch, sequence, num_heads * head_size].</dd>
<dt><tt>present_key</tt> (optional) : T</dt>
<dt><tt>present_key</tt> (optional) : TCache</dt>
<dd>Updated main key cache in BNSH layout.</dd>
<dt><tt>present_value</tt> (optional) : T</dt>
<dt><tt>present_value</tt> (optional) : TCache</dt>
<dd>Updated main value cache in BNSH layout.</dd>
<dt><tt>validation_status</tt> (optional) : M</dt>
<dd>Optional CUDA-only int32 status: 0 when valid, otherwise a validation code (1-8: selection, sequence-length or rotary-position errors; 9: inconsistent committed_length cursor; 10: invalid INT8 cache scale).</dd>
</dl>

#### Type Constraints

<dl>
<dt><tt>T</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
<dd>Constrain all floating-point inputs and outputs to one element type.</dd>
<dt><tt>TCache</tt> : tensor(float), tensor(float16), tensor(bfloat16), tensor(int8)</dt>
<dd>Cache type must match T unless INT8 per-channel scales are supplied.</dd>
<dt><tt>S</tt> : tensor(float)</dt>
<dd>INT8 cache scales use float32.</dd>
<dt><tt>M</tt> : tensor(int32)</dt>
<dd>Constrain selection and sequence metadata to int32.</dd>
</dl>
Expand Down Expand Up @@ -5280,8 +5300,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 the request-local logical cache positions derived from
past_sequence_lengths and cumulative_sequence_lengths; position_ids must be omitted.
incomplete block. There is no position_ids input: QSA positions are the request-local logical
cache positions derived from past_sequence_lengths and cumulative_sequence_lengths.

policy_mode = "csa" ("compressed sparse attention" block indexer)
Applies the same window-plan arithmetic as SparseAttentionIndexer (overlap/leftover/new window
Expand Down Expand Up @@ -5375,8 +5395,8 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Only for policy_mode 'csa': per-head score weights with shape (total_tokens, num_heads).</dd>
<dt><tt>position_ids</tt> (optional) : I</dt>
<dd>Only for policy_mode 'csa': absolute position of every packed query, with shape (total_tokens).</dd>
<dt><tt>past_key_state</tt> : T</dt>
<dd>Generic fixed-capacity state: policy_mode 'qsa' stores prepared complete-block keys; policy_mode 'csa' stores compressed keys. Shape is (batch_size, state_capacity, head_size) and never changes across calls.</dd>
<dt><tt>past_key_state</tt> : T_KEY</dt>
<dd>Generic fixed-capacity state: policy_mode 'qsa' stores prepared complete-block keys; policy_mode 'csa' stores compressed keys. Shape is (batch_size, state_capacity, head_size) and never changes across calls. CUDA QSA also supports FLOAT key state with BFLOAT16 query/key tensors. WebGPU requires key state and activations to have the same type.</dd>
<dt><tt>past_kv_buffer</tt> : T</dt>
<dd>Generic fixed-capacity pending-token buffer. Shape is (batch_size, 2 * compress_ratio - 1, head_size) for policy_mode 'qsa' (which only ever uses up to compress_ratio - 1 of these entries) and (batch_size, 2 * compress_ratio - 1, 2 * head_size) for policy_mode 'csa'.</dd>
<dt><tt>past_gate_buffer</tt> (optional) : T</dt>
Expand Down Expand Up @@ -5404,7 +5424,7 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Selected entries with shape (total_tokens, capacity). capacity is token_budget + compress_ratio - 1 for policy_mode 'qsa' (request-local token positions into the main key/value cache) and index_topk for policy_mode 'csa' (compressed entry indices into key_state). Unused entries are -1.</dd>
<dt><tt>selected_counts</tt> : M</dt>
<dd>Exact number of used (non -1) entries of selected_indices for every query, with shape (total_tokens).</dd>
<dt><tt>present_key_state</tt> : T</dt>
<dt><tt>present_key_state</tt> : T_KEY</dt>
<dd>Updated generic key state, with the same fixed shape as past_key_state.</dd>
<dt><tt>present_kv_buffer</tt> : T</dt>
<dd>Updated generic pending-token buffer, with the same fixed shape as past_kv_buffer.</dd>
Expand All @@ -5425,6 +5445,8 @@ This version of the operator has been available since version 1 of the 'com.micr
<dl>
<dt><tt>T</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
<dd>Constrain floating point tensors to float, float16 and bfloat16.</dd>
<dt><tt>T_KEY</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
<dd>Constrain the generic key-state tensor to float, float16 and bfloat16.</dd>
<dt><tt>I</tt> : tensor(int64)</dt>
<dd>Constrain position ids to 64-bit integer tensors.</dd>
<dt><tt>M</tt> : tensor(int32)</dt>
Expand Down
4 changes: 2 additions & 2 deletions docs/OperatorKernels.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1112,7 +1112,7 @@ The **OpSet Version** column uses the following notation:
|DecoderMaskedSelfAttention|*in* input:**T**<br> *in* weights:**T**<br> *in* bias:**T**<br> *in* mask_index:**M**<br> *in* past:**T**<br> *in* attention_bias:**T**<br> *in* past_sequence_length:**M**<br> *in* beam_width:**M**<br> *in* cache_indirection:**M**<br> *out* output:**T**<br> *out* present:**T**|1+|**T** = tensor(float), tensor(float16)|
|DequantizeLinear|*in* x:**T1**<br> *in* x_scale:**T2**<br> *in* x_zero_point:**T1**<br> *out* y:**T2**|1+|**T1** = tensor(int8), tensor(uint8)<br/> **T2** = tensor(float16)|
|DequantizeWithOrder|*in* input:**Q**<br> *in* scale_input:**S**<br> *out* output:**F**|1+|**F** = tensor(float), tensor(float16)<br/> **Q** = tensor(int8)<br/> **S** = tensor(float)|
|DynamicSparseAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* auxiliary_key:**T**<br> *in* auxiliary_value:**T**<br> *in* selected_indices:**M**<br> *in* selected_counts:**M**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* position_ids:**tensor(int64)**<br> *in* q_norm_weight:**T**<br> *in* k_norm_weight:**T**<br> *in* head_sink:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)|
|DynamicSparseAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**TCache**<br> *in* past_value:**TCache**<br> *in* auxiliary_key:**T**<br> *in* auxiliary_value:**T**<br> *in* selected_indices:**M**<br> *in* selected_counts:**M**<br> *in* seqlens_k:**M**<br> *in* total_sequence_length:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* position_ids:**tensor(int64)**<br> *in* q_norm_weight:**T**<br> *in* k_norm_weight:**T**<br> *in* head_sink:**T**<br> *in* committed_length:**M**<br> *in* key_scale:**S**<br> *in* value_scale:**S**<br> *out* output:**T**<br> *out* present_key:**TCache**<br> *out* present_value:**TCache**<br> *out* validation_status:**M**|1+|**M** = tensor(int32)<br/> **S** = tensor(float)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **TCache** = tensor(bfloat16), tensor(float), tensor(float16), tensor(int8)|
|DynamicTimeWarping|*in* input:**F**<br> *out* output:**I**|1+|**F** = tensor(float)<br/> **I** = tensor(int32)|
|EmbedLayerNormalization|*in* input_ids:**T1**<br> *in* segment_ids:**T1**<br> *in* word_embedding:**T**<br> *in* position_embedding:**T**<br> *in* segment_embedding:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* mask:**T1**<br> *in* position_ids:**T1**<br> *out* output:**T**<br> *out* mask_index:**T1**<br> *out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|EngramGate|*in* key:**T**<br> *in* query:**T**<br> *in* value:**T**<br> *in* key_norm_scale:**T**<br> *in* query_norm_scale:**T**<br> *in* conv_norm_scale:**T**<br> *out* output:**T**<br> *out* gated_value_normed:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
Expand Down Expand Up @@ -1151,7 +1151,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> *in* state_update_capture_count:**M**<br> *in* state_update_active:**M**<br> *in* indexshare_mode:**M**<br> *in* merged_indices:**M**<br> *in* merged_counts:**M**<br> *in* merged_status:**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**<br> *out* indexshare_status:**M**<br> *out* state_status:**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_KEY**<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> *in* indexshare_mode:**M**<br> *in* merged_indices:**M**<br> *in* merged_counts:**M**<br> *in* merged_status:**M**<br> *out* selected_indices:**M**<br> *out* selected_counts:**M**<br> *out* present_key_state:**T_KEY**<br> *out* present_kv_buffer:**T**<br> *out* present_gate_buffer:**T**<br> *out* present_state_lengths:**M**<br> *out* state_update:**T**<br> *out* indexshare_status:**M**<br> *out* state_status:**M**|1+|**I** = tensor(int64)<br/> **M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **T_KEY** = tensor(bfloat16), tensor(float), tensor(float16)|
|PackedSparseAttentionIndexerMerge|*in* base_indices:**T**<br> *in* base_counts:**T**<br> *in* base_row_indices:**T**<br> *in* range_starts:**T**<br> *in* range_ends:**T**<br> *in* additional_indices:**T**<br> *in* additional_counts:**T**<br> *out* selected_indices:**T**<br> *out* selected_counts:**T**<br> *out* status:**T**|1+|**T** = tensor(int32)|
|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)|
Expand Down
Loading
Loading