diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index f84d7593141f2..459e4cf71a98b 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -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
@@ -1825,7 +1833,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Add a zero-valued sink logit when no explicit head_sink is supplied.
-#### Inputs (11 - 17)
+#### Inputs (11 - 20)
- query : T
@@ -1834,10 +1842,10 @@ This version of the operator has been available since version 1 of the 'com.micr
- Current main key [batch, sequence, kv_num_heads * head_size].
- value (optional) : T
- Current main value [batch, sequence, kv_num_heads * head_size].
-- past_key (optional) : T
-- Main key cache in BNSH layout.
-- past_value (optional) : T
-- Main value cache in BNSH layout.
+- past_key (optional) : TCache
+- Main key cache in BNSH layout, same type as query or INT8.
+- past_value (optional) : TCache
+- Main value cache in BNSH layout, same type as past_key.
- auxiliary_key (optional) : T
- Read-only auxiliary key sequence in BNSH layout.
- auxiliary_value (optional) : T
@@ -1862,17 +1870,25 @@ This version of the operator has been available since version 1 of the 'com.micr
- Optional K RMSNorm weight [head_size].
- head_sink (optional) : T
- Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.
+- committed_length (optional) : M
+- 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.
+- key_scale (optional) : S
+- Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].
+- value_scale (optional) : S
+- Positive finite per-channel INT8 value-cache scales [kv_num_heads * head_size].
-#### Outputs (1 - 3)
+#### Outputs (1 - 4)
- output : T
- Attention output [batch, sequence, num_heads * head_size].
-- present_key (optional) : T
+- present_key (optional) : TCache
- Updated main key cache in BNSH layout.
-- present_value (optional) : T
+- present_value (optional) : TCache
- Updated main value cache in BNSH layout.
+- validation_status (optional) : M
+- 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).
#### Type Constraints
@@ -1880,6 +1896,10 @@ This version of the operator has been available since version 1 of the 'com.micr
- T : tensor(float), tensor(float16), tensor(bfloat16)
- Constrain all floating-point inputs and outputs to one element type.
+- TCache : tensor(float), tensor(float16), tensor(bfloat16), tensor(int8)
+- Cache type must match T unless INT8 per-channel scales are supplied.
+- S : tensor(float)
+- INT8 cache scales use float32.
- M : tensor(int32)
- Constrain selection and sequence metadata to int32.
@@ -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
@@ -5375,8 +5395,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Only for policy_mode 'csa': per-head score weights with shape (total_tokens, num_heads).
position_ids (optional) : I
Only for policy_mode 'csa': absolute position of every packed query, with shape (total_tokens).
-past_key_state : T
-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.
+past_key_state : T_KEY
+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.
past_kv_buffer : T
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'.
past_gate_buffer (optional) : T
@@ -5404,7 +5424,7 @@ This version of the operator has been available since version 1 of the 'com.micr
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.
selected_counts : M
Exact number of used (non -1) entries of selected_indices for every query, with shape (total_tokens).
-present_key_state : T
+present_key_state : T_KEY
Updated generic key state, with the same fixed shape as past_key_state.
present_kv_buffer : T
Updated generic pending-token buffer, with the same fixed shape as past_kv_buffer.
@@ -5425,6 +5445,8 @@ This version of the operator has been available since version 1 of the 'com.micr
- T : tensor(float), tensor(float16), tensor(bfloat16)
- Constrain floating point tensors to float, float16 and bfloat16.
+- T_KEY : tensor(float), tensor(float16), tensor(bfloat16)
+- Constrain the generic key-state tensor to float, float16 and bfloat16.
- I : tensor(int64)
- Constrain position ids to 64-bit integer tensors.
- M : tensor(int32)
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 2497a7e6fd31a..cd0be24517121 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -1112,7 +1112,7 @@ The **OpSet Version** column uses the following notation:
|DecoderMaskedSelfAttention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* mask_index:**M**
*in* past:**T**
*in* attention_bias:**T**
*in* past_sequence_length:**M**
*in* beam_width:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present:**T**|1+|**T** = tensor(float), tensor(float16)|
|DequantizeLinear|*in* x:**T1**
*in* x_scale:**T2**
*in* x_zero_point:**T1**
*out* y:**T2**|1+|**T1** = tensor(int8), tensor(uint8)
**T2** = tensor(float16)|
|DequantizeWithOrder|*in* input:**Q**
*in* scale_input:**S**
*out* output:**F**|1+|**F** = tensor(float), tensor(float16)
**Q** = tensor(int8)
**S** = tensor(float)|
-|DynamicSparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* selected_indices:**M**
*in* selected_counts:**M**
*in* seqlens_k:**M**
*in* total_sequence_length:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**tensor(int64)**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* head_sink:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)|
+|DynamicSparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**TCache**
*in* past_value:**TCache**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* selected_indices:**M**
*in* selected_counts:**M**
*in* seqlens_k:**M**
*in* total_sequence_length:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**tensor(int64)**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* head_sink:**T**
*in* committed_length:**M**
*in* key_scale:**S**
*in* value_scale:**S**
*out* output:**T**
*out* present_key:**TCache**
*out* present_value:**TCache**
*out* validation_status:**M**|1+|**M** = tensor(int32)
**S** = tensor(float)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TCache** = tensor(bfloat16), tensor(float), tensor(float16), tensor(int8)|
|DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)|
|EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
@@ -1151,7 +1151,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**
*in* state_update_capture_count:**M**
*in* state_update_active:**M**
*in* indexshare_mode:**M**
*in* merged_indices:**M**
*in* merged_counts:**M**
*in* merged_status:**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**
*out* indexshare_status:**M**
*out* state_status:**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_KEY**
*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**
*in* indexshare_mode:**M**
*in* merged_indices:**M**
*in* merged_counts:**M**
*in* merged_status:**M**
*out* selected_indices:**M**
*out* selected_counts:**M**
*out* present_key_state:**T_KEY**
*out* present_kv_buffer:**T**
*out* present_gate_buffer:**T**
*out* present_state_lengths:**M**
*out* state_update:**T**
*out* indexshare_status:**M**
*out* state_status:**M**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**T_KEY** = tensor(bfloat16), tensor(float), tensor(float16)|
|PackedSparseAttentionIndexerMerge|*in* base_indices:**T**
*in* base_counts:**T**
*in* base_row_indices:**T**
*in* range_starts:**T**
*in* range_ends:**T**
*in* additional_indices:**T**
*in* additional_counts:**T**
*out* selected_indices:**T**
*out* selected_counts:**T**
*out* status:**T**|1+|**T** = tensor(int32)|
|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)|
diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md
new file mode 100644
index 0000000000000..86ca09455b0a3
--- /dev/null
+++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md
@@ -0,0 +1,100 @@
+# CUDA DynamicSparseAttention grouped decode
+
+`UseGroupedDecodeAttention` enables the split-K grouped decode attention kernel
+for decode and multi-row speculative verification (`1 <= sequence_length <= 8`)
+across both floating-point (`FP16`/`BF16`) and quantized (`INT8`) KV caches,
+with `max_splits = 32` for quantized KV caches.
+
+The INT8 grouped-decode and vector-staging reference tests cover one through
+four query rows (`sequence_length in {1, 2, 3, 4}`).
+
+## INT8 KV cache encoding
+
+The INT8 main cache is `tensor(int8)`: symmetric signed storage with no zero point,
+`q = clamp(rint(x / scale), -128, 127)` and `x = q * scale`, with one float32 scale per
+`kv_head * head_size` channel (`key_scale`, `value_scale`). This is the same encoding as the 8-bit
+KV cache of `GroupQueryAttention` (`group_query_attention_qdq.cuh`) and `PagedAttention` /
+`SparsePagedAttention` (`paged_attention_impl.cu`), which also use `int8`. `uint8` is used by those
+operators only for packed 4-bit caches (`kv_cache_bit_width=4`), so `DynamicSparseAttention` does not
+accept `tensor(uint8)`.
+
+## Validation status codes
+
+`validation_status` (and the CUDA strict-validation host check enabled with
+`ORT_DYNAMIC_SPARSE_ATTENTION_STRICT_VALIDATION=1`) uses these codes; the first error found wins,
+except that a cursor error (9) is replaced by any selection/sequence/position/scale error.
+
+| Code | Enumerator | Meaning |
+|---|---|---|
+| 0 | `kDynamicSparseAttentionValidationOk` | Valid. |
+| 1 | `kDynamicSparseAttentionInvalidCount` | `selected_counts` outside `[0, max_selected]`. |
+| 2 | `kDynamicSparseAttentionInvalidPadding` | Entries after `selected_counts` are not `-1`. |
+| 3 | `kDynamicSparseAttentionNegativeIndex` | Active selected index is negative. |
+| 4 | `kDynamicSparseAttentionIndexOutOfBounds` | Selected index outside the source sequence. |
+| 5 | `kDynamicSparseAttentionDuplicateIndex` | Duplicate active selected index. |
+| 6 | `kDynamicSparseAttentionNonCausalIndex` | Main selected index refers to a future key. |
+| 7 | `kDynamicSparseAttentionInvalidSequenceLength` | `seqlens_k` incompatible with the sequence and cache capacity. |
+| 8 | `kDynamicSparseAttentionInvalidPosition` | Rotary position outside the rotary cache. |
+| 9 | `kDynamicSparseAttentionInvalidCommittedLength` | `committed_length` disagrees with `seqlens_k`, `sequence_length` or `cache_capacity`. |
+| 10 | `kDynamicSparseAttentionInvalidCacheScale` | INT8 scale is not positive and finite. |
+
+With strict validation enabled every non-zero code, including 9, fails the kernel with an
+`INVALID_ARGUMENT` status even when `validation_status` is not connected. Without strict validation
+the code is only reported through `validation_status`.
+
+## CUDA dispatch switches
+
+`ORT_DSA_INT8_VECTOR_STAGING` and `ORT_DSA_SHARED_SCALE_MERGE` are read once per process (cached in
+function-local statics) when the grouped decode path is first launched, so set them before the first
+operator invocation. Only `0` and `1` are accepted. Neither changes schemas or cache layouts.
+`ORT_DSA_PREFILL_KERNEL` (default `1`) is read on every BF16 call.
+
+Grouped decode is used when `1 <= sequence_length <= 8`, `num_heads / kv_num_heads >= 4` (and
+divisible), `head_size == 256`, 16-bit `T` (FP16/BF16), more than 256 candidates
+(`max_selected` plus the local window for `local_plus_selected`) and the grouped shared-memory
+footprint fits the per-block limit. The kernel uses 12 query-head warps when the ratio is exactly 12
+and 8 otherwise.
+
+| Variable | Default | Applies to | Value | Kernels selected |
+|---|---|---|---|---|
+| `ORT_DSA_INT8_VECTOR_STAGING` | `1` | Grouped decode with an INT8 main cache | `1` | `GroupedDecodeDynamicSparseAttentionKernel<..., kQuantizedCache=true, kVectorizedStaging=true>`: 8-byte `uint2` cache loads with `float4` scale loads. |
+| | | | `0` | Same kernel with `kVectorizedStaging=false`: scalar per-channel dequantization into shared memory. |
+| `ORT_DSA_SHARED_SCALE_MERGE` | `1` | Grouped decode (INT8 and floating-point caches) when `split_count <= 32` (always for INT8, whose `max_splits` is 32; for floating-point caches `split_count = min(ceil(candidates / 64), ceil(128 / grouped_blocks))`) | `1` | `MergeGroupedDecodeDynamicSparseAttentionKernel`: one `expf` per split shared across channels; a split is active when its partial sum is non-zero. |
+| | | | `0` | `MergeGroupedDecodeDynamicSparseAttentionKernel`: per-channel `expf` over all splits. |
+| `ORT_DSA_PREFILL_KERNEL` | `1` | BF16, not grouped decode, `sequence_length >= 8`, `head_size == 256`, `selected_only` from the main cache, `smooth_softmax = 0`, SM80+ with enough opt-in shared memory | `1` | `DynamicSparseAttentionPrefillKernel` tensor-core prefill. |
+| | | | `0` | Generic fused/split or grouped prefill kernels. |
+
+When `split_count > 32` (floating-point caches only) the generic
+`MergeDynamicSparseAttentionKernel` is used regardless of `ORT_DSA_SHARED_SCALE_MERGE`. The tests
+`Int8CacheGroupedDecodeAlternateSwitchesInFreshProcess_CUDA` re-run the INT8 grouped-decode and merge
+cases in fresh processes with each non-default combination.
+
+### Buffer alignment
+
+The wide-load fast paths assume more alignment than a tensor's element type guarantees, and tensors
+bound through IOBinding can start at any element offset of a larger allocation. The launcher therefore
+checks the caller-provided buffers and uses an equivalent narrower path, with identical results, when a
+buffer is misaligned:
+
+| Fast path | Buffers that must be aligned | Fallback |
+|---|---|---|
+| INT8 grouped decode vector staging (`ORT_DSA_INT8_VECTOR_STAGING=1`) | `key_scale`, `value_scale`, `auxiliary_key`, `auxiliary_value`: 16 bytes; `present_key`, `present_value`: 8 bytes | Scalar per-channel staging (`kVectorizedStaging=false`). |
+| Floating-point grouped decode and grouped prefill | Each staged K/V row: 16 bytes (checked per row in the kernel) | Element-wise copy into shared memory. |
+| Tensor-core BF16 prefill (`ORT_DSA_PREFILL_KERNEL=1`) | `present_key`, `present_value`: 16 bytes (8 for INT8); `key_scale`, `value_scale`: 16 bytes; `output`: 4 bytes | Generic fused/split or grouped prefill kernels. |
+
+Buffers allocated by ONNX Runtime are 256-byte aligned, so only IOBinding with offset buffers takes a
+fallback. The `*MisalignedBuffers_CUDA` tests cover each case.
+
+### Benchmark
+
+rc3 model (INT4), H200, clean machine, 5 repetitions, baseline noise about +-0.03 ms, DFlash2 S=8
+decode ms/cycle at context 2K/8K/32K. Both defaults are the faster setting, so the switches remain
+only as fallbacks:
+
+| Setting | Decode vs default |
+|---|---|
+| `ORT_DSA_INT8_VECTOR_STAGING=0` | about +0.3 ms (+2%) slower |
+| `ORT_DSA_SHARED_SCALE_MERGE=0` | about +0.05 to +0.1 ms (+0.3 to 0.5%) slower |
+
+The corresponding indexer switch results are in
+[../packed_sparse_attention_indexer.md](../packed_sparse_attention_indexer.md#benchmark).
diff --git a/docs/contrib_ops/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md
index 658930a42d55e..6c47de3386eb9 100644
--- a/docs/contrib_ops/packed_sparse_attention_indexer.md
+++ b/docs/contrib_ops/packed_sparse_attention_indexer.md
@@ -1,5 +1,44 @@
# PackedSparseAttentionIndexer โ Operator Documentation
+## CUDA QSA dispatch switches
+
+Two process-level environment variables select alternate CUDA QSA kernels. Each value is read once
+per process on the first eligible call and cached, so set them before the first operator invocation.
+None changes schemas, cache layouts, causal visibility, stable tie ordering, selected counts or
+state-update semantics. The companion `DynamicSparseAttention` switches are documented in
+[cuda/dynamic_sparse_attention.md](cuda/dynamic_sparse_attention.md).
+
+Hierarchical QSA TopK is eligible when `1 <= total_tokens <= 64`, `state_capacity >= 2048`,
+`head_size == 128`, `compress_ratio == 4`, `num_heads == 4` and `32 < block_topk <= 512`. Ragged QSA
+prefill is eligible when `total_tokens > 64`, `head_size == 128`, `compress_ratio == 4` and
+`num_heads == 4`. Other geometries never read these switches.
+
+| Variable | Default | Applies to | Value | Kernels selected |
+|---|---|---|---|---|
+| `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK` | `1` (only `0`/`1` accepted) | hierarchical QSA TopK | `1` | Compacts each 2048-key chunk to 512 keys (`QsaCompactTileTopKKernel`). Up to five chunks use the fused `QsaFinalMergeAndEmitKernel`; more chunks use pairwise `QsaMergeTileTopKKernel` followed by `QsaEmitHierarchicalTopKKernel`. |
+| | | | `0` | Tile-level merge: `QsaMergeTwoLevelsTopKKernel` (4-way) while two levels remain, then `QsaMergeTileTopKKernel`, then `QsaEmitHierarchicalTopKKernel`. |
+| `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS` | `1` (boolean) | ragged QSA prefill | `1` | `QsaRaggedPrefillScoreKernel<..., true>`: one warp per state block keeps all four heads (256 threads per CTA). |
+| | | | `0` | `QsaRaggedPrefillScoreKernel<...>` with one warp per (state block, head) and CTA-wide barriers (1024 threads per CTA). |
+
+The packed-QSA tests re-run the hierarchical, stable-tie, partial-merge-group
+and ragged-prefill cases in fresh processes for every non-default mode
+(`COMPACT_TOPK=0`, `WARP_HEADS=0`; `QsaAlternateTopKModesInFreshProcess`).
+
+### Benchmark
+
+rc3 model (INT4), H200, clean machine, 5 repetitions, baseline noise about +-0.03 ms. DFlash2 S=8
+decode is ms/cycle at context 2K/8K/32K; TTFT is at 2K/8K/32K prompt. Defaults are fastest, so
+both switches default to the fast path and remain as fallbacks:
+
+| Setting | Result vs default |
+|---|---|
+| `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=0` | decode about +0.2 to +0.27 ms (+1.5%) slower |
+| `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=0` | TTFT +3.7% at 8K, +14% at 32K |
+
+A live-key radix selection alternative was benchmarked (default 14.93/15.07/15.30 ms versus
+15.08/15.22/15.64 ms with it, i.e. +0.15 to +0.34 ms slower) and removed. Qualify accuracy and
+performance on the target runtime, hardware and workload before changing a default.
+
This document describes the `com.microsoft::PackedSparseAttentionIndexer` contrib operator: the
packed/variable-length counterpart of `com.microsoft::SparseAttentionIndexer`, built for
continuous-batching (paged) inference engines such as an OgaEngine-style `PagedAttention` model.
@@ -71,8 +110,8 @@ later slot keeps its fixed index.
| 8 | `gate` | `(total_tokens, 2*head_size)` | T | csa only |
| 9 | `position_bias` | `(compress_ratio, 2*head_size)` | T | csa only |
| 10 | `head_weights` | `(total_tokens, num_heads)` | T | csa only |
-| 11 | `position_ids` | `(total_tokens)` | int64 | omitted qsa / required csa |
-| 12 | `past_key_state` | `(batch_size, state_capacity, head_size)` | T | both (generic) |
+| 11 | `position_ids` | `(total_tokens)` | int64 | csa only |
+| 12 | `past_key_state` | `(batch_size, state_capacity, head_size)` | T_KEY | both (generic) |
| 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) |
@@ -86,13 +125,17 @@ empty output name for `qsa`; `state_update` is declared only when QSA capture is
|---|---|---|---|---|
| 0 | `selected_indices` | `(total_tokens, capacity)` | int32 | both |
| 1 | `selected_counts` | `(total_tokens)` | int32 | both |
-| 2 | `present_key_state` | same shape as `past_key_state` | T | both |
+| 2 | `present_key_state` | same shape as `past_key_state` | T_KEY | both |
| 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 |
+| 7 | `index_share_status` | `(total_tokens)` | int32 | index sharing |
| 8 | `state_status` | `(batch_size)` | int32 | optional, CUDA only |
+CUDA supports `T_KEY = T` and, for QSA only, BF16 activations with FP32 key state.
+The pending buffers and captured state updates remain in the activation type.
+
`capacity` is `token_budget + compress_ratio - 1` for `qsa` and `index_topk` for `csa`, exactly the
same formula (`SelectedCapacity`) used by `SparseAttentionIndexer`.
@@ -208,9 +251,7 @@ dense `qsa` policy but against fixed-capacity state instead of a growing cache:
`past_sequence_lengths[b] + local_offset` โ i.e. the request's own absolute token position, which
is exactly what a per-request main paged KV cache is addressed by. The output is therefore directly
consumable by `SparsePagedAttention` configured with `attention_mode="selected_only"`,
-`selected_kv_source="main"`. `position_ids` must be omitted for `qsa`; allowing an unrelated
-absolute-position space would make these cache indices ambiguous and would rotate queries
-inconsistently with prepared block keys.
+`selected_kv_source="main"`.
## 6. Policy `csa`
@@ -228,7 +269,7 @@ mapping are packed:
rejects (caps) new windows beyond `state_capacity` as described in [ยง3](#3-generic-state-shared-by-both-policies).
2. **Score** (one launch per query token, per compressed entry): every entry is scored with
`sum_h w_h * ReLU(q_h . k)` and masked by the causal threshold from `position_ids` (required for
- `csa`; `qsa` always derives request-local logical positions from `past_sequence_lengths`).
+ `csa`; must be omitted for `qsa`).
3. **Select** (one launch per query token): keeps the `index_topk` highest scoring, causally
visible entries and writes the exact active count to `selected_counts`.
@@ -293,6 +334,9 @@ mechanics with no packed equivalent):
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.
+- The CUDA dispatch switches `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK`
+ and `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS` are described in
+ the "CUDA QSA dispatch switches" section at the top of this document.
- 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/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc
index 878659302a19a..d3e020eb16e2e 100644
--- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc
@@ -19,19 +19,21 @@ namespace cuda {
using namespace onnxruntime::cuda;
-#define REGISTER_KERNEL_TYPED(T) \
- ONNX_OPERATOR_TYPED_KERNEL_EX( \
- DynamicSparseAttention, \
- kMSDomain, \
- 1, \
- T, \
- kCudaExecutionProvider, \
- (*KernelDefBuilder::Create()) \
- .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
- .TypeConstraint("M", DataTypeImpl::GetTensorType()) \
- .MayInplace(3, 1) \
- .MayInplace(4, 2) \
- .InputMemoryType(OrtMemTypeCPUInput, 10), \
+#define REGISTER_KERNEL_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ DynamicSparseAttention, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCudaExecutionProvider, \
+ (*KernelDefBuilder::Create()) \
+ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("TCache", BuildKernelDefConstraints()) \
+ .TypeConstraint("S", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("M", DataTypeImpl::GetTensorType()) \
+ .MayInplace(3, 1) \
+ .MayInplace(4, 2) \
+ .InputMemoryType(OrtMemTypeCPUInput, 10), \
DynamicSparseAttention);
REGISTER_KERNEL_TYPED(float)
@@ -129,6 +131,19 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons
const Tensor* q_norm_weight = context->Input(14);
const Tensor* k_norm_weight = context->Input(15);
const Tensor* head_sink = context->Input(16);
+ const Tensor* committed_length =
+ context->InputCount() > 17 ? context->Input(17) : nullptr;
+ const Tensor* key_scale = context->InputCount() > 18 ? context->Input(18) : nullptr;
+ const Tensor* value_scale = context->InputCount() > 19 ? context->Input(19) : nullptr;
+ const bool quantized_cache = past_key != nullptr && past_key->IsDataType();
+ ORT_RETURN_IF_NOT((key_scale != nullptr) == (value_scale != nullptr),
+ "DynamicSparseAttention: key_scale and value_scale must be provided together.");
+ ORT_RETURN_IF_NOT(quantized_cache == (key_scale != nullptr),
+ "DynamicSparseAttention: INT8 caches require per-channel key_scale and value_scale.");
+ ORT_RETURN_IF_NOT((past_key == nullptr || past_key->IsDataType() || quantized_cache) &&
+ (past_value == nullptr || (quantized_cache ? past_value->IsDataType()
+ : past_value->IsDataType())),
+ "DynamicSparseAttention: cache types must match activations or both be INT8.");
DynamicSparseAttentionParameters parameters;
ORT_RETURN_IF_ERROR(dynamic_sparse_attention_helper::CheckInputs(
@@ -140,9 +155,19 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons
qk_norm_epsilon_, parameters));
parameters.rotary_interleaved = rotary_interleaved_;
parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr;
+ if (quantized_cache) {
+ ORT_RETURN_IF_NOT(past_value != nullptr &&
+ key_scale->Shape() == TensorShape({parameters.kv_num_heads * parameters.head_size}) &&
+ value_scale->Shape() == key_scale->Shape(),
+ "DynamicSparseAttention: INT8 cache scales must have shape [kv_num_heads * head_size].");
+ }
ORT_RETURN_IF_NOT(parameters.head_size <= GetDeviceProp().maxThreadsPerBlock,
"DynamicSparseAttention: head_size exceeds the CUDA device thread-block limit.");
+ ORT_RETURN_IF(committed_length != nullptr &&
+ committed_length->Shape().Size() != 1 &&
+ committed_length->Shape().Size() != parameters.batch_size,
+ "DynamicSparseAttention: committed_length must have 1 or batch_size int32 elements.");
TensorShape output_shape({parameters.batch_size, parameters.sequence_length,
parameters.query_hidden_size});
@@ -153,6 +178,8 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons
parameters.cache_capacity, parameters.head_size});
Tensor* present_key_output = context->OutputCount() > 1 ? context->Output(1, present_shape) : nullptr;
Tensor* present_value_output = context->OutputCount() > 2 ? context->Output(2, present_shape) : nullptr;
+ Tensor* validation_status_output =
+ context->OutputCount() > 3 ? context->Output(3, TensorShape({1})) : nullptr;
const size_t query_elements = SafeInt(parameters.batch_size) *
parameters.sequence_length * parameters.query_hidden_size;
auto prepared_query = GetScratchBuffer(query_elements, GetComputeStream(context));
@@ -160,16 +187,17 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons
const size_t cache_elements = SafeInt(parameters.batch_size) *
parameters.kv_num_heads * parameters.cache_capacity *
parameters.head_size;
- IAllocatorUniquePtr key_cache_scratch;
- IAllocatorUniquePtr value_cache_scratch;
- T* key_cache = present_key_output == nullptr ? nullptr : present_key_output->MutableData();
- T* value_cache = present_value_output == nullptr ? nullptr : present_value_output->MutableData();
+ const size_t cache_bytes = SafeInt(cache_elements) * (quantized_cache ? 1 : sizeof(T));
+ IAllocatorUniquePtr key_cache_scratch;
+ IAllocatorUniquePtr value_cache_scratch;
+ void* key_cache = present_key_output == nullptr ? nullptr : present_key_output->MutableDataRaw();
+ void* value_cache = present_value_output == nullptr ? nullptr : present_value_output->MutableDataRaw();
if (key_cache == nullptr) {
- key_cache_scratch = GetScratchBuffer(cache_elements, GetComputeStream(context));
+ key_cache_scratch = GetScratchBuffer(cache_bytes, GetComputeStream(context));
key_cache = key_cache_scratch.get();
}
if (value_cache == nullptr) {
- value_cache_scratch = GetScratchBuffer(cache_elements, GetComputeStream(context));
+ value_cache_scratch = GetScratchBuffer(cache_bytes, GetComputeStream(context));
value_cache = value_cache_scratch.get();
}
@@ -178,8 +206,11 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons
data.query = reinterpret_cast(query->Data());
data.key = key == nullptr ? nullptr : reinterpret_cast(key->Data());
data.value = value == nullptr ? nullptr : reinterpret_cast(value->Data());
- data.past_key = past_key == nullptr ? nullptr : reinterpret_cast(past_key->Data());
- data.past_value = past_value == nullptr ? nullptr : reinterpret_cast(past_value->Data());
+ data.past_key = past_key == nullptr ? nullptr : reinterpret_cast(past_key->DataRaw());
+ data.past_value = past_value == nullptr ? nullptr : reinterpret_cast(past_value->DataRaw());
+ data.quantized_cache = quantized_cache;
+ data.key_scale = key_scale == nullptr ? nullptr : key_scale->Data();
+ data.value_scale = value_scale == nullptr ? nullptr : value_scale->Data();
data.auxiliary_key =
auxiliary_key == nullptr ? nullptr : reinterpret_cast(auxiliary_key->Data());
data.auxiliary_value = auxiliary_value == nullptr
@@ -207,22 +238,38 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons
CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(stream, &capture_status));
ORT_RETURN_IF_NOT(capture_status == cudaStreamCaptureStatusNone,
"DynamicSparseAttention strict validation is not supported during CUDA graph capture.");
+ }
- auto validation_error = GetScratchBuffer(1, GetComputeStream(context));
+ IAllocatorUniquePtr validation_error_scratch;
+ int32_t* validation_error = nullptr;
+ if (validation_status_output != nullptr) {
+ validation_error = validation_status_output->MutableData();
+ } else if (strict_validation_) {
+ validation_error_scratch = GetScratchBuffer(1, GetComputeStream(context));
+ validation_error = validation_error_scratch.get();
+ }
+ if (validation_error != nullptr) {
const size_t validation_workspace_elements =
GetDynamicSparseAttentionValidationWorkspaceSize(
parameters, GetDeviceProp().sharedMemPerBlock);
auto validation_workspace =
GetScratchBuffer(validation_workspace_elements, GetComputeStream(context));
+ const int committed_length_stride =
+ (committed_length != nullptr && committed_length->Shape().Size() == parameters.batch_size &&
+ parameters.batch_size > 1)
+ ? 1
+ : 0;
ORT_RETURN_IF_ERROR(ValidateDynamicSparseAttentionOnDevice(
stream, data.selected_indices, data.selected_counts, data.seqlens_k,
- data.position_ids, parameters, validation_error.get(), validation_workspace.get(),
- GetDeviceProp().sharedMemPerBlock, true));
+ data.position_ids, parameters, validation_error, validation_workspace.get(),
+ GetDeviceProp().sharedMemPerBlock, strict_validation_, data.key_scale, data.value_scale,
+ committed_length == nullptr ? nullptr : committed_length->Data(),
+ committed_length_stride));
}
const size_t attention_workspace_elements =
GetDynamicSparseAttentionWorkspaceSize(
- parameters, sizeof(CudaT), GetDeviceProp().sharedMemPerBlock);
+ parameters, sizeof(CudaT), GetDeviceProp().sharedMemPerBlock, quantized_cache);
auto attention_workspace =
GetScratchBuffer(attention_workspace_elements, GetComputeStream(context));
data.attention_workspace = attention_workspace.get();
diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu
index 25da163a21f97..a612c7071aab9 100644
--- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu
@@ -6,9 +6,11 @@
#include
#include
#include
+#include
#include
#include
+#include
#include
#include
@@ -16,6 +18,8 @@
#include "core/common/safeint.h"
#include "core/providers/cuda/cuda_common.h"
+// Include env_var_utils.h after cuda_common.h: the latter transitively pulls in provider_api.h,
+// which defines SHARED_PROVIDER and suppresses env_var_utils.h's own logging.h include.
#include "core/platform/env_var_utils.h"
namespace onnxruntime {
@@ -28,6 +32,13 @@ constexpr int kSplitCandidateSize = 256;
constexpr int kTargetSplitBlocks = 1024;
constexpr size_t kMaxFusedSharedBytes = 16 * 1024;
constexpr size_t kMaxValidationSharedBytes = 32 * 1024;
+constexpr int kGroupedPrefillCandidatesPerTile = 8;
+constexpr int kGroupedPrefillQueryHeadsPerBlock = 6;
+constexpr int kGroupedDecodeQueryHeadsPerBlock = 8;
+constexpr int kGroupedDecodeCandidatesPerSplit = 64;
+constexpr int kTargetGroupedDecodeBlocks = 128;
+constexpr int kGroupedQuantizedDecodeCandidatesPerSplit = 32;
+constexpr int kTargetGroupedQuantizedDecodeBlocks = 512;
template
__device__ __forceinline__ float ToFloat(T value) {
@@ -39,9 +50,54 @@ __device__ __forceinline__ T FromFloat(float value) {
return static_cast(value);
}
+__device__ __forceinline__ int8_t QuantizeCacheValue(float value, float scale) {
+ return static_cast(__float2int_rn(fminf(127.0f, fmaxf(-128.0f, value / scale))));
+}
+
+__device__ __forceinline__ uint32_t PackMmaHalf2(half low, half high) {
+ return static_cast(__half_as_ushort(low)) |
+ (static_cast(__half_as_ushort(high)) << 16);
+}
+
+__device__ __forceinline__ void MmaM16N8K16(float (&output)[4],
+ const uint32_t (&query_fragment)[4],
+ const uint32_t (&key_fragment)[2]) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
+ : "+f"(output[0]), "+f"(output[1]), "+f"(output[2]), "+f"(output[3])
+ : "r"(query_fragment[0]), "r"(query_fragment[1]),
+ "r"(query_fragment[2]), "r"(query_fragment[3]),
+ "r"(key_fragment[0]), "r"(key_fragment[1]));
+#else
+ ORT_UNUSED_PARAMETER(output);
+ ORT_UNUSED_PARAMETER(query_fragment);
+ ORT_UNUSED_PARAMETER(key_fragment);
+#endif
+}
+
+// Copies one head_size == 256 K/V row into shared memory with one 16-byte chunk per lane. Tensors bound
+// through IOBinding (cache outputs, auxiliary inputs) can start at any element offset, so a source that is
+// not 16-byte aligned is copied element-wise instead of faulting with cudaErrorMisalignedAddress.
+template
+__device__ __forceinline__ void StageHeadRow(T* shared_row, const T* global_row, int lane, int head_size) {
+ if ((reinterpret_cast(global_row) & 15) == 0) {
+ reinterpret_cast(shared_row)[lane] = reinterpret_cast(global_row)[lane];
+ } else {
+ for (int d = lane; d < head_size; d += 32) {
+ shared_row[d] = global_row[d];
+ }
+ }
+}
+
__device__ __forceinline__ void SetValidationError(int32_t* error_flag,
DynamicSparseAttentionValidationError error) {
- atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, static_cast(error));
+ const int32_t previous =
+ atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, static_cast(error));
+ if (previous == kDynamicSparseAttentionInvalidCommittedLength) {
+ atomicCAS(error_flag, kDynamicSparseAttentionInvalidCommittedLength, static_cast(error));
+ }
}
__global__ void ValidateInputsKernel(const int32_t* selected_indices,
@@ -57,21 +113,34 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices,
int rotary_max_position,
bool do_rotary,
bool use_auxiliary,
- uint32_t* validation_bitmap,
- size_t validation_bitmap_words,
- bool use_shared_bitmap,
- int32_t* error_flag) {
+ uint32_t* validation_hash_table,
+ size_t validation_hash_entries,
+ bool use_shared_hash_table,
+ int32_t* error_flag,
+ const float* key_scale,
+ const float* value_scale,
+ int cache_channels,
+ const int32_t* committed_length,
+ int committed_length_stride) {
const int row = static_cast(blockIdx.x);
const int b = row / sequence_length;
const int s = row - b * sequence_length;
if (b >= batch_size) {
return;
}
+ if (row == 0 && key_scale != nullptr) {
+ for (int i = threadIdx.x; i < cache_channels; i += blockDim.x) {
+ if (!(key_scale[i] > 0.0f) || !isfinite(key_scale[i]) ||
+ !(value_scale[i] > 0.0f) || !isfinite(value_scale[i])) {
+ SetValidationError(error_flag, kDynamicSparseAttentionInvalidCacheScale);
+ }
+ }
+ }
- extern __shared__ uint32_t shared_bitmap[];
- if (use_shared_bitmap) {
- for (size_t i = threadIdx.x; i < validation_bitmap_words; i += blockDim.x) {
- shared_bitmap[i] = 0;
+ extern __shared__ uint32_t shared_hash_table[];
+ if (use_shared_hash_table) {
+ for (size_t i = threadIdx.x; i < validation_hash_entries; i += blockDim.x) {
+ shared_hash_table[i] = std::numeric_limits::max();
}
}
__syncthreads();
@@ -120,10 +189,10 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices,
const int32_t* row_indices =
max_selected == 0 ? selected_indices : selected_indices + static_cast(row) * max_selected;
- uint32_t* row_bitmap =
- use_shared_bitmap
- ? shared_bitmap
- : validation_bitmap + static_cast(row) * validation_bitmap_words;
+ uint32_t* row_hash_table =
+ use_shared_hash_table
+ ? shared_hash_table
+ : validation_hash_table + static_cast(row) * validation_hash_entries;
for (int i = static_cast(threadIdx.x); i < max_selected; i += static_cast(blockDim.x)) {
const int index = row_indices[i];
DynamicSparseAttentionValidationError error = kDynamicSparseAttentionValidationOk;
@@ -138,10 +207,19 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices,
} else if (!use_auxiliary && index > query_position) {
error = kDynamicSparseAttentionNonCausalIndex;
} else {
- const uint32_t mask = 1U << (index & 31);
- const uint32_t previous = atomicOr(row_bitmap + (index >> 5), mask);
- if ((previous & mask) != 0) {
- error = kDynamicSparseAttentionDuplicateIndex;
+ const uint32_t hash_mask = static_cast(validation_hash_entries - 1);
+ uint32_t slot = static_cast(index) * 0x9E3779B9U & hash_mask;
+ for (size_t probe = 0; probe < validation_hash_entries; ++probe) {
+ const uint32_t previous = atomicCAS(
+ row_hash_table + slot, std::numeric_limits::max(), static_cast(index));
+ if (previous == std::numeric_limits::max()) {
+ break;
+ }
+ if (previous == static_cast(index)) {
+ error = kDynamicSparseAttentionDuplicateIndex;
+ break;
+ }
+ slot = (slot + 1) & hash_mask;
}
}
if (error != kDynamicSparseAttentionValidationOk) {
@@ -169,6 +247,15 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices,
}
return;
}
+
+ if (s == 0 && threadIdx.x == 0 && committed_length != nullptr) {
+ const int64_t committed = committed_length[b * committed_length_stride];
+ const int64_t total = committed + sequence_length;
+ if (committed < 0 || total > main_capacity || seqlens_k[b] != total - 1) {
+ atomicCAS(error_flag, kDynamicSparseAttentionValidationOk,
+ kDynamicSparseAttentionInvalidCommittedLength);
+ }
+ }
}
template
@@ -305,7 +392,10 @@ __global__ void AppendKvKernel(const T* query,
int rotary_offset,
bool rotary_interleaved,
bool is_packed,
- float epsilon) {
+ float epsilon,
+ bool quantized_cache,
+ const float* key_scale,
+ const float* value_scale) {
const int block = static_cast(blockIdx.x);
if (block >= block_count) {
return;
@@ -371,8 +461,16 @@ __global__ void AppendKvKernel(const T* query,
destination_sequence) *
head_size +
h;
- present_key[cache_offset] = result;
- present_value[cache_offset] = (is_packed ? query : value)[value_offset];
+ if (quantized_cache) {
+ const int channel = kv_head * head_size + h;
+ reinterpret_cast(present_key)[cache_offset] =
+ QuantizeCacheValue(ToFloat(result), key_scale[channel]);
+ reinterpret_cast(present_value)[cache_offset] =
+ QuantizeCacheValue(ToFloat((is_packed ? query : value)[value_offset]), value_scale[channel]);
+ } else {
+ present_key[cache_offset] = result;
+ present_value[cache_offset] = (is_packed ? query : value)[value_offset];
+ }
}
}
@@ -400,7 +498,10 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
float scale,
bool local_plus_selected,
bool selected_from_auxiliary,
- bool use_smooth_softmax) {
+ bool use_smooth_softmax,
+ bool quantized_cache,
+ const float* key_scale,
+ const float* value_scale) {
const int block = static_cast(blockIdx.x);
if (block >= block_count) {
return;
@@ -436,8 +537,7 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
extern __shared__ float shared[];
float* logits = shared;
- float* valid_candidates = logits + candidate_capacity;
- float* reduction = valid_candidates + candidate_capacity;
+ float* reduction = logits + candidate_capacity;
T* shared_query = reinterpret_cast(reduction + blockDim.x);
if (h < head_size) {
shared_query[h] = query_head[h];
@@ -452,7 +552,7 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
index = local_start + candidate;
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- key_head = main_key + offset;
+ key_head = quantized_cache ? nullptr : main_key + offset;
} else {
index = row_indices[candidate - local_count];
if (selected_from_auxiliary) {
@@ -470,7 +570,7 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
if (valid) {
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- key_head = main_key + offset;
+ key_head = quantized_cache ? nullptr : main_key + offset;
}
}
}
@@ -478,15 +578,21 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
float dot = 0.0f;
if (valid) {
for (int d = lane; d < head_size; d += 32) {
- dot += ToFloat(shared_query[d]) * ToFloat(key_head[d]);
+ const bool quantized = quantized_cache && (candidate < local_count || !selected_from_auxiliary);
+ const int64_t offset =
+ ((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size + d;
+ const float key_value = quantized
+ ? static_cast(reinterpret_cast(main_key)[offset]) *
+ key_scale[kv_head * head_size + d]
+ : ToFloat(key_head[d]);
+ dot += ToFloat(shared_query[d]) * key_value;
}
for (int offset = 16; offset > 0; offset >>= 1) {
dot += __shfl_down_sync(0xffffffff, dot, offset);
}
}
if (lane == 0) {
- logits[candidate] = valid ? dot * scale : 0.0f;
- valid_candidates[candidate] = valid ? 1.0f : 0.0f;
+ logits[candidate] = valid ? dot * scale : -CUDART_INF_F;
}
}
__syncthreads();
@@ -495,9 +601,7 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
? (head_sink == nullptr ? 0.0f : ToFloat(head_sink[head]))
: -FLT_MAX;
for (int candidate = h; candidate < candidate_count; candidate += static_cast(blockDim.x)) {
- if (valid_candidates[candidate] != 0.0f) {
- thread_max = fmaxf(thread_max, logits[candidate]);
- }
+ thread_max = fmaxf(thread_max, logits[candidate]);
}
reduction[h] = thread_max;
__syncthreads();
@@ -517,7 +621,7 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
}
for (int candidate = h; candidate < candidate_count; candidate += static_cast(blockDim.x)) {
const float logit = logits[candidate];
- const float weight = valid_candidates[candidate] != 0.0f ? expf(logit - max_logit) : 0.0f;
+ const float weight = logit == -CUDART_INF_F ? 0.0f : expf(logit - max_logit);
logits[candidate] = weight;
thread_sum += weight;
}
@@ -536,19 +640,21 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
float accumulator = 0.0f;
for (int candidate = 0; candidate < candidate_count; ++candidate) {
const float weight = logits[candidate];
- if (valid_candidates[candidate] == 0.0f) {
- continue;
- }
-
int index;
const T* value_head;
if (candidate < local_count) {
index = local_start + candidate;
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- value_head = main_value + offset;
+ value_head = quantized_cache ? nullptr : main_value + offset;
} else {
index = row_indices[candidate - local_count];
+ if (selected_from_auxiliary
+ ? index < 0 || index >= auxiliary_sequence_length
+ : index < 0 || index >= total_length || index > query_position || index >= main_capacity ||
+ (local_plus_selected && index >= local_start)) {
+ continue;
+ }
if (selected_from_auxiliary) {
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * auxiliary_sequence_length + index) * head_size;
@@ -556,10 +662,17 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
} else {
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- value_head = main_value + offset;
+ value_head = quantized_cache ? nullptr : main_value + offset;
}
}
- accumulator += weight * ToFloat(value_head[h]);
+ const bool quantized = quantized_cache && (candidate < local_count || !selected_from_auxiliary);
+ const int64_t offset =
+ ((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size + h;
+ const float value = quantized
+ ? static_cast(reinterpret_cast(main_value)[offset]) *
+ value_scale[kv_head * head_size + h]
+ : ToFloat(value_head[h]);
+ accumulator += weight * value;
}
const int64_t output_offset =
@@ -569,6 +682,599 @@ __global__ void FusedDynamicSparseAttentionKernel(const T* query,
}
}
+template
+__global__ void GroupedPrefillDynamicSparseAttentionKernel(
+ const T* query,
+ const T* main_key,
+ const T* main_value,
+ const T* auxiliary_key,
+ const T* auxiliary_value,
+ const int32_t* selected_indices,
+ const int32_t* selected_counts,
+ const int32_t* seqlens_k,
+ const T* head_sink,
+ T* output,
+ int block_count,
+ int sequence_length,
+ int num_heads,
+ int kv_num_heads,
+ int head_size,
+ int main_capacity,
+ int auxiliary_sequence_length,
+ int max_selected,
+ int local_window_size,
+ float scale,
+ bool local_plus_selected,
+ bool selected_from_auxiliary,
+ bool use_smooth_softmax) {
+ const int block = static_cast(blockIdx.x);
+ if (block >= block_count) {
+ return;
+ }
+
+ const int query_heads_per_kv = num_heads / kv_num_heads;
+ const int head_tiles = (query_heads_per_kv + kQueryHeadsPerBlock - 1) / kQueryHeadsPerBlock;
+ const int row = block / (kv_num_heads * head_tiles);
+ const int kv_tile = block - row * kv_num_heads * head_tiles;
+ const int kv_head = kv_tile / head_tiles;
+ const int head_tile = kv_tile - kv_head * head_tiles;
+ const int group_head = head_tile * kQueryHeadsPerBlock + static_cast(threadIdx.x) / 32;
+ const bool active_head = group_head < query_heads_per_kv;
+ const int head = kv_head * query_heads_per_kv + group_head;
+ const int lane = static_cast(threadIdx.x) & 31;
+ const int b = row / sequence_length;
+ const int s = row - b * sequence_length;
+ const int64_t total_length_64 = static_cast(seqlens_k[b]) + 1;
+ const bool valid_length = total_length_64 >= sequence_length && total_length_64 <= main_capacity;
+ const int total_length = valid_length ? static_cast(total_length_64) : 0;
+ const int query_position = valid_length ? total_length - sequence_length + s : -1;
+
+ int local_start = 0;
+ int local_count = 0;
+ if (local_plus_selected) {
+ local_start = max(0, query_position - local_window_size + 1);
+ const int local_end = min(query_position, min(total_length, main_capacity) - 1);
+ local_count = max(0, local_end - local_start + 1);
+ }
+ const int selected_count = max(0, min(selected_counts[row], max_selected));
+ const int candidate_count = local_count + selected_count;
+ const int32_t* row_indices =
+ max_selected == 0 ? selected_indices : selected_indices + static_cast(row) * max_selected;
+
+ extern __shared__ unsigned char shared_bytes[];
+ T* shared_key = reinterpret_cast(shared_bytes);
+ T* shared_value = shared_key + kGroupedPrefillCandidatesPerTile * head_size;
+ __shared__ bool candidate_valid[kGroupedPrefillCandidatesPerTile];
+ __shared__ half mma_query[kGroupedPrefillQueryHeadsPerBlock * 256];
+ __shared__ float mma_scores[kGroupedPrefillQueryHeadsPerBlock * kGroupedPrefillCandidatesPerTile];
+
+ constexpr int kValuesPerLane = 8;
+ float query_values[kValuesPerLane]{};
+ float accumulator[kValuesPerLane]{};
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ if constexpr (std::is_same_v) {
+ const int group_head_base = head_tile * kQueryHeadsPerBlock;
+ for (int index = static_cast(threadIdx.x);
+ index < kGroupedPrefillQueryHeadsPerBlock * head_size;
+ index += static_cast(blockDim.x)) {
+ const int local_head = index / head_size;
+ const int head_offset = index - local_head * head_size;
+ const int staged_group_head = group_head_base + local_head;
+ mma_query[index] = staged_group_head < query_heads_per_kv
+ ? query[(static_cast(row) * num_heads +
+ kv_head * query_heads_per_kv + staged_group_head) *
+ head_size +
+ head_offset]
+ : half{0.0f};
+ }
+ __syncthreads();
+ } else
+#endif
+ if (active_head) {
+ const T* query_head = query + (static_cast(row) * num_heads + head) * head_size;
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ query_values[i] = ToFloat(query_head[lane + i * 32]);
+ }
+ }
+
+ float running_max = -FLT_MAX;
+ float running_sum = 0.0f;
+ if (active_head && use_smooth_softmax) {
+ running_max = head_sink == nullptr ? 0.0f : ToFloat(head_sink[head]);
+ running_sum = 1.0f;
+ }
+
+ for (int candidate_start = 0; candidate_start < candidate_count;
+ candidate_start += kGroupedPrefillCandidatesPerTile) {
+ const int tile_count = min(kGroupedPrefillCandidatesPerTile, candidate_count - candidate_start);
+ const int load_warp = static_cast(threadIdx.x) / 32;
+ const int load_lane = static_cast(threadIdx.x) & 31;
+ for (int load_candidate = load_warp; load_candidate < tile_count;
+ load_candidate += kQueryHeadsPerBlock) {
+ int candidate_index = 0;
+ int candidate_is_auxiliary = 0;
+ int is_valid = 0;
+ if (load_lane == 0) {
+ const int candidate = candidate_start + load_candidate;
+ is_valid = valid_length;
+ if (candidate < local_count) {
+ candidate_index = local_start + candidate;
+ } else {
+ candidate_index = row_indices[candidate - local_count];
+ candidate_is_auxiliary = selected_from_auxiliary;
+ if (selected_from_auxiliary) {
+ is_valid = is_valid && candidate_index >= 0 &&
+ candidate_index < auxiliary_sequence_length;
+ } else {
+ is_valid = is_valid && candidate_index >= 0 && candidate_index < total_length &&
+ candidate_index <= query_position && candidate_index < main_capacity;
+ if (local_plus_selected && candidate_index >= local_start &&
+ candidate_index <= query_position) {
+ is_valid = false;
+ }
+ }
+ }
+ candidate_valid[load_candidate] = is_valid;
+ }
+ candidate_index = __shfl_sync(0xffffffff, candidate_index, 0);
+ candidate_is_auxiliary = __shfl_sync(0xffffffff, candidate_is_auxiliary, 0);
+ is_valid = __shfl_sync(0xffffffff, is_valid, 0);
+
+ const int shared_offset = load_candidate * head_size;
+ if (is_valid) {
+ const int64_t source_length = candidate_is_auxiliary
+ ? auxiliary_sequence_length
+ : main_capacity;
+ const int64_t offset =
+ ((static_cast(b) * kv_num_heads + kv_head) * source_length +
+ candidate_index) *
+ head_size;
+ const T* key_head = candidate_is_auxiliary
+ ? auxiliary_key + offset
+ : main_key + offset;
+ const T* value_head = candidate_is_auxiliary
+ ? auxiliary_value + offset
+ : main_value + offset;
+ StageHeadRow(shared_key + shared_offset, key_head, load_lane, head_size);
+ StageHeadRow(shared_value + shared_offset, value_head, load_lane, head_size);
+ }
+ }
+ __syncthreads();
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ if constexpr (std::is_same_v) {
+ if (threadIdx.x < 32) {
+ const int mma_lane = static_cast(threadIdx.x);
+ const int row_pair = mma_lane >> 2;
+ const int k_pair = (mma_lane & 3) * 2;
+ float score_fragment[4]{};
+ for (int k_start = 0; k_start < head_size; k_start += 16) {
+ const int first_k = k_start + k_pair;
+ const int second_k = first_k + 8;
+ uint32_t query_fragment[4]{};
+ if (row_pair < kGroupedPrefillQueryHeadsPerBlock) {
+ query_fragment[0] = *reinterpret_cast(
+ mma_query + row_pair * head_size + first_k);
+ query_fragment[2] = *reinterpret_cast(
+ mma_query + row_pair * head_size + second_k);
+ }
+ const uint32_t key_fragment[2]{
+ *reinterpret_cast(
+ shared_key + row_pair * head_size + first_k),
+ *reinterpret_cast(
+ shared_key + row_pair * head_size + second_k)};
+ MmaM16N8K16(score_fragment, query_fragment, key_fragment);
+ }
+ const int score_column = (mma_lane & 3) * 2;
+ if (row_pair < kGroupedPrefillQueryHeadsPerBlock) {
+ mma_scores[row_pair * kGroupedPrefillCandidatesPerTile + score_column] = score_fragment[0];
+ mma_scores[row_pair * kGroupedPrefillCandidatesPerTile + score_column + 1] = score_fragment[1];
+ }
+ }
+ __syncthreads();
+ }
+#endif
+
+ for (int tile_candidate = 0; tile_candidate < tile_count; ++tile_candidate) {
+ if (active_head && candidate_valid[tile_candidate]) {
+ const int shared_offset = tile_candidate * head_size;
+ float score;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ if constexpr (std::is_same_v) {
+ score = mma_scores[static_cast(threadIdx.x) / 32 *
+ kGroupedPrefillCandidatesPerTile +
+ tile_candidate] *
+ scale;
+ } else
+#endif
+ {
+ float dot = 0.0f;
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ dot += query_values[i] * ToFloat(shared_key[shared_offset + lane + i * 32]);
+ }
+ for (int offset = 16; offset > 0; offset >>= 1) {
+ dot += __shfl_down_sync(0xffffffff, dot, offset);
+ }
+ score = __shfl_sync(0xffffffff, dot, 0) * scale;
+ }
+ const float next_max = fmaxf(running_max, score);
+ const float old_scale = running_sum == 0.0f ? 0.0f : expf(running_max - next_max);
+ const float weight = expf(score - next_max);
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ accumulator[i] = accumulator[i] * old_scale +
+ weight * ToFloat(shared_value[shared_offset + lane + i * 32]);
+ }
+ running_max = next_max;
+ running_sum = running_sum * old_scale + weight;
+ }
+ }
+ __syncthreads();
+ }
+
+ if (active_head) {
+ T* output_head = output + (static_cast(row) * num_heads + head) * head_size;
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ output_head[lane + i * 32] = running_sum == 0.0f
+ ? FromFloat(0.0f)
+ : FromFloat(accumulator[i] / running_sum);
+ }
+ }
+}
+
+__device__ __forceinline__ float4 DequantizeCacheWord(uint32_t packed, float4 scales) {
+ return make_float4(
+ static_cast(static_cast(packed << 24) >> 24) * scales.x,
+ static_cast(static_cast(packed << 16) >> 24) * scales.y,
+ static_cast(static_cast(packed << 8) >> 24) * scales.z,
+ static_cast(static_cast(packed) >> 24) * scales.w);
+}
+
+template
+__global__ void GroupedDecodeDynamicSparseAttentionKernel(
+ const T* query,
+ const T* main_key,
+ const T* main_value,
+ const T* auxiliary_key,
+ const T* auxiliary_value,
+ const int32_t* selected_indices,
+ const int32_t* selected_counts,
+ const int32_t* seqlens_k,
+ const T* head_sink,
+ float* partial_max,
+ float* partial_sum,
+ float* partial_output,
+ int grouped_block_count,
+ int split_count,
+ int sequence_length,
+ int num_heads,
+ int kv_num_heads,
+ int head_size,
+ int main_capacity,
+ int auxiliary_sequence_length,
+ int max_selected,
+ int local_window_size,
+ float scale,
+ bool local_plus_selected,
+ bool selected_from_auxiliary,
+ bool use_smooth_softmax,
+ const float* key_scale,
+ const float* value_scale) {
+ const int block = static_cast(blockIdx.x);
+ if (block >= grouped_block_count) {
+ return;
+ }
+
+ const int split = block % split_count;
+ const int grouped_block = block / split_count;
+ const int query_heads_per_kv = num_heads / kv_num_heads;
+ const int head_tiles = (query_heads_per_kv + kQueryHeadsPerBlock - 1) / kQueryHeadsPerBlock;
+ const int row = grouped_block / (kv_num_heads * head_tiles);
+ const int b = row / sequence_length;
+ const int s = row - b * sequence_length;
+ const int kv_tile = grouped_block - row * kv_num_heads * head_tiles;
+ const int kv_head = kv_tile / head_tiles;
+ const int head_tile = kv_tile - kv_head * head_tiles;
+ const int group_head = head_tile * kQueryHeadsPerBlock + static_cast(threadIdx.x) / 32;
+ const bool active_head = group_head < query_heads_per_kv;
+ const int head = kv_head * query_heads_per_kv + group_head;
+ const int lane = static_cast(threadIdx.x) & 31;
+ const int64_t total_length_64 = static_cast(seqlens_k[b]) + 1;
+ const bool valid_length = total_length_64 >= sequence_length && total_length_64 <= main_capacity;
+ const int total_length = valid_length ? static_cast(total_length_64) : 0;
+ const int query_position = valid_length ? total_length - sequence_length + s : -1;
+
+ int local_start = 0;
+ int local_count = 0;
+ if (local_plus_selected) {
+ local_start = max(0, query_position - local_window_size + 1);
+ const int local_end = min(query_position, min(total_length, main_capacity) - 1);
+ local_count = max(0, local_end - local_start + 1);
+ }
+ const int selected_count = max(0, min(selected_counts[row], max_selected));
+ const int candidate_count = local_count + selected_count;
+ constexpr int kCandidatesPerSplit = kQuantizedCache ? kGroupedQuantizedDecodeCandidatesPerSplit
+ : kGroupedDecodeCandidatesPerSplit;
+ const int candidate_begin = split * kCandidatesPerSplit;
+ const int32_t* row_indices =
+ max_selected == 0 ? selected_indices : selected_indices + static_cast(row) * max_selected;
+
+ extern __shared__ unsigned char shared_bytes[];
+ using SharedT = std::conditional_t;
+ SharedT* shared_key = reinterpret_cast(shared_bytes);
+ SharedT* shared_value = shared_key + kGroupedPrefillCandidatesPerTile * head_size;
+ __shared__ int candidate_indices[kGroupedPrefillCandidatesPerTile];
+ __shared__ bool candidate_valid[kGroupedPrefillCandidatesPerTile];
+ __shared__ bool candidate_sources[kGroupedPrefillCandidatesPerTile];
+ constexpr int kMmaHeads = (std::is_same_v && !kQuantizedCache) ? kQueryHeadsPerBlock : 1;
+ __shared__ half mma_query[kMmaHeads * 256];
+ __shared__ float mma_scores[kMmaHeads * kGroupedPrefillCandidatesPerTile];
+
+ constexpr int kValuesPerLane = 8;
+ float query_values[kValuesPerLane]{};
+ float accumulator[kValuesPerLane]{};
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ if constexpr (std::is_same_v && !kQuantizedCache) {
+ const int group_head_base = head_tile * kQueryHeadsPerBlock;
+ for (int index = static_cast(threadIdx.x);
+ index < kQueryHeadsPerBlock * head_size;
+ index += static_cast(blockDim.x)) {
+ const int local_head = index / head_size;
+ const int head_offset = index - local_head * head_size;
+ const int staged_group_head = group_head_base + local_head;
+ mma_query[index] = staged_group_head < query_heads_per_kv
+ ? query[(static_cast(row) * num_heads +
+ kv_head * query_heads_per_kv + staged_group_head) *
+ head_size +
+ head_offset]
+ : half{0.0f};
+ }
+ __syncthreads();
+ } else
+#endif
+ if (active_head) {
+ const T* query_head = query + (static_cast(row) * num_heads + head) * head_size;
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ query_values[i] = ToFloat(query_head[lane + i * 32]);
+ }
+ }
+
+ float running_max = -FLT_MAX;
+ float running_sum = 0.0f;
+ if (active_head && split == 0 && use_smooth_softmax) {
+ running_max = head_sink == nullptr ? 0.0f : ToFloat(head_sink[head]);
+ running_sum = 1.0f;
+ }
+
+ for (int candidate_start = candidate_begin; candidate_start < candidate_count;) {
+ const int tile_count = min(kGroupedPrefillCandidatesPerTile, candidate_count - candidate_start);
+ const int load_candidate = static_cast(threadIdx.x) / 32;
+ const int load_lane = static_cast(threadIdx.x) & 31;
+ int candidate_index = 0;
+ int candidate_is_auxiliary = 0;
+ int is_valid = 0;
+ if constexpr (!kWarpLocalMetadata) {
+ if (threadIdx.x < tile_count) {
+ const int tile_candidate = static_cast(threadIdx.x);
+ const int candidate = candidate_start + tile_candidate;
+ candidate_valid[tile_candidate] = total_length > 0 && total_length <= main_capacity;
+ candidate_sources[tile_candidate] = false;
+ if (candidate < local_count) {
+ candidate_indices[tile_candidate] = local_start + candidate;
+ } else {
+ candidate_indices[tile_candidate] = row_indices[candidate - local_count];
+ candidate_sources[tile_candidate] = selected_from_auxiliary;
+ if (selected_from_auxiliary) {
+ candidate_valid[tile_candidate] =
+ candidate_valid[tile_candidate] && candidate_indices[tile_candidate] >= 0 &&
+ candidate_indices[tile_candidate] < auxiliary_sequence_length;
+ } else {
+ candidate_valid[tile_candidate] =
+ candidate_valid[tile_candidate] && candidate_indices[tile_candidate] >= 0 &&
+ candidate_indices[tile_candidate] < total_length &&
+ candidate_indices[tile_candidate] <= query_position;
+ if (local_plus_selected && candidate_indices[tile_candidate] >= local_start) {
+ candidate_valid[tile_candidate] = false;
+ }
+ }
+ }
+ }
+ __syncthreads();
+ }
+ if (load_candidate < tile_count) {
+ if constexpr (kWarpLocalMetadata) {
+ if (load_lane == 0) {
+ const int candidate = candidate_start + load_candidate;
+ is_valid = total_length > 0 && total_length <= main_capacity;
+ if (candidate < local_count) {
+ candidate_index = local_start + candidate;
+ } else {
+ candidate_index = row_indices[candidate - local_count];
+ candidate_is_auxiliary = selected_from_auxiliary;
+ if (selected_from_auxiliary) {
+ is_valid = is_valid && candidate_index >= 0 &&
+ candidate_index < auxiliary_sequence_length;
+ } else {
+ is_valid = is_valid && candidate_index >= 0 && candidate_index < total_length &&
+ candidate_index <= query_position;
+ if (local_plus_selected && candidate_index >= local_start) {
+ is_valid = false;
+ }
+ }
+ }
+ candidate_valid[load_candidate] = is_valid;
+ }
+ candidate_index = __shfl_sync(0xffffffff, candidate_index, 0);
+ candidate_is_auxiliary = __shfl_sync(0xffffffff, candidate_is_auxiliary, 0);
+ is_valid = __shfl_sync(0xffffffff, is_valid, 0);
+ } else {
+ candidate_index = candidate_indices[load_candidate];
+ candidate_is_auxiliary = candidate_sources[load_candidate];
+ is_valid = candidate_valid[load_candidate];
+ }
+
+ const int shared_offset = load_candidate * head_size;
+ if (is_valid) {
+ const int64_t source_length = candidate_is_auxiliary
+ ? auxiliary_sequence_length
+ : main_capacity;
+ const int64_t offset =
+ ((static_cast(b) * kv_num_heads + kv_head) * source_length +
+ candidate_index) *
+ head_size;
+ if constexpr (kQuantizedCache && kVectorizedStaging) {
+ const int d = load_lane * 8;
+ float4* staged_key = reinterpret_cast(shared_key + shared_offset + d);
+ float4* staged_value = reinterpret_cast(shared_value + shared_offset + d);
+ if (candidate_is_auxiliary) {
+ const uint4 packed_key = __ldg(reinterpret_cast(auxiliary_key + offset + d));
+ const uint4 packed_value = __ldg(reinterpret_cast(auxiliary_value + offset + d));
+ const T* key = reinterpret_cast(&packed_key);
+ const T* value = reinterpret_cast(&packed_value);
+ staged_key[0] = make_float4(ToFloat(key[0]), ToFloat(key[1]), ToFloat(key[2]), ToFloat(key[3]));
+ staged_key[1] = make_float4(ToFloat(key[4]), ToFloat(key[5]), ToFloat(key[6]), ToFloat(key[7]));
+ staged_value[0] = make_float4(ToFloat(value[0]), ToFloat(value[1]), ToFloat(value[2]), ToFloat(value[3]));
+ staged_value[1] = make_float4(ToFloat(value[4]), ToFloat(value[5]), ToFloat(value[6]), ToFloat(value[7]));
+ } else {
+ const uint2 packed_key = __ldg(reinterpret_cast(
+ reinterpret_cast(main_key) + offset + d));
+ const uint2 packed_value = __ldg(reinterpret_cast(
+ reinterpret_cast(main_value) + offset + d));
+ const float4* key_scales = reinterpret_cast(key_scale + kv_head * head_size + d);
+ const float4* value_scales = reinterpret_cast(value_scale + kv_head * head_size + d);
+ staged_key[0] = DequantizeCacheWord(packed_key.x, __ldg(key_scales));
+ staged_key[1] = DequantizeCacheWord(packed_key.y, __ldg(key_scales + 1));
+ staged_value[0] = DequantizeCacheWord(packed_value.x, __ldg(value_scales));
+ staged_value[1] = DequantizeCacheWord(packed_value.y, __ldg(value_scales + 1));
+ }
+ } else if constexpr (kQuantizedCache) {
+ for (int d = load_lane; d < head_size; d += 32) {
+ const int channel = kv_head * head_size + d;
+ shared_key[shared_offset + d] =
+ candidate_is_auxiliary
+ ? ToFloat(auxiliary_key[offset + d])
+ : static_cast(reinterpret_cast(main_key)[offset + d]) *
+ key_scale[channel];
+ shared_value[shared_offset + d] =
+ candidate_is_auxiliary
+ ? ToFloat(auxiliary_value[offset + d])
+ : static_cast(reinterpret_cast(main_value)[offset + d]) *
+ value_scale[channel];
+ }
+ } else {
+ const T* key_head = candidate_is_auxiliary ? auxiliary_key + offset : main_key + offset;
+ const T* value_head = candidate_is_auxiliary ? auxiliary_value + offset : main_value + offset;
+ StageHeadRow(shared_key + shared_offset, key_head, load_lane, head_size);
+ StageHeadRow(shared_value + shared_offset, value_head, load_lane, head_size);
+ }
+ }
+ }
+ __syncthreads();
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ if constexpr (std::is_same_v && !kQuantizedCache) {
+ if (threadIdx.x < 32) {
+ const int mma_lane = static_cast(threadIdx.x);
+ const int row_pair = mma_lane >> 2;
+ const int second_row = row_pair + 8;
+ const int k_pair = (mma_lane & 3) * 2;
+ float score_fragment[4]{};
+ for (int k_start = 0; k_start < head_size; k_start += 16) {
+ const int first_k = k_start + k_pair;
+ const int second_k = first_k + 8;
+ const uint32_t query_fragment[4]{
+ *reinterpret_cast(
+ mma_query + row_pair * head_size + first_k),
+ second_row < kQueryHeadsPerBlock
+ ? *reinterpret_cast(
+ mma_query + second_row * head_size + first_k)
+ : 0u,
+ *reinterpret_cast(
+ mma_query + row_pair * head_size + second_k),
+ second_row < kQueryHeadsPerBlock
+ ? *reinterpret_cast(
+ mma_query + second_row * head_size + second_k)
+ : 0u};
+ const uint32_t key_fragment[2]{
+ *reinterpret_cast(
+ shared_key + row_pair * head_size + first_k),
+ *reinterpret_cast(
+ shared_key + row_pair * head_size + second_k)};
+ MmaM16N8K16(score_fragment, query_fragment, key_fragment);
+ }
+ const int score_column = (mma_lane & 3) * 2;
+ mma_scores[row_pair * kGroupedPrefillCandidatesPerTile + score_column] = score_fragment[0];
+ mma_scores[row_pair * kGroupedPrefillCandidatesPerTile + score_column + 1] = score_fragment[1];
+ if (second_row < kQueryHeadsPerBlock) {
+ mma_scores[second_row * kGroupedPrefillCandidatesPerTile + score_column] = score_fragment[2];
+ mma_scores[second_row * kGroupedPrefillCandidatesPerTile + score_column + 1] = score_fragment[3];
+ }
+ }
+ __syncthreads();
+ }
+#endif
+
+ for (int tile_candidate = 0; tile_candidate < tile_count; ++tile_candidate) {
+ if (active_head && candidate_valid[tile_candidate]) {
+ const int shared_offset = tile_candidate * head_size;
+ float score;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ if constexpr (std::is_same_v && !kQuantizedCache) {
+ score = mma_scores[static_cast(threadIdx.x) / 32 *
+ kGroupedPrefillCandidatesPerTile +
+ tile_candidate] *
+ scale;
+ } else
+#endif
+ {
+ float dot = 0.0f;
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ dot += query_values[i] * ToFloat(shared_key[shared_offset + lane + i * 32]);
+ }
+ for (int offset = 16; offset > 0; offset >>= 1) {
+ dot += __shfl_down_sync(0xffffffff, dot, offset);
+ }
+ score = __shfl_sync(0xffffffff, dot, 0) * scale;
+ }
+ const float next_max = fmaxf(running_max, score);
+ const float old_scale = running_sum == 0.0f ? 0.0f : expf(running_max - next_max);
+ const float weight = expf(score - next_max);
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ accumulator[i] = accumulator[i] * old_scale +
+ weight * ToFloat(shared_value[shared_offset + lane + i * 32]);
+ }
+ running_max = next_max;
+ running_sum = running_sum * old_scale + weight;
+ }
+ }
+ __syncthreads();
+
+ candidate_start += kGroupedPrefillCandidatesPerTile;
+ if (candidate_start % kCandidatesPerSplit == 0) {
+ candidate_start += (split_count - 1) * kCandidatesPerSplit;
+ }
+ }
+
+ if (active_head) {
+ const int partial_block = (row * num_heads + head) * split_count + split;
+ if (lane == 0) {
+ partial_max[partial_block] = running_max;
+ partial_sum[partial_block] = running_sum;
+ }
+#pragma unroll
+ for (int i = 0; i < kValuesPerLane; ++i) {
+ partial_output[static_cast(partial_block) * head_size + lane + i * 32] = accumulator[i];
+ }
+ }
+}
+
template
__global__ void SplitDynamicSparseAttentionKernel(const T* query,
const T* main_key,
@@ -596,7 +1302,10 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
float scale,
bool local_plus_selected,
bool selected_from_auxiliary,
- bool use_smooth_softmax) {
+ bool use_smooth_softmax,
+ bool quantized_cache,
+ const float* key_scale,
+ const float* value_scale) {
const int partial_block = static_cast(blockIdx.x);
if (partial_block >= partial_block_count) {
return;
@@ -633,8 +1342,7 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
extern __shared__ float shared[];
float* logits = shared;
- float* valid_candidates = logits + kSplitCandidateSize;
- float* reduction = valid_candidates + kSplitCandidateSize;
+ float* reduction = logits + kSplitCandidateSize;
T* shared_query = reinterpret_cast(reduction + blockDim.x);
if (h < head_size) {
shared_query[h] = query_head[h];
@@ -663,7 +1371,7 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
index = local_start + candidate;
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- key_head = main_key + offset;
+ key_head = quantized_cache ? nullptr : main_key + offset;
} else {
index = row_indices[candidate - local_count];
if (selected_from_auxiliary) {
@@ -681,7 +1389,7 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
if (valid) {
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- key_head = main_key + offset;
+ key_head = quantized_cache ? nullptr : main_key + offset;
}
}
}
@@ -689,15 +1397,21 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
float dot = 0.0f;
if (valid) {
for (int d = lane; d < head_size; d += 32) {
- dot += ToFloat(shared_query[d]) * ToFloat(key_head[d]);
+ const bool quantized = quantized_cache && (candidate < local_count || !selected_from_auxiliary);
+ const int64_t offset =
+ ((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size + d;
+ const float key_value = quantized
+ ? static_cast(reinterpret_cast(main_key)[offset]) *
+ key_scale[kv_head * head_size + d]
+ : ToFloat(key_head[d]);
+ dot += ToFloat(shared_query[d]) * key_value;
}
for (int offset = 16; offset > 0; offset >>= 1) {
dot += __shfl_down_sync(0xffffffff, dot, offset);
}
}
if (lane == 0) {
- logits[split_candidate] = valid ? dot * scale : 0.0f;
- valid_candidates[split_candidate] = valid ? 1.0f : 0.0f;
+ logits[split_candidate] = valid ? dot * scale : -CUDART_INF_F;
}
}
__syncthreads();
@@ -705,9 +1419,7 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
float thread_max = -FLT_MAX;
for (int candidate = h; candidate < split_candidate_count;
candidate += static_cast(blockDim.x)) {
- if (valid_candidates[candidate] != 0.0f) {
- thread_max = fmaxf(thread_max, logits[candidate]);
- }
+ thread_max = fmaxf(thread_max, logits[candidate]);
}
reduction[h] = thread_max;
__syncthreads();
@@ -726,7 +1438,7 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
for (int candidate = h; candidate < split_candidate_count;
candidate += static_cast(blockDim.x)) {
const float logit = logits[candidate];
- const float weight = valid_candidates[candidate] != 0.0f ? expf(logit - next_max) : 0.0f;
+ const float weight = logit == -CUDART_INF_F ? 0.0f : expf(logit - next_max);
logits[candidate] = weight;
thread_sum += weight;
}
@@ -745,10 +1457,6 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
float tile_accumulator = 0.0f;
for (int split_candidate = 0; split_candidate < split_candidate_count; ++split_candidate) {
const float weight = logits[split_candidate];
- if (valid_candidates[split_candidate] == 0.0f) {
- continue;
- }
-
const int candidate = candidate_start + split_candidate;
int index;
const T* value_head;
@@ -756,9 +1464,15 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
index = local_start + candidate;
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- value_head = main_value + offset;
+ value_head = quantized_cache ? nullptr : main_value + offset;
} else {
index = row_indices[candidate - local_count];
+ if (selected_from_auxiliary
+ ? index < 0 || index >= auxiliary_sequence_length
+ : index < 0 || index >= total_length || index > query_position || index >= main_capacity ||
+ (local_plus_selected && index >= local_start)) {
+ continue;
+ }
if (selected_from_auxiliary) {
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * auxiliary_sequence_length + index) * head_size;
@@ -766,10 +1480,17 @@ __global__ void SplitDynamicSparseAttentionKernel(const T* query,
} else {
const int64_t offset =
((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size;
- value_head = main_value + offset;
+ value_head = quantized_cache ? nullptr : main_value + offset;
}
}
- tile_accumulator += weight * ToFloat(value_head[h]);
+ const bool quantized = quantized_cache && (candidate < local_count || !selected_from_auxiliary);
+ const int64_t offset =
+ ((static_cast(b) * kv_num_heads + kv_head) * main_capacity + index) * head_size + h;
+ const float value = quantized
+ ? static_cast(reinterpret_cast(main_value)[offset]) *
+ value_scale[kv_head * head_size + h]
+ : ToFloat(value_head[h]);
+ tile_accumulator += weight * value;
}
running_accumulator =
running_accumulator * old_scale + tile_accumulator;
@@ -832,7 +1553,9 @@ __global__ void MergeDynamicSparseAttentionKernel(const float* partial_max,
float thread_sum = 0.0f;
for (int split = h; split < split_count; split += static_cast(blockDim.x)) {
const float split_max = partial_max[partial_start + split];
- thread_sum += partial_sum[partial_start + split] * expf(split_max - max_logit);
+ if (partial_sum[partial_start + split] != 0.0f) {
+ thread_sum += partial_sum[partial_start + split] * expf(split_max - max_logit);
+ }
}
reduction[h] = thread_sum;
__syncthreads();
@@ -849,10 +1572,117 @@ __global__ void MergeDynamicSparseAttentionKernel(const float* partial_max,
float accumulator = 0.0f;
for (int split = 0; split < split_count; ++split) {
const float split_max = partial_max[partial_start + split];
- const float merge_scale = expf(split_max - max_logit);
- const int64_t partial_offset =
- (static_cast(partial_start + split) * head_size) + h;
- accumulator += partial_output[partial_offset] * merge_scale;
+ if (partial_sum[partial_start + split] != 0.0f) {
+ const float merge_scale = expf(split_max - max_logit);
+ const int64_t partial_offset =
+ (static_cast(partial_start + split) * head_size) + h;
+ accumulator += partial_output[partial_offset] * merge_scale;
+ }
+ }
+ output[static_cast(query_block) * head_size + h] =
+ denominator == 0.0f ? FromFloat(0.0f) : FromFloat(accumulator / denominator);
+ }
+}
+
+template
+__global__ void MergeGroupedDecodeDynamicSparseAttentionKernel(const float* partial_max,
+ const float* partial_sum,
+ const float* partial_output,
+ T* output,
+ int query_block_count,
+ int split_count,
+ int head_size) {
+ const int query_block = static_cast(blockIdx.x);
+ if (query_block >= query_block_count) {
+ return;
+ }
+
+ const int h = static_cast(threadIdx.x);
+ const int lane = h & 31;
+ const int partial_start = query_block * split_count;
+ __shared__ float max_logit;
+ __shared__ float denominator;
+ __shared__ float split_scales[kSharedScales ? 32 : 1];
+ __shared__ unsigned active_splits;
+
+ if constexpr (kSharedScales) {
+ if (h < 32) {
+ const float split_max = lane < split_count ? partial_max[partial_start + lane] : -FLT_MAX;
+ float value = split_max;
+ for (int offset = 16; offset > 0; offset >>= 1) {
+ value = fmaxf(value, __shfl_down_sync(0xffffffff, value, offset));
+ }
+ const float maximum = __shfl_sync(0xffffffff, value, 0);
+ const float sum = lane < split_count ? partial_sum[partial_start + lane] : 0.0f;
+ const bool active = lane < split_count && sum != 0.0f;
+ const float factor = active ? expf(split_max - maximum) : 0.0f;
+ split_scales[lane] = factor;
+ const unsigned mask = __ballot_sync(0xffffffff, active);
+ value = active ? sum * factor : 0.0f;
+ for (int offset = 16; offset > 0; offset >>= 1) {
+ value += __shfl_down_sync(0xffffffff, value, offset);
+ }
+ if (lane == 0) {
+ denominator = value;
+ active_splits = mask;
+ }
+ }
+ __syncthreads();
+
+ if (h < head_size) {
+ float accumulator = 0.0f;
+ const unsigned mask = active_splits;
+ const int end = mask == 0 ? 0 : 32 - __clz(mask);
+ for (int split = 0; split < end; ++split) {
+ // Do not skip an active split whose scale underflows: zero times a
+ // nonfinite partial must retain the generic path's behavior.
+ if (mask & (1u << split)) {
+ const int64_t partial_offset =
+ static_cast(partial_start + split) * head_size + h;
+ accumulator += partial_output[partial_offset] * split_scales[split];
+ }
+ }
+ output[static_cast(query_block) * head_size + h] =
+ denominator == 0.0f ? FromFloat(0.0f) : FromFloat(accumulator / denominator);
+ }
+ return;
+ }
+
+ if (h < 32) {
+ float value = lane < split_count ? partial_max[partial_start + lane] : -FLT_MAX;
+ for (int offset = 16; offset > 0; offset >>= 1) {
+ value = fmaxf(value, __shfl_down_sync(0xffffffff, value, offset));
+ }
+ if (lane == 0) {
+ max_logit = value;
+ }
+ }
+ __syncthreads();
+
+ if (h < 32) {
+ float value = 0.0f;
+ if (lane < split_count && partial_sum[partial_start + lane] != 0.0f) {
+ value = partial_sum[partial_start + lane] *
+ expf(partial_max[partial_start + lane] - max_logit);
+ }
+ for (int offset = 16; offset > 0; offset >>= 1) {
+ value += __shfl_down_sync(0xffffffff, value, offset);
+ }
+ if (lane == 0) {
+ denominator = value;
+ }
+ }
+ __syncthreads();
+
+ if (h < head_size) {
+ float accumulator = 0.0f;
+ for (int split = 0; split < split_count; ++split) {
+ const float split_max = partial_max[partial_start + split];
+ if (partial_sum[partial_start + split] != 0.0f) {
+ const int64_t partial_offset =
+ (static_cast(partial_start + split) * head_size) + h;
+ accumulator += partial_output[partial_offset] * expf(split_max - max_logit);
+ }
}
output[static_cast(query_block) * head_size + h] =
denominator == 0.0f ? FromFloat(0.0f) : FromFloat(accumulator / denominator);
@@ -892,6 +1722,11 @@ int GetAttentionSplitCount(const DynamicSparseAttentionParameters& parameters) {
return static_cast(std::min(candidate_tiles, occupancy_splits));
}
+// A null pointer is treated as aligned because kernels do not dereference it.
+bool IsAligned(const void* pointer, size_t alignment) {
+ return reinterpret_cast(pointer) % alignment == 0;
+}
+
bool UseFusedAttention(const DynamicSparseAttentionParameters& parameters,
size_t element_size,
size_t max_shared_memory_per_block) {
@@ -902,27 +1737,92 @@ bool UseFusedAttention(const DynamicSparseAttentionParameters& parameters,
const int threads = GetThreadsPerBlock(parameters.head_size);
const int fused_threads = threads < 128 ? 128 : threads;
const size_t fused_shared_bytes =
- (2 * static_cast(GetCandidateCapacity(parameters)) + static_cast(fused_threads)) * sizeof(float) +
+ (static_cast(GetCandidateCapacity(parameters)) + static_cast(fused_threads)) * sizeof(float) +
static_cast(parameters.head_size) * element_size;
return fused_shared_bytes <= kMaxFusedSharedBytes &&
fused_shared_bytes < max_shared_memory_per_block;
}
-size_t GetValidationBitmapWords(const DynamicSparseAttentionParameters& parameters) {
+int GetGroupedHeadTiles(const DynamicSparseAttentionParameters& parameters,
+ int query_heads_per_block) {
+ const int query_heads_per_kv = parameters.num_heads / parameters.kv_num_heads;
+ return (query_heads_per_kv + query_heads_per_block - 1) / query_heads_per_block;
+}
+
+size_t GetGroupedAttentionSharedBytes(const DynamicSparseAttentionParameters& parameters,
+ size_t element_size) {
+ return 2 * kGroupedPrefillCandidatesPerTile *
+ static_cast(parameters.head_size) * element_size;
+}
+
+bool SupportsGroupedAttention(const DynamicSparseAttentionParameters& parameters,
+ size_t element_size,
+ size_t max_shared_memory_per_block) {
+ return element_size == sizeof(half) &&
+ parameters.head_size == 256 &&
+ parameters.num_heads > parameters.kv_num_heads &&
+ parameters.num_heads % parameters.kv_num_heads == 0 &&
+ GetGroupedAttentionSharedBytes(parameters, element_size) < max_shared_memory_per_block;
+}
+
+bool UseGroupedPrefillAttention(const DynamicSparseAttentionParameters& parameters,
+ size_t element_size,
+ size_t max_shared_memory_per_block) {
+ return parameters.sequence_length > 1 &&
+ SupportsGroupedAttention(parameters, element_size, max_shared_memory_per_block);
+}
+
+bool UseGroupedDecodeAttention(const DynamicSparseAttentionParameters& parameters,
+ size_t element_size,
+ size_t max_shared_memory_per_block,
+ bool quantized_cache = false) {
+ return parameters.sequence_length >= 1 &&
+ parameters.sequence_length <= 8 &&
+ parameters.num_heads / parameters.kv_num_heads >= 4 &&
+ GetCandidateCapacity(parameters) > kSplitCandidateSize &&
+ SupportsGroupedAttention(parameters, element_size, max_shared_memory_per_block) &&
+ (!quantized_cache ||
+ GetGroupedAttentionSharedBytes(parameters, sizeof(float)) < max_shared_memory_per_block);
+}
+
+int GetGroupedDecodeHeadsPerBlock(const DynamicSparseAttentionParameters& parameters) {
+ const int query_heads_per_kv = parameters.num_heads / parameters.kv_num_heads;
+ return (query_heads_per_kv == 12) ? 12 : kGroupedDecodeQueryHeadsPerBlock;
+}
+
+int GetGroupedDecodeSplitCount(const DynamicSparseAttentionParameters& parameters, bool quantized_cache = false) {
+ const int candidates_per_split = quantized_cache ? kGroupedQuantizedDecodeCandidatesPerSplit
+ : kGroupedDecodeCandidatesPerSplit;
+ const int target_blocks = quantized_cache ? kTargetGroupedQuantizedDecodeBlocks : kTargetGroupedDecodeBlocks;
+ const int64_t candidate_splits =
+ (GetCandidateCapacity(parameters) + candidates_per_split - 1) / candidates_per_split;
+ const int64_t grouped_head_blocks =
+ static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads *
+ GetGroupedHeadTiles(parameters, GetGroupedDecodeHeadsPerBlock(parameters));
+ const int64_t occupancy_splits =
+ std::max(1, (target_blocks + grouped_head_blocks - 1) /
+ grouped_head_blocks);
+ const int64_t max_splits = quantized_cache ? 32 : candidate_splits;
+ return static_cast(std::min({candidate_splits, occupancy_splits, max_splits}));
+}
+
+size_t GetValidationHashEntries(const DynamicSparseAttentionParameters& parameters) {
if (parameters.max_selected == 0) {
return 0;
}
- const size_t source_length =
- parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary
- ? static_cast(parameters.auxiliary_sequence_length)
- : static_cast(parameters.cache_capacity);
- return (source_length + 31) / 32;
+ const size_t max_selected = static_cast(parameters.max_selected);
+ const size_t target_entries = max_selected + (max_selected + 1) / 2;
+ size_t hash_entries = 1;
+ while (hash_entries < target_entries) {
+ hash_entries <<= 1;
+ }
+ return hash_entries;
}
-bool UseSharedValidationBitmap(const DynamicSparseAttentionParameters& parameters,
- size_t max_shared_memory_per_block) {
- const size_t shared_bytes = GetValidationBitmapWords(parameters) * sizeof(uint32_t);
+bool UseSharedValidationHashTable(const DynamicSparseAttentionParameters& parameters,
+ size_t max_shared_memory_per_block) {
+ const size_t shared_bytes = GetValidationHashEntries(parameters) * sizeof(uint32_t);
return shared_bytes <= kMaxValidationSharedBytes &&
shared_bytes < max_shared_memory_per_block;
}
@@ -932,27 +1832,49 @@ bool UseSharedValidationBitmap(const DynamicSparseAttentionParameters& parameter
size_t GetDynamicSparseAttentionValidationWorkspaceSize(
const DynamicSparseAttentionParameters& parameters,
size_t max_shared_memory_per_block) {
- if (UseSharedValidationBitmap(parameters, max_shared_memory_per_block)) {
+ if (UseSharedValidationHashTable(parameters, max_shared_memory_per_block)) {
return 0;
}
const size_t row_count =
SafeInt(parameters.batch_size) * parameters.sequence_length;
- return SafeInt(row_count) * GetValidationBitmapWords(parameters);
+ return SafeInt(row_count) * GetValidationHashEntries(parameters);
}
size_t GetDynamicSparseAttentionWorkspaceSize(
const DynamicSparseAttentionParameters& parameters,
size_t element_size,
- size_t max_shared_memory_per_block) {
- if (UseFusedAttention(parameters, element_size, max_shared_memory_per_block)) {
- return 0;
+ size_t max_shared_memory_per_block,
+ bool quantized_cache) {
+ if (quantized_cache) {
+ if (GetCandidateCapacity(parameters) <= kSplitCandidateSize) return 0;
+ const size_t query_blocks =
+ SafeInt(parameters.batch_size) * parameters.sequence_length * parameters.num_heads;
+ const int split_count =
+ UseGroupedDecodeAttention(parameters, element_size, max_shared_memory_per_block, true)
+ ? GetGroupedDecodeSplitCount(parameters, true)
+ : GetAttentionSplitCount(parameters);
+ return SafeInt(query_blocks) * split_count *
+ (static_cast(parameters.head_size) + 2);
+ }
+ const bool grouped_decode =
+ UseGroupedDecodeAttention(parameters, element_size, max_shared_memory_per_block);
+ if (!grouped_decode) {
+ if (UseGroupedPrefillAttention(parameters, element_size, max_shared_memory_per_block)) {
+ return 0;
+ }
+ if (UseFusedAttention(parameters, element_size, max_shared_memory_per_block)) {
+ return 0;
+ }
}
const size_t query_blocks =
static_cast(parameters.batch_size) * parameters.sequence_length * parameters.num_heads;
- const size_t split_count = GetAttentionSplitCount(parameters);
- if (split_count == 1) {
+ const size_t split_count =
+ grouped_decode
+ ? GetGroupedDecodeSplitCount(parameters)
+ : GetAttentionSplitCount(parameters);
+ if (split_count == 1 && !grouped_decode) {
return 0;
}
return SafeInt(query_blocks) * split_count *
@@ -969,24 +1891,31 @@ Status ValidateDynamicSparseAttentionOnDevice(
int32_t* error_flag,
uint32_t* validation_bitmap,
size_t max_shared_memory_per_block,
- bool copy_result_to_host) {
+ bool copy_result_to_host,
+ const float* key_scale,
+ const float* value_scale,
+ const int32_t* committed_length,
+ int committed_length_stride) {
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(error_flag, 0, sizeof(int32_t), stream));
const int64_t row_count =
static_cast