Skip to content

[CUDA][WebGPU] Add packed sparse attention indexer for continuous batching #32618

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

Open
Copilot wants to merge 16 commits into from
Open

[CUDA][WebGPU] Add packed sparse attention indexer for continuous batching #32618

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
dbaeb30
Initial plan
Copilot Sep 15, 2026
cf54693
Implement packed sparse attention indexer
Copilot Sep 16, 2026
6f8041c
Harden packed indexer state validation
Copilot Sep 16, 2026
5b9d10e
Validate packed indexer capacity bounds
Copilot Sep 16, 2026
de0e587
Fix packed indexer output-count test
Copilot Sep 16, 2026
2125b41
Merge remote-tracking branch 'origin/copilot/implement-sparse-attenti…
Copilot Sep 16, 2026
df26284
Address packed indexer review feedback
Copilot Sep 16, 2026
8ece210
Fix packed indexer CI failures
Copilot Sep 16, 2026
d0b6faf
Fix packed indexer schema validation
Copilot Sep 17, 2026
e81906b
Reduce packed indexer WebGPU bindings
Copilot Sep 17, 2026
a73ff17
Merge branch 'stack-32528' into stack-32618
kunal-vaishnavi Sep 18, 2026
f32b9bf
Fuse query normalization into PackedSparseAttentionIndexer
kunal-vaishnavi Sep 18, 2026
96f00c6
Merge branch 'stack-32528' into stack-32618
kunal-vaishnavi Sep 18, 2026
324128f
Merge branch 'stack-32528' into stack-32618
kunal-vaishnavi Sep 18, 2026
b150814
Merge branch 'stack-32528' into stack-32618
kunal-vaishnavi Sep 19, 2026
7f41361
Merge branch 'stack-32528' into stack-32618
kunal-vaishnavi Sep 20, 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
171 changes: 171 additions & 0 deletions docs/ContribOperators.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ Do not modify directly.*
* <a href="#com.microsoft.NhwcMaxPool">com.microsoft.NhwcMaxPool</a>
* <a href="#com.microsoft.PackedAttention">com.microsoft.PackedAttention</a>
* <a href="#com.microsoft.PackedMultiHeadAttention">com.microsoft.PackedMultiHeadAttention</a>
* <a href="#com.microsoft.PackedSparseAttentionIndexer">com.microsoft.PackedSparseAttentionIndexer</a>
* <a href="#com.microsoft.Pad">com.microsoft.Pad</a>
* <a href="#com.microsoft.PagedAttention">com.microsoft.PagedAttention</a>
* <a href="#com.microsoft.QAttention">com.microsoft.QAttention</a>
Expand Down Expand Up @@ -4326,6 +4327,17 @@ This version of the operator has been available since version 1 of the 'com.micr
past_ids and present_ids may use the same allocation. Such in-place execution is transaction-safe
only when the whole operator call is unconditionally committed; a caller that may select a prefix or
roll back must preserve past_ids.

Optional inputs add packed-sequence and Qwen4-Exp-style n-gram embedding support:

- eos_token_id, when provided together with reset_on_eos != 0, causes causal history to reset at EOS
boundaries: any shifted position at or before the most recent EOS strictly before the current
position is replaced with eos_token_id instead of the real token.
- segment_ids, when provided, additionally resets causal history at any position whose segment id
differs from the immediately preceding position's segment id within input_ids. Segment boundaries
are not checked against past_ids history.
- head_offsets, when provided, adds a fixed per-output-head offset after the modulo by the head's
vocabulary size, letting all heads across all n-gram orders share one flat embedding table.

Optional inputs add packed-sequence and Qwen4-Exp-style n-gram embedding support:

Expand Down Expand Up @@ -4730,6 +4742,163 @@ This version of the operator has been available since version 1 of the 'com.micr
</dl>


### <a name="com.microsoft.PackedSparseAttentionIndexer"></a><a name="com.microsoft.packedsparseattentionindexer">**com.microsoft.PackedSparseAttentionIndexer**</a>

Packed/variable-length counterpart of SparseAttentionIndexer, for continuous-batching engines
(such as an OgaEngine-style PagedAttention model) that flatten every request's tokens into one
[total_tokens, ...] axis instead of a dense [batch_size, sequence_length, ...] axis. It selects, for
every packed query token, the sparse-attention candidates that a following SparsePagedAttention (or
similar) operator is allowed to read.

Unlike SparseAttentionIndexer, this operator:
* takes packed query/key tensors plus cumulative_sequence_lengths (request boundaries) and
past_sequence_lengths (per-request past length) instead of a dense batch and a dense mask;
* derives ordinary causal visibility purely from that packed metadata -- there is no mask input;
* uses a single generic set of state slots (past_key_state / past_kv_buffer / past_gate_buffer /
past_state_lengths) for both policy_mode values, each with a shape that is fixed across calls
(state never grows and is never concatenated); a step that would overflow the fixed capacity is
rejected as a deterministic no-op on state rather than truncated or allowed to corrupt memory;
* additionally emits selected_counts, the exact number of active (non -1) entries per query, so
that no downstream consumer needs to scan selected_indices for its query's true count.

Both policy_mode values keep the semantics of SparseAttentionIndexer, applied independently to each
request's own packed token range and fixed-capacity state slice:

policy_mode = "qsa" ("query sparse attention" token indexer)
Processes each request's new tokens sequentially: appends raw indexer keys to the generic
pending buffer, and whenever it reaches compress_ratio tokens, mean-pools it, applies RMSNorm
and key_norm_weight, applies the leading/split-half rotary convention at the block's first
logical token position, and appends the prepared (already normalized and rotated) key to
key_state. Queries are scored against every causally visible complete block with
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.

policy_mode = "csa" ("compressed sparse attention" block indexer)
Applies the same window-plan arithmetic as SparseAttentionIndexer (overlap/leftover/new window
count) independently per request, using that request's own buffer_length and new token count;
every newly closed window is compressed with the softmax-gated Ca/Cb pooling, normalized,
rotated and appended to key_state. Queries are scored against every causally visible compressed
entry with sum_h w_h * ReLU(q_h . k) and the index_topk highest scoring entry indices are
emitted.

Common contract:
* selected_indices is int32 with a fixed capacity that only depends on attributes:
token_budget + compress_ratio - 1 for "qsa" (values are request-local token positions into the
main key/value cache, directly consumable by SparsePagedAttention configured with
attention_mode="selected_only", selected_kv_source="main") and index_topk for "csa" (values are
compressed-entry indices into key_state, directly consumable by SparsePagedAttention configured
with attention_mode="local_plus_selected", selected_kv_source="auxiliary"; key_state is
layout-compatible with a [batch_size, capacity, 1, head_size] auxiliary cache when K = V).
Unused entries are -1 and selected_counts holds the exact number of used entries.
* key_norm_weight is the effective RMSNorm multiplier, exactly as in SparseAttentionIndexer.
* Accumulation, pooling, softmax, normalization and scoring are performed in float32 and the
result is rounded once to the tensor element type.
* Ties in the top-k selection are broken by the smaller entry index, and the emitted entries are
ordered by decreasing score, so the result is deterministic.
* cos_cache / sin_cache may be shared across the batch ([max_position, rotary_width]) or
request-specific ([batch_size, max_position, rotary_width]).
* cumulative_sequence_lengths, past_sequence_lengths and past_state_lengths are read directly by
the device kernel; a zero-token request row (a repeated cumulative offset) is valid and simply
contributes no query rows for that request.

OgaEngine integration note: this operator only defines the ORT contrib op; wiring
past_key_state / past_kv_buffer / past_gate_buffer / past_state_lengths as Engine-managed,
per-request fixed-size state (analogous to a paged auxiliary cache) is expected to happen in the
OgaEngine / Model Builder integration, which is out of scope for this operator definition.

#### Version

This version of the operator has been available since version 1 of the 'com.microsoft' operator set.

#### Attributes

<dl>
<dt><tt>compress_ratio</tt> : int (required)</dt>
<dd>Number of consecutive tokens folded into one compressed/pooled entry. Must be > 0 and 2 * compress_ratio - 1 must not exceed INT_MAX.</dd>
<dt><tt>epsilon</tt> : float</dt>
<dd>Epsilon of the RMS normalization applied to queries and compressed keys. Default is 1e-6.</dd>
<dt><tt>head_weight_scale</tt> : float</dt>
<dd>Only for policy_mode 'csa': scale applied to head_weights. Default is 1/sqrt(num_heads). Must be omitted when policy_mode is 'qsa'.</dd>
<dt><tt>index_topk</tt> : int</dt>
<dd>Only for policy_mode 'csa': number of compressed entries selected per query. Must be > 0. Must be omitted when policy_mode is 'qsa'.</dd>
<dt><tt>policy_mode</tt> : string (required)</dt>
<dd>Indexer policy. Must be exactly 'qsa' (token indexer) or 'csa' (compressed block indexer).</dd>
<dt><tt>scale</tt> : float</dt>
<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>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

<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>
<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>
<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, 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>cumulative_sequence_lengths</tt> : M</dt>
<dd>Device-resident packed request boundaries with shape (batch_size + 1); cumulative_sequence_lengths[0] must be 0 and cumulative_sequence_lengths[batch_size] must equal total_tokens. Request b owns rows [cumulative_sequence_lengths[b], cumulative_sequence_lengths[b + 1]) of query/key (a repeated offset is a valid zero-token row).</dd>
<dt><tt>past_sequence_lengths</tt> : M</dt>
<dd>Device-resident number of tokens already processed for each request before this call, with shape (batch_size). Used as the default absolute query position when position_ids is omitted (policy_mode 'qsa'), and to validate state consistency.</dd>
<dt><tt>gate</tt> (optional) : T</dt>
<dd>Only for policy_mode 'csa': gate projection of the new tokens with shape (total_tokens, 2 * head_size).</dd>
<dt><tt>position_bias</tt> (optional) : T</dt>
<dd>Only for policy_mode 'csa': per-slot gate bias with shape (compress_ratio, 2 * head_size).</dd>
<dt><tt>head_weights</tt> (optional) : T</dt>
<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>Optional for policy_mode 'qsa', required 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_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>
<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>
</dl>

#### Outputs

<dl>
<dt><tt>selected_indices</tt> : M</dt>
<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>
<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>
<dt><tt>present_gate_buffer</tt> (optional) : T</dt>
<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>
</dl>

#### Type Constraints

<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>I</tt> : tensor(int64)</dt>
<dd>Constrain position ids to 64-bit integer tensors.</dd>
<dt><tt>M</tt> : tensor(int32)</dt>
<dd>Constrain packed metadata, generic state lengths and selected indices/counts to 32-bit integer tensors.</dd>
</dl>


### <a name="com.microsoft.Pad"></a><a name="com.microsoft.pad">**com.microsoft.Pad**</a>

Given `data` tensor, pads, mode, and value.
Expand Down Expand Up @@ -7868,3 +8037,5 @@ No versioning maintained for experimental ops.
<dt><tt>T</tt> : tensor(float)</dt>
<dd>Constrain input and output types to float32 tensors.</dd>
</dl>


1 change: 1 addition & 0 deletions docs/OperatorKernels.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1122,6 +1122,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* 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)|
|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)|
Expand Down
Loading
Loading