From 52c73b515ec75a8cc6983ad15434bab3909f53cd Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:18:01 +0000 Subject: [PATCH 01/23] Extend sparse attention caches and optimize packed indexing --- docs/ContribOperators.md | 35 +- .../packed_sparse_attention_indexer_common.h | 4 + .../cuda/bert/dynamic_sparse_attention.cc | 94 +- .../bert/dynamic_sparse_attention_impl.cu | 1115 +++++++++++++++-- .../cuda/bert/dynamic_sparse_attention_impl.h | 19 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 14 +- .../sparse/packed_sparse_attention_indexer.cc | 90 +- .../sparse/packed_sparse_attention_indexer.h | 2 +- .../packed_sparse_attention_indexer_impl.cu | 533 +++++--- .../packed_sparse_attention_indexer_impl.h | 15 +- .../webgpu/bert/dynamic_sparse_attention.cc | 7 + .../bert/packed_sparse_attention_indexer.cc | 45 +- .../core/graph/contrib_ops/bert_defs.cc | 89 +- .../dynamic_sparse_attention_op_test.cc | 351 +++++- ...packed_sparse_attention_indexer_op_test.cc | 131 +- 15 files changed, 2173 insertions(+), 371 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2703de7cc1fae..70be4dc1cba47 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1723,6 +1723,12 @@ This version of the operator has been available since version 1 of the 'com.micr auxiliary_key is used as both K and V. DeepSeek-style post-attention inverse/conjugate RoPE and output projection remain the exporter's responsibility. + CUDA also supports INT8 main caches with float32 per-channel key/value scales. Appended rows are + quantized with round-to-nearest-even and saturation to [-128, 127]; attention reads the quantized + representation, including current rows. Auxiliary caches retain the query's floating-point type. + Scales must be positive and finite; requested validation reports invalid scales as status 10, + and CUDA strict validation raises an error. WebGPU supports only query-typed floating caches. + #### Version This version of the operator has been available since version 1 of the 'com.microsoft' operator set. @@ -1758,7 +1764,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
@@ -1767,10 +1773,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, query-typed or INT8.
+
past_value (optional) : TCache
+
Main value cache in BNSH layout, matching past_key.
auxiliary_key (optional) : T
Read-only auxiliary key sequence in BNSH layout.
auxiliary_value (optional) : T
@@ -1795,17 +1801,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 single-request committed cursor. Inconsistency with sequence metadata produces validation status 9.
+
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]. Both scales are required for INT8 caches.
-#### 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
+
One int32 validation status: zero for valid inputs; 9 for cursor mismatch; 10 for invalid INT8 scales.
#### Type Constraints @@ -1813,6 +1827,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 matches 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.
@@ -8771,4 +8789,3 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
- diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h index 8ed3e7adb01f0..785d2ed4250e2 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h @@ -71,7 +71,11 @@ enum OutputIndex : int { kPresentStateLengths = 5, // [batch_size, 2], int32 kStateUpdate = 6, // qsa only, optional: [batch_size, state_update_capacity, head_size] kIndexShareStatus = 7, + kMidKvBuffer = 8, // qsa only, optional: buffer snapshot after the first request token + kMidStateLengths = 9, // qsa only, optional: state lengths after the first request token + kValidationStatus = 10, // qsa only, optional: per-request overflow/metadata status kOutputCount = 8, + kMaxOutputCount = 11, }; // Every PackedSparseAttentionIndexer node declares all 6 fixed outputs; present_gate_buffer is an diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index 878659302a19a..f7a65fbad3663 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,18 @@ 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 && + (parameters.batch_size != 1 || committed_length->Shape().Size() != 1), + "DynamicSparseAttention: committed_length requires batch_size=1 and one int32 element."); TensorShape output_shape({parameters.batch_size, parameters.sequence_length, parameters.query_hidden_size}); @@ -153,6 +177,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 +186,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 +205,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,8 +237,17 @@ 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); @@ -216,13 +255,18 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons GetScratchBuffer(validation_workspace_elements, GetComputeStream(context)); 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)); + if (committed_length != nullptr) { + ORT_RETURN_IF_ERROR(ValidateDynamicSparseAttentionCursor( + stream, data.seqlens_k, committed_length->Data(), + parameters.sequence_length, parameters.cache_capacity, validation_error)); + } } 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 83750ba4e0eee..ccdfdf76bb7ee 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -6,10 +6,12 @@ #include #include #include +#include #include #include #include +#include #include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" @@ -24,6 +26,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 = 256; template __device__ __forceinline__ float ToFloat(T value) { @@ -35,6 +44,33 @@ __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 +} + __device__ __forceinline__ void SetValidationError(int32_t* error_flag, DynamicSparseAttentionValidationError error) { atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, static_cast(error)); @@ -53,21 +89,32 @@ __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 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(); @@ -116,10 +163,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; @@ -134,10 +181,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) { @@ -301,7 +357,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; @@ -367,8 +426,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]; + } } } @@ -396,7 +463,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; @@ -432,8 +502,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]; @@ -448,7 +517,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) { @@ -466,7 +535,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; } } } @@ -474,15 +543,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(); @@ -491,9 +566,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(); @@ -513,7 +586,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; } @@ -532,19 +605,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; @@ -552,10 +627,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 = @@ -565,6 +647,557 @@ __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; + reinterpret_cast(shared_key + shared_offset)[load_lane] = + reinterpret_cast(key_head)[load_lane]; + reinterpret_cast(shared_value + shared_offset)[load_lane] = + reinterpret_cast(value_head)[load_lane]; + } + } + __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); + } + } +} + +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]; + __shared__ half mma_query[kGroupedDecodeQueryHeadsPerBlock * 256]; + __shared__ float mma_scores[kGroupedDecodeQueryHeadsPerBlock * 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 < kGroupedDecodeQueryHeadsPerBlock * 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) { + 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; + reinterpret_cast(shared_key + shared_offset)[load_lane] = + reinterpret_cast(key_head)[load_lane]; + reinterpret_cast(shared_value + shared_offset)[load_lane] = + reinterpret_cast(value_head)[load_lane]; + } + } + } + __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 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), + 0, + *reinterpret_cast( + mma_query + row_pair * head_size + second_k), + 0}; + 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]; + } + __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, @@ -592,7 +1225,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; @@ -629,8 +1265,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]; @@ -659,7 +1294,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) { @@ -677,7 +1312,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; } } } @@ -685,15 +1320,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(); @@ -701,9 +1342,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(); @@ -722,7 +1361,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; } @@ -741,10 +1380,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; @@ -752,9 +1387,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; @@ -762,10 +1403,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; @@ -828,7 +1476,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(); @@ -845,10 +1495,72 @@ __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; + + 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); @@ -898,56 +1610,150 @@ 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 || + (quantized_cache && parameters.sequence_length == 2)) && + 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 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, kGroupedDecodeQueryHeadsPerBlock); + const int64_t occupancy_splits = + std::max(1, (target_blocks + grouped_head_blocks - 1) / + grouped_head_blocks); + return static_cast(std::min(candidate_splits, occupancy_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; } +__global__ void ValidateCommittedLengthKernel( + const int32_t* seqlens_k, + const int32_t* committed_length, + int32_t sequence_length, + int32_t cache_capacity, + int32_t* status) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + const int64_t committed = *committed_length; + const int64_t total = committed + sequence_length; + if (*status == kDynamicSparseAttentionValidationOk && + (committed < 0 || total > cache_capacity || *seqlens_k != total - 1)) { + *status = 9; + } +} + } // namespace 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) { + 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); + } + 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); + const size_t split_count = + UseGroupedDecodeAttention(parameters, element_size, max_shared_memory_per_block) + ? GetGroupedDecodeSplitCount(parameters) + : GetAttentionSplitCount(parameters); if (split_count == 1) { return 0; } @@ -965,24 +1771,29 @@ 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) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(error_flag, 0, sizeof(int32_t), stream)); const int64_t row_count = static_cast(parameters.batch_size) * parameters.sequence_length; - const size_t validation_bitmap_words = GetValidationBitmapWords(parameters); - const bool use_shared_bitmap = - UseSharedValidationBitmap(parameters, max_shared_memory_per_block); - const size_t validation_bitmap_elements = - use_shared_bitmap ? 0 : static_cast(row_count) * validation_bitmap_words; - if (validation_bitmap_elements > 0) { + const size_t validation_hash_entries = GetValidationHashEntries(parameters); + const bool use_shared_hash_table = + UseSharedValidationHashTable(parameters, max_shared_memory_per_block); + const size_t validation_hash_elements = + use_shared_hash_table ? 0 : static_cast(row_count) * validation_hash_entries; + if (validation_hash_elements > 0) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync( - validation_bitmap, 0, validation_bitmap_elements * sizeof(uint32_t), stream)); + validation_bitmap, 0xff, validation_hash_elements * sizeof(uint32_t), stream)); } ORT_RETURN_IF_ERROR(CheckBlockCount(row_count, "validation")); - constexpr int kValidationThreads = 256; + int validation_threads = 32; + while (validation_threads < parameters.max_selected && validation_threads < 1024) { + validation_threads <<= 1; + } const size_t validation_shared_bytes = - use_shared_bitmap ? validation_bitmap_words * sizeof(uint32_t) : 0; - ValidateInputsKernel<<(row_count), kValidationThreads, + use_shared_hash_table ? validation_hash_entries * sizeof(uint32_t) : 0; + ValidateInputsKernel<<(row_count), validation_threads, validation_shared_bytes, stream>>>( selected_indices, selected_counts, seqlens_k, position_ids, parameters.batch_size, parameters.sequence_length, parameters.max_selected, @@ -990,7 +1801,8 @@ Status ValidateDynamicSparseAttentionOnDevice( parameters.cache_capacity, parameters.auxiliary_sequence_length, parameters.rotary_max_position, parameters.do_rotary, parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, - validation_bitmap, validation_bitmap_words, use_shared_bitmap, error_flag); + validation_bitmap, validation_hash_entries, use_shared_hash_table, error_flag, + key_scale, value_scale, parameters.kv_num_heads * parameters.head_size); CUDA_RETURN_IF_ERROR(cudaGetLastError()); if (!copy_result_to_host) { @@ -1031,6 +1843,9 @@ Status ValidateDynamicSparseAttentionOnDevice( case kDynamicSparseAttentionInvalidPosition: message = "a rotary position is outside the rotary cache"; break; + case kDynamicSparseAttentionInvalidCacheScale: + message = "INT8 cache scales must be positive and finite"; + break; default: break; } @@ -1038,6 +1853,18 @@ Status ValidateDynamicSparseAttentionOnDevice( "DynamicSparseAttention: ", message, "."); } +Status ValidateDynamicSparseAttentionCursor( + cudaStream_t stream, + const int32_t* seqlens_k, + const int32_t* committed_length, + int32_t sequence_length, + int32_t cache_capacity, + int32_t* status) { + ValidateCommittedLengthKernel<<<1, 1, 0, stream>>>( + seqlens_k, committed_length, sequence_length, cache_capacity, status); + return CUDA_CALL(cudaGetLastError()); +} + template Status LaunchDynamicSparseAttention( cudaStream_t stream, @@ -1055,7 +1882,7 @@ Status LaunchDynamicSparseAttention( static_cast(parameters.batch_size) * parameters.kv_num_heads * parameters.cache_capacity * parameters.head_size; const size_t cache_bytes = - SafeInt(cache_elements) * sizeof(T); + SafeInt(cache_elements) * (data.quantized_cache ? 1 : sizeof(T)); ORT_RETURN_IF_NOT( data.past_key == nullptr || parameters.past_cache_capacity == parameters.cache_capacity, @@ -1107,7 +1934,7 @@ Status LaunchDynamicSparseAttention( parameters.kv_hidden_size, packed_stride, parameters.cache_capacity, parameters.rotary_dim, parameters.rotary_max_position, parameters.rotary_offset, parameters.rotary_interleaved, parameters.is_packed_qkv, - parameters.qk_norm_epsilon); + parameters.qk_norm_epsilon, data.quantized_cache, data.key_scale, data.value_scale); CUDA_RETURN_IF_ERROR(cudaGetLastError()); ORT_RETURN_IF_ERROR(CheckBlockCount(query_blocks, "attention")); @@ -1121,9 +1948,107 @@ Status LaunchDynamicSparseAttention( const int candidate_capacity = static_cast(candidate_capacity_64); const int fused_threads = threads < 128 && max_threads_per_block >= 128 ? 128 : threads; const size_t fused_shared_bytes = - (2 * static_cast(candidate_capacity) + static_cast(fused_threads)) * sizeof(float) + + (static_cast(candidate_capacity) + static_cast(fused_threads)) * sizeof(float) + static_cast(parameters.head_size) * sizeof(T); - if (UseFusedAttention(parameters, sizeof(T), max_shared_memory_per_block)) { + constexpr int kGroupedDecodeThreads = 32 * kGroupedDecodeQueryHeadsPerBlock; + constexpr int kGroupedPrefillThreads = 32 * kGroupedPrefillQueryHeadsPerBlock; + const bool grouped_decode_threads_supported = kGroupedDecodeThreads <= max_threads_per_block; + const bool grouped_prefill_threads_supported = kGroupedPrefillThreads <= max_threads_per_block; + if (grouped_decode_threads_supported && + UseGroupedDecodeAttention(parameters, sizeof(T), max_shared_memory_per_block, data.quantized_cache)) { + const int split_count = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); + const int64_t partial_blocks = query_blocks * split_count; + const int64_t grouped_blocks = + static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads * + GetGroupedHeadTiles(parameters, kGroupedDecodeQueryHeadsPerBlock) * split_count; + ORT_RETURN_IF_ERROR(CheckBlockCount(partial_blocks, "grouped decode partials")); + ORT_RETURN_IF_ERROR(CheckBlockCount(grouped_blocks, "grouped decode attention")); + ORT_RETURN_IF_NOT(data.attention_workspace != nullptr, + "DynamicSparseAttention: grouped decode workspace is required."); + float* partial_max = data.attention_workspace; + float* partial_sum = partial_max + partial_blocks; + float* partial_output = partial_sum + partial_blocks; + const size_t grouped_shared_bytes = + GetGroupedAttentionSharedBytes(parameters, data.quantized_cache ? sizeof(float) : sizeof(T)); + if (data.quantized_cache) { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, data.key_scale, data.value_scale); + } else if (parameters.cache_capacity <= 131072) { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, nullptr, nullptr); + } else { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, nullptr, nullptr); + } + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + + if (split_count <= 32) { + MergeGroupedDecodeDynamicSparseAttentionKernel<<(query_blocks), threads, 0, stream>>>( + partial_max, partial_sum, partial_output, data.output, + static_cast(query_blocks), split_count, parameters.head_size); + } else { + const size_t merge_shared_bytes = static_cast(threads) * sizeof(float); + MergeDynamicSparseAttentionKernel<<(query_blocks), threads, merge_shared_bytes, stream>>>( + partial_max, partial_sum, partial_output, data.output, + static_cast(query_blocks), split_count, parameters.head_size); + } + } else if (!data.quantized_cache && grouped_prefill_threads_supported && + UseGroupedPrefillAttention(parameters, sizeof(T), max_shared_memory_per_block)) { + const int64_t grouped_blocks = + static_cast(parameters.batch_size) * parameters.sequence_length * + parameters.kv_num_heads * GetGroupedHeadTiles(parameters, kGroupedPrefillQueryHeadsPerBlock); + ORT_RETURN_IF_ERROR(CheckBlockCount(grouped_blocks, "grouped prefill attention")); + const size_t grouped_shared_bytes = + GetGroupedAttentionSharedBytes(parameters, sizeof(T)); + GroupedPrefillDynamicSparseAttentionKernel + <<(grouped_blocks), kGroupedPrefillThreads, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, data.output, + static_cast(grouped_blocks), parameters.sequence_length, parameters.num_heads, + parameters.kv_num_heads, parameters.head_size, parameters.cache_capacity, + parameters.auxiliary_sequence_length, parameters.max_selected, + parameters.local_window_size, parameters.scale, local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax); + } else if (data.quantized_cache + ? candidate_capacity <= kSplitCandidateSize + : UseFusedAttention(parameters, sizeof(T), max_shared_memory_per_block)) { + ORT_RETURN_IF_NOT(fused_shared_bytes <= max_shared_memory_per_block, + "DynamicSparseAttention: quantized cache attention exceeds shared-memory capacity."); FusedDynamicSparseAttentionKernel<<(query_blocks), fused_threads, fused_shared_bytes, stream>>>( data.prepared_query, data.present_key, data.present_value, data.auxiliary_key, data.auxiliary_value, data.selected_indices, @@ -1134,7 +2059,7 @@ Status LaunchDynamicSparseAttention( parameters.local_window_size, candidate_capacity, parameters.scale, local_plus_selected, parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, - parameters.use_smooth_softmax); + parameters.use_smooth_softmax, data.quantized_cache, data.key_scale, data.value_scale); } else { const int split_count = GetAttentionSplitCount(parameters); const int64_t partial_blocks = query_blocks * split_count; @@ -1146,7 +2071,7 @@ Status LaunchDynamicSparseAttention( float* partial_output = split_count == 1 ? nullptr : partial_sum + partial_blocks; const int split_threads = threads < 128 && max_threads_per_block >= 128 ? 128 : threads; const size_t split_shared_bytes = - (2 * static_cast(kSplitCandidateSize) + static_cast(split_threads)) * sizeof(float) + + (static_cast(kSplitCandidateSize) + static_cast(split_threads)) * sizeof(float) + static_cast(parameters.head_size) * sizeof(T); ORT_RETURN_IF_NOT(split_shared_bytes < max_shared_memory_per_block, "DynamicSparseAttention: split attention exceeds the CUDA shared-memory limit."); @@ -1162,7 +2087,7 @@ Status LaunchDynamicSparseAttention( parameters.local_window_size, parameters.scale, local_plus_selected, parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, - parameters.use_smooth_softmax); + parameters.use_smooth_softmax, data.quantized_cache, data.key_scale, data.value_scale); CUDA_RETURN_IF_ERROR(cudaGetLastError()); if (split_count > 1) { diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index d60bc7b9c69f7..2367f77b89630 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -22,6 +22,7 @@ enum DynamicSparseAttentionValidationError : int32_t { kDynamicSparseAttentionNonCausalIndex = 6, kDynamicSparseAttentionInvalidSequenceLength = 7, kDynamicSparseAttentionInvalidPosition = 8, + kDynamicSparseAttentionInvalidCacheScale = 10, }; template @@ -47,6 +48,9 @@ struct DynamicSparseAttentionData { T* present_value = nullptr; T* output = nullptr; float* attention_workspace = nullptr; + bool quantized_cache = false; + const float* key_scale = nullptr; + const float* value_scale = nullptr; }; size_t GetDynamicSparseAttentionValidationWorkspaceSize( @@ -56,7 +60,8 @@ size_t GetDynamicSparseAttentionValidationWorkspaceSize( size_t GetDynamicSparseAttentionWorkspaceSize( const DynamicSparseAttentionParameters& parameters, size_t element_size, - size_t max_shared_memory_per_block); + size_t max_shared_memory_per_block, + bool quantized_cache = false); Status ValidateDynamicSparseAttentionOnDevice( cudaStream_t stream, @@ -68,7 +73,17 @@ 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 = nullptr, + const float* value_scale = nullptr); + +Status ValidateDynamicSparseAttentionCursor( + cudaStream_t stream, + const int32_t* seqlens_k, + const int32_t* committed_length, + int32_t sequence_length, + int32_t cache_capacity, + int32_t* status); template Status LaunchDynamicSparseAttention( diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 546b5b12e178f..dcfef1542edb6 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -277,9 +277,10 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, SparseAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, SparseAttentionIndexer); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, SparseAttentionIndexer); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, SparseAttentionIndexer); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, PackedSparseAttentionIndexer); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, PackedSparseAttentionIndexer); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, PackedSparseAttentionIndexer); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float_float, PackedSparseAttentionIndexer); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16_MLFloat16, PackedSparseAttentionIndexer); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16_BFloat16, PackedSparseAttentionIndexer); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16_float, PackedSparseAttentionIndexer); class CUDA_MS_OP_CLASS_NAME(1, PackedSparseAttentionIndexerMerge); class CUDA_MS_OP_THREE_TYPED_CLASS_NAME(1, uint8_t, float, int32_t, GatherBlockQuantized); @@ -647,9 +648,10 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index 01bbefdf685af..eaedf62a4ffca 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -7,6 +7,7 @@ #include #include #include +#include #include "contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h" #include "core/providers/cuda/cuda_common.h" @@ -19,15 +20,16 @@ namespace cuda { using namespace onnxruntime::cuda; namespace psai = onnxruntime::contrib::packed_sparse_attention_indexer; -#define REGISTER_KERNEL_TYPED(T) \ +#define REGISTER_KERNEL_TYPED(K, T, TState) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ PackedSparseAttentionIndexer, \ kMSDomain, \ 1, \ - T, \ + K, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_KEY", DataTypeImpl::GetTensorType()) \ .TypeConstraint("I", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ .InputMemoryType(OrtMemTypeCPUInput, psai::kStateUpdateActive) \ @@ -36,14 +38,13 @@ namespace psai = onnxruntime::contrib::packed_sparse_attention_indexer; .MayInplace(psai::kMergedCounts, psai::kSelectedCounts) \ .MayInplace(psai::kMergedStatus, psai::kIndexShareStatus) \ .MayInplace(psai::kPastKeyState, psai::kPresentKeyState) \ - .MayInplace(psai::kPastKvBuffer, psai::kPresentKvBuffer) \ - .MayInplace(psai::kPastGateBuffer, psai::kPresentGateBuffer) \ .MayInplace(psai::kPastStateLengths, psai::kPresentStateLengths), \ - PackedSparseAttentionIndexer); + PackedSparseAttentionIndexer); -REGISTER_KERNEL_TYPED(float) -REGISTER_KERNEL_TYPED(MLFloat16) -REGISTER_KERNEL_TYPED(BFloat16) +REGISTER_KERNEL_TYPED(float_float, float, float) +REGISTER_KERNEL_TYPED(MLFloat16_MLFloat16, MLFloat16, MLFloat16) +REGISTER_KERNEL_TYPED(BFloat16_BFloat16, BFloat16, BFloat16) +REGISTER_KERNEL_TYPED(BFloat16_float, BFloat16, float) #undef REGISTER_KERNEL_TYPED @@ -94,13 +95,17 @@ Status CheckRotaryCache(const Tensor* cos_cache, const Tensor* sin_cache, int64_ } // namespace -template -PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo& info) : CudaKernel(info) { +template +PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo& info) : CudaKernel(info) { std::string policy_mode; ORT_ENFORCE(info.GetAttr("policy_mode", &policy_mode).IsOK(), "PackedSparseAttentionIndexer: policy_mode is required"); ORT_ENFORCE(psai::TryParsePolicy(policy_mode, policy_), "PackedSparseAttentionIndexer: policy_mode must be '", psai::kPolicyModeQsa, "' or '", psai::kPolicyModeCsa, "', got '", policy_mode, "'"); + if constexpr (!std::is_same_v) { + ORT_ENFORCE(policy_ == psai::Policy::kQsa, + "PackedSparseAttentionIndexer: mixed input/state types are supported only for policy_mode 'qsa'"); + } ORT_ENFORCE(info.GetAttr("compress_ratio", &compress_ratio_).IsOK(), "PackedSparseAttentionIndexer: compress_ratio is required"); @@ -163,8 +168,8 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo "PackedSparseAttentionIndexer: max_output_entries must cover the original selection capacity"); } -template -Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context) const { +template +Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context) const { bool reuse = false; const Tensor* mode = context->InputCount() > psai::kIndexShareMode ? context->Input(psai::kIndexShareMode) @@ -211,9 +216,9 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context } const bool position_ids_provided = psai::kPositionIds < context->InputCount() && context->Input(psai::kPositionIds) != nullptr; - ORT_RETURN_IF(position_ids_provided == is_qsa, - "PackedSparseAttentionIndexer: position_ids is required for " - "policy_mode 'csa' and must be omitted for policy_mode 'qsa'"); + ORT_RETURN_IF(!is_qsa && !position_ids_provided, + "PackedSparseAttentionIndexer: input ", psai::kPositionIds, + " (position_ids) is required for policy_mode 'csa'"); const bool gate_buffer_provided = psai::kPastGateBuffer < context->InputCount() && context->Input(psai::kPastGateBuffer) != nullptr; ORT_RETURN_IF(gate_buffer_provided != !is_qsa, "PackedSparseAttentionIndexer: input ", psai::kPastGateBuffer, @@ -233,9 +238,10 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context return is_qsa ? ComputeQsa(context) : ComputeCsa(context); } -template -Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) const { +template +Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) const { using CudaT = typename OrtToCudaType::type; + using CudaStateT = typename OrtToCudaType::type; const Tensor* query = context->Input(psai::kQuery); const Tensor* key = context->Input(psai::kKey); @@ -304,6 +310,9 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con } ORT_RETURN_IF_ERROR(CheckShape(query_norm_weight, "query_norm_weight", {head_size})); ORT_RETURN_IF_ERROR(CheckShape(key_norm_weight, "key_norm_weight", {head_size})); + if (position_ids != nullptr) { + ORT_RETURN_IF_ERROR(CheckShape(position_ids, "position_ids", {total_tokens})); + } RotaryCacheShape rotary; ORT_RETURN_IF_ERROR(CheckRotaryCache(cos_cache, sin_cache, batch_size, rotary)); @@ -353,10 +362,24 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con ? context->Output( psai::kStateUpdate, TensorShape({batch_size, state_update_capacity_, head_size})) : nullptr; + Tensor* mid_kv_buffer = + context->OutputCount() > psai::kMidKvBuffer + ? context->Output(psai::kMidKvBuffer, TensorShape({batch_size, buffer_capacity, head_size})) + : nullptr; + Tensor* mid_state_lengths = + context->OutputCount() > psai::kMidStateLengths + ? context->Output(psai::kMidStateLengths, TensorShape({batch_size, psai::kStateLengthColumns})) + : nullptr; + Tensor* validation_status = + context->OutputCount() > psai::kValidationStatus + ? context->Output(psai::kValidationStatus, TensorShape({batch_size})) + : nullptr; ORT_RETURN_IF(selected_indices == nullptr || selected_counts == nullptr || present_key_state == nullptr || present_kv_buffer == nullptr || present_state_lengths == nullptr, "PackedSparseAttentionIndexer: policy_mode 'qsa' requires selected_indices, selected_counts, " "present_key_state, present_kv_buffer and present_state_lengths outputs"); + ORT_RETURN_IF((mid_kv_buffer == nullptr) != (mid_state_lengths == nullptr), + "PackedSparseAttentionIndexer: mid_kv_buffer and mid_state_lengths must be connected together"); if (state_update != nullptr && state_update->SizeInBytes() > 0) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(state_update->MutableDataRaw(), 0, state_update->SizeInBytes(), Stream(context))); @@ -382,6 +405,10 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size)); params.block_topk = static_cast(token_budget_ / compress_ratio_); params.state_update_capacity = static_cast(state_update_capacity_); + params.mid_kv_buffer = + mid_kv_buffer == nullptr ? nullptr : reinterpret_cast(mid_kv_buffer->MutableData()); + params.mid_state_lengths = + mid_state_lengths == nullptr ? nullptr : mid_state_lengths->MutableData(); auto float_workspace = GetScratchBuffer(GetQsaPackedWorkspaceFloatCount(params), GetComputeStream(context)); auto overflow_flags = @@ -391,7 +418,7 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con const CudaT* key_data = has_separate_key ? reinterpret_cast(key->Data()) : query_data + query_width; - return LaunchQsaPackedSparseAttentionIndexer( + ORT_RETURN_IF_ERROR((LaunchQsaPackedSparseAttentionIndexer( Stream(context), params, query_data, key_data, @@ -402,23 +429,29 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con cumulative_sequence_lengths->Data(), past_sequence_lengths->Data(), position_ids != nullptr ? position_ids->Data() : nullptr, - reinterpret_cast(past_key_state->Data()), + reinterpret_cast(past_key_state->Data()), reinterpret_cast(past_kv_buffer->Data()), past_state_lengths->Data(), state_update_capture_count != nullptr ? state_update_capture_count->Data() : nullptr, state_update_active == nullptr || state_update_active->Data()[0] != 0, selected_indices->MutableData(), selected_counts->MutableData(), - reinterpret_cast(present_key_state->MutableData()), + reinterpret_cast(present_key_state->MutableData()), reinterpret_cast(present_kv_buffer->MutableData()), present_state_lengths->MutableData(), state_update != nullptr ? reinterpret_cast(state_update->MutableData()) : nullptr, float_workspace.get(), - overflow_flags.get()); + overflow_flags.get()))); + if (validation_status != nullptr) { + return LaunchMapQsaValidationStatus( + Stream(context), cumulative_sequence_lengths->Data(), overflow_flags.get(), + static_cast(batch_size), validation_status->MutableData()); + } + return Status::OK(); } -template -Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, const bool state_only) const { +template +Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, const bool state_only) const { using CudaT = typename OrtToCudaType::type; const Tensor* query = context->Input(psai::kQuery); @@ -611,8 +644,8 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, con static_cast(compress_ratio_), false, overflow_flags.get()); } -template -Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { +template +Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { const Tensor* query = context->Input(psai::kQuery); const Tensor* lengths = context->Input(psai::kPastSequenceLengths); const Tensor* indices = context->Input(psai::kMergedIndices); @@ -714,9 +747,10 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c static_cast(state_capacity_), static_cast(compress_ratio_), is_qsa); } -template class PackedSparseAttentionIndexer; -template class PackedSparseAttentionIndexer; -template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h index 9330edc01b0a8..1ce56f0fd229f 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h @@ -11,7 +11,7 @@ namespace onnxruntime { namespace contrib { namespace cuda { -template +template class PackedSparseAttentionIndexer final : public onnxruntime::cuda::CudaKernel { public: explicit PackedSparseAttentionIndexer(const OpKernelInfo& info); diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 3304a7839dd3d..b61b2d8c80b47 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -36,6 +36,7 @@ #include "contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h" #include "contrib_ops/cuda/sparse/sparse_attention_indexer_device_math.cuh" +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cu_inc/topk_warp_sort.cuh" @@ -56,7 +57,10 @@ constexpr int kBoundedTopKItemsPerThread = kBoundedTopKMax / kThreads; constexpr int kDistributedTopKMaxRows = 64; constexpr int kDistributedTopKMinBlocks = 2048; constexpr int kHierarchicalTileBlocks = 8; +constexpr int kHierarchicalTopKTileBlocks = 32; constexpr int kHierarchicalScoreThreads = 1024; +constexpr int kHierarchicalTileScoreThreads = 512; +constexpr int kHierarchicalScoreMaxBlocksPerRow = 128; constexpr int kHierarchicalMergeThreads = 256; constexpr int kQwenHeadSize = 128; constexpr int kQwenCompressRatio = 4; @@ -67,10 +71,8 @@ Status ConfigureDynamicSharedMemory(Kernel kernel, size_t bytes, const char* ker int device = 0; int max_bytes = 0; CUDA_RETURN_IF_ERROR(cudaGetDevice(&device)); - CUDA_RETURN_IF_ERROR( - cudaDeviceGetAttribute(&max_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); - ORT_RETURN_IF_ERROR(psai::ValidateDynamicSharedMemory( - bytes, static_cast(max_bytes), kernel_name)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&max_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + ORT_RETURN_IF_ERROR(psai::ValidateDynamicSharedMemory(bytes, static_cast(max_bytes), kernel_name)); return CUDA_CALL(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_bytes)); } @@ -89,7 +91,7 @@ bool UseRaggedQsaPrefill(const PackedSparseAttentionIndexerParams& params) { } int GetHierarchicalTileCount(const PackedSparseAttentionIndexerParams& params) { - return (params.state_capacity + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; + return (params.state_capacity + kHierarchicalTopKTileBlocks - 1) / kHierarchicalTopKTileBlocks; } template @@ -138,13 +140,13 @@ __global__ void ElementwiseCopyKernel(const T* src, T* dst, int64_t count) { // leading RoPE -> append into fixed-capacity key_state) and publishes the raw trailing buffer. // Reads only past_kv_buffer / key (never present_kv_buffer / present_key_state), so it is correct // whether or not the present/past tensors are the same aliased allocation. -template +template __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, const T* cos_cache, const T* sin_cache, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const T* past_kv_buffer, const int32_t* past_state_lengths, const int32_t* state_update_capture_count, bool state_update_active, - T* present_key_state, T* present_kv_buffer, + TState* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths, T* state_update, int32_t* overflow_flags, PackedSparseAttentionIndexerParams params) { extern __shared__ float shared[]; @@ -196,6 +198,36 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con overflow_flags[b] = rejected ? 1 : 0; } + if (params.mid_kv_buffer != nullptr && params.mid_state_lengths != nullptr) { + const int mid_req = min(req_len, 1); + const int mid_pending = old_buf_len + mid_req; + const int mid_blocks = rejected ? 0 : mid_pending / params.compress_ratio; + const int mid_buf_len = rejected ? old_buf_len : mid_pending % params.compress_ratio; + if (threadIdx.x == 0) { + params.mid_state_lengths[b * 2 + psai::kKeyStateLength] = old_key_len + mid_blocks; + params.mid_state_lengths[b * 2 + psai::kBufferLength] = mid_buf_len; + } + T* mid_buffer = static_cast(params.mid_kv_buffer); + const int mid_total = params.buffer_capacity * params.head_size; + for (int i = static_cast(threadIdx.x); i < mid_total; i += static_cast(blockDim.x)) { + const int t = i / params.head_size; + const int d = i - t * params.head_size; + const int virtual_pos = mid_blocks * params.compress_ratio + t; + const int64_t slot = static_cast(b) * params.buffer_capacity * params.head_size + i; + if (t >= mid_buf_len || rejected) { + mid_buffer[slot] = past_kv_buffer[slot]; + } else { + mid_buffer[slot] = + virtual_pos < old_buf_len + ? past_kv_buffer[(static_cast(b) * params.buffer_capacity + virtual_pos) * + params.head_size + + d] + : key[(static_cast(req_start) + (virtual_pos - old_buf_len)) * params.key_row_stride + + d]; + } + } + } + for (int t = static_cast(threadIdx.x); t < capture_count; t += static_cast(blockDim.x)) { if ((old_buf_len + t + 1) % params.compress_ratio != 0) { const int64_t capture_base = @@ -253,7 +285,7 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con const int64_t capture_base = (static_cast(b) * params.state_update_capacity + completion_token) * params.head_size; for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { - present_key_state[out_base + d] = from_float(rotated[d]); + present_key_state[out_base + d] = from_float(rotated[d]); if (completion_token < capture_count) { state_update[capture_base + d] = from_float(rotated[d]); } @@ -345,8 +377,8 @@ __global__ void PackedRotateQueryKernel(const T* query, const T* query_norm_weig // One block per (token, key_state slot). Scores are directly dotted against the already-prepared // present_key_state entry (unlike the dense op, no per-query pooling/normalize/rotate is repeated // here because the packed contract stores fully-prepared blocks in key_state). -template -__global__ void QsaBlockScoreKernel(const T* present_key_state, const float* query_rotated, +template +__global__ void QsaBlockScoreKernel(const TState* present_key_state, const float* query_rotated, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, const int32_t* present_state_lengths, float* block_scores, @@ -386,7 +418,7 @@ __global__ void QsaBlockScoreKernel(const T* present_key_state, const float* que query_rotated + (static_cast(token) * params.num_heads + head) * params.head_size; float partial = 0.0f; for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { - partial += query_head[d] * to_float(present_key_state[key_base + d]); + partial += query_head[d] * to_float(present_key_state[key_base + d]); } score += fmaxf(SaiBlockSum(partial, reduction), 0.0f); } @@ -400,8 +432,9 @@ __global__ void QsaBlockScoreKernel(const T* present_key_state, const float* que // One CTA per packed query row. Eight state blocks are scored concurrently, with one warp per // (state block, query head). Unlike the generic rectangular launch, each CTA stops at its row's // live causal block count and never launches or writes work for unused fixed-capacity slots. -template -__global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const float* query_rotated, +// WarpHeads instead keeps all heads in one warp to reuse keys without CTA-wide barriers. +template +__global__ void QsaRaggedPrefillScoreKernel(const TState* present_key_state, const float* query_rotated, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, const int32_t* present_state_lengths, const int32_t* overflow_flags, @@ -410,7 +443,7 @@ __global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const fl const int warp = static_cast(threadIdx.x) / kWarpSize; const int lane = static_cast(threadIdx.x) % kWarpSize; - const int block_slot = warp / kQwenNumHeads; + const int block_slot = WarpHeads ? warp : warp / kQwenNumHeads; const int head = warp % kQwenNumHeads; for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) { const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); @@ -423,55 +456,84 @@ __global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const fl const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; - for (int tile_start = 0; tile_start < block_count; tile_start += kHierarchicalTileBlocks) { - const int block_index = tile_start + block_slot; - float score = 0.0f; - if (block_index < block_count) { - const int d = lane * 4; - const int64_t query_base = - (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; + if constexpr (WarpHeads) { + float4 queries[kQwenNumHeads]; +#pragma unroll + for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) { + queries[query_head] = LoadQsaVector4( + query_rotated + (static_cast(token) * kQwenNumHeads + query_head) * kQwenHeadSize + lane * 4); + } + for (int block_index = block_slot; block_index < block_count; block_index += kHierarchicalTileBlocks) { const int64_t key_base = (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; - const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); - const float4 key_value = LoadQsaVector4(present_key_state + key_base + d); - score = query_value.x * key_value.x + query_value.y * key_value.y + - query_value.z * key_value.z + query_value.w * key_value.w; - score = topk::WarpReduceSum(score); - } - if (lane == 0) { - head_scores[block_slot][head] = fmaxf(score, 0.0f); - } - __syncthreads(); - if (threadIdx.x < kHierarchicalTileBlocks && tile_start + threadIdx.x < block_count) { + const float4 key_value = LoadQsaVector4(present_key_state + key_base + lane * 4); float total_score = 0.0f; #pragma unroll for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) { - total_score += head_scores[threadIdx.x][query_head]; + const float4 query_value = queries[query_head]; + float score = query_value.x * key_value.x + query_value.y * key_value.y + + query_value.z * key_value.z + query_value.w * key_value.w; + score = topk::WarpReduceSum(score); + total_score += fmaxf(score, 0.0f); + } + if (lane == 0) { + const float scaled_score = total_score * params.scale; + block_scores[static_cast(token) * params.state_capacity + block_index] = + scaled_score == 0.0f ? 0.0f : scaled_score; } - const float scaled_score = total_score * params.scale; - block_scores[static_cast(token) * params.state_capacity + tile_start + threadIdx.x] = - scaled_score == 0.0f ? 0.0f : scaled_score; } - __syncthreads(); + } else { + for (int tile_start = 0; tile_start < block_count; tile_start += kHierarchicalTileBlocks) { + const int block_index = tile_start + block_slot; + float score = 0.0f; + if (block_index < block_count) { + const int d = lane * 4; + const int64_t query_base = + (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; + const int64_t key_base = + (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; + const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); + const float4 key_value = LoadQsaVector4(present_key_state + key_base + d); + score = query_value.x * key_value.x + query_value.y * key_value.y + + query_value.z * key_value.z + query_value.w * key_value.w; + score = topk::WarpReduceSum(score); + } + if (lane == 0) { + head_scores[block_slot][head] = fmaxf(score, 0.0f); + } + __syncthreads(); + if (threadIdx.x < kHierarchicalTileBlocks && tile_start + threadIdx.x < block_count) { + float total_score = 0.0f; +#pragma unroll + for (int query_head = 0; query_head < kQwenNumHeads; ++query_head) { + total_score += head_scores[threadIdx.x][query_head]; + } + const float scaled_score = total_score * params.scale; + block_scores[static_cast(token) * params.state_capacity + tile_start + threadIdx.x] = + scaled_score == 0.0f ? 0.0f : scaled_score; + } + __syncthreads(); + } } } } -template -__global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* query_rotated, +template +__global__ void QsaScoreTileTopKKernel(const TState* present_key_state, const float* query_rotated, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, const int32_t* present_state_lengths, const int32_t* overflow_flags, uint64_t* tile_keys, int tile_count, PackedSparseAttentionIndexerParams params) { - __shared__ uint64_t keys[kHierarchicalTileBlocks]; - __shared__ float head_scores[kHierarchicalTileBlocks][kQwenNumHeads]; + __shared__ uint64_t keys[kHierarchicalTopKTileBlocks]; + __shared__ float head_scores[kHierarchicalTopKTileBlocks][kQwenNumHeads]; const int warp = static_cast(threadIdx.x) / kWarpSize; const int lane = static_cast(threadIdx.x) % kWarpSize; - const int block_slot = warp / kQwenNumHeads; - const int head = warp % kQwenNumHeads; - const int token = static_cast(blockIdx.y); + const int block_slot = warp; + const int blocks_per_row = min(tile_count, kHierarchicalScoreMaxBlocksPerRow); + const int first_tile = static_cast(blockIdx.x) % blocks_per_row; + const int token = static_cast(blockIdx.x) / blocks_per_row; const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength]; const int64_t abs_position = @@ -481,35 +543,54 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* const int64_t causal_count = SaiCausalThreshold(abs_position, kQwenCompressRatio); const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; - const int active_tiles = (block_count + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; - const int key_stride = tile_count * kHierarchicalTileBlocks; + + const int live_tiles = (block_count + kHierarchicalTopKTileBlocks - 1) / kHierarchicalTopKTileBlocks; + const int key_stride = tile_count * kHierarchicalTopKTileBlocks; const int64_t row_base = static_cast(token) * key_stride; - for (int index = active_tiles * kHierarchicalTileBlocks + blockIdx.x * blockDim.x + threadIdx.x; - index < key_stride; index += gridDim.x * blockDim.x) { + // Initialize only the fully padded suffix; scoring writes every key in live tiles. + for (int index = live_tiles * kHierarchicalTopKTileBlocks + first_tile * blockDim.x + threadIdx.x; + index < key_stride; index += blocks_per_row * blockDim.x) { tile_keys[row_base + index] = topk::kPaddingSortKey; } - - for (int tile = static_cast(blockIdx.x); tile < active_tiles; tile += static_cast(gridDim.x)) { - const int block_index = tile * kHierarchicalTileBlocks + block_slot; - float score = 0.0f; - if (block_index < block_count) { - const int d = lane * 4; - const int64_t query_base = - (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; - const int64_t key_base = - (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; - const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); - const float4 key_value = LoadQsaVector4(present_key_state + key_base + d); - score = query_value.x * key_value.x + query_value.y * key_value.y + - query_value.z * key_value.z + query_value.w * key_value.w; - score = topk::WarpReduceSum(score); - } - if (lane == 0) { - head_scores[block_slot][head] = fmaxf(score, 0.0f); + if (first_tile >= live_tiles) { + return; + } + const int d = lane * 4; + float4 query_values[kQwenNumHeads]; +#pragma unroll + for (int head = 0; head < kQwenNumHeads; ++head) { + const int64_t query_base = + (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; + query_values[head] = LoadQsaVector4(query_rotated + query_base + d); + } + for (int tile = first_tile; tile < live_tiles; tile += blocks_per_row) { + const int64_t output_base = row_base + tile * kHierarchicalTopKTileBlocks; + for (int slot = block_slot; slot < kHierarchicalTopKTileBlocks; + slot += kHierarchicalTileScoreThreads / kWarpSize) { + const int block_index = tile * kHierarchicalTopKTileBlocks + slot; + float4 key_value = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if (block_index < block_count) { + const int64_t key_base = + (static_cast(batch) * params.state_capacity + block_index) * kQwenHeadSize; + key_value = LoadQsaVector4(present_key_state + key_base + d); + } +#pragma unroll + for (int head = 0; head < kQwenNumHeads; ++head) { + float score = 0.0f; + if (block_index < block_count) { + const float4 query_value = query_values[head]; + score = query_value.x * key_value.x + query_value.y * key_value.y + + query_value.z * key_value.z + query_value.w * key_value.w; + score = topk::WarpReduceSum(score); + } + if (lane == 0) { + head_scores[slot][head] = fmaxf(score, 0.0f); + } + } } __syncthreads(); - if (threadIdx.x < kHierarchicalTileBlocks) { - const int tile_block = tile * kHierarchicalTileBlocks + static_cast(threadIdx.x); + if (threadIdx.x < kHierarchicalTopKTileBlocks) { + const int tile_block = tile * kHierarchicalTopKTileBlocks + static_cast(threadIdx.x); uint64_t key = topk::kPaddingSortKey; if (tile_block < block_count) { float total_score = 0.0f; @@ -523,20 +604,21 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* keys[threadIdx.x] = key; } __syncthreads(); - if (threadIdx.x == 0) { - for (int index = 1; index < kHierarchicalTileBlocks; ++index) { - const uint64_t key = keys[index]; - int insertion = index; - while (insertion > 0 && key > keys[insertion - 1]) { - keys[insertion] = keys[insertion - 1]; - --insertion; + if (warp == 0) { + uint64_t key = keys[lane]; +#pragma unroll + for (int width = 2; width <= kHierarchicalTopKTileBlocks; width *= 2) { +#pragma unroll + for (int stride = width / 2; stride > 0; stride /= 2) { + const uint64_t other = __shfl_xor_sync(0xffffffffU, key, stride); + const bool take_maximum = ((lane & width) == 0) == ((lane & stride) == 0); + key = take_maximum ? max(key, other) : min(key, other); } - keys[insertion] = key; } + keys[lane] = key; } __syncthreads(); - const int64_t output_base = row_base + tile * kHierarchicalTileBlocks; - if (threadIdx.x < kHierarchicalTileBlocks) { + if (threadIdx.x < kHierarchicalTopKTileBlocks) { tile_keys[output_base + threadIdx.x] = keys[threadIdx.x]; } __syncthreads(); @@ -719,11 +801,78 @@ __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* out const int right_count = min(min(list_width, kBoundedTopKMax), max(0, key_stride - right_start)); const int output_count = min(left_count + right_count, kBoundedTopKMax); const int64_t row_base = row * key_stride; + const uint64_t* left = input_keys + row_base + left_start; + const uint64_t* right = input_keys + row_base + right_start; + const bool left_empty = left_count == 0 || left[0] == topk::kPaddingSortKey; + const bool right_empty = right_count == 0 || right[0] == topk::kPaddingSortKey; + for (int rank = threadIdx.x; rank < output_count; rank += blockDim.x) { + uint64_t key = topk::kPaddingSortKey; + if (left_empty) { + key = rank < right_count ? right[rank] : topk::kPaddingSortKey; + } else if (right_empty) { + key = rank < left_count ? left[rank] : topk::kPaddingSortKey; + } else { + key = QsaMergeRank(left, left_count, right, right_count, rank); + } + output_keys[row_base + left_start + rank] = key; + } + } +} + +__global__ void QsaMergeTwoLevelsTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, + int key_stride, int list_width, int rows) { + __shared__ uint64_t inputs[4 * kBoundedTopKMax]; + __shared__ uint64_t intermediate[2 * kBoundedTopKMax]; + const int groups_per_row = (key_stride + 4 * list_width - 1) / (4 * list_width); + const int input_width = min(list_width, kBoundedTopKMax); + const int intermediate_width = min(2 * list_width, kBoundedTopKMax); + const int64_t total_groups = static_cast(rows) * groups_per_row; + for (int64_t work = blockIdx.x; work < total_groups; work += gridDim.x) { + const int start = static_cast(work % groups_per_row) * 4 * list_width; + const int64_t row_base = (work / groups_per_row) * key_stride; + for (int index = threadIdx.x; index < 4 * input_width; index += blockDim.x) { + const int source = start + (index / input_width) * list_width + index % input_width; + inputs[index] = source < key_stride ? input_keys[row_base + source] : topk::kPaddingSortKey; + } + __syncthreads(); + for (int index = threadIdx.x; index < 2 * intermediate_width; index += blockDim.x) { + const int half = index / intermediate_width; + const int rank = index % intermediate_width; + const int left_start = start + half * 2 * list_width; + const int right_start = left_start + list_width; + const int left_count = min(min(list_width, kBoundedTopKMax), max(0, key_stride - left_start)); + const int right_count = min(min(list_width, kBoundedTopKMax), max(0, key_stride - right_start)); + uint64_t key = topk::kPaddingSortKey; + if (rank < left_count + right_count) { + const uint64_t* left = inputs + half * 2 * input_width; + const uint64_t* right = left + input_width; + if (left_count == 0 || left[0] == topk::kPaddingSortKey) { + key = rank < right_count ? right[rank] : topk::kPaddingSortKey; + } else if (right_count == 0 || right[0] == topk::kPaddingSortKey) { + key = rank < left_count ? left[rank] : topk::kPaddingSortKey; + } else { + key = QsaMergeRank(left, left_count, right, right_count, rank); + } + } + intermediate[index] = key; + } + __syncthreads(); + const int left_count = min(intermediate_width, key_stride - start); + const int right_count = min(intermediate_width, max(0, key_stride - start - 2 * list_width)); + const int output_count = min(left_count + right_count, kBoundedTopKMax); + const uint64_t* right = intermediate + intermediate_width; for (int rank = threadIdx.x; rank < output_count; rank += blockDim.x) { - output_keys[row_base + left_start + rank] = - QsaMergeRank(input_keys + row_base + left_start, left_count, - input_keys + row_base + right_start, right_count, rank); + uint64_t key; + if (intermediate[0] == topk::kPaddingSortKey) { + key = rank < right_count ? right[rank] : topk::kPaddingSortKey; + } else if (right_count == 0 || right[0] == topk::kPaddingSortKey) { + key = rank < left_count ? intermediate[rank] : topk::kPaddingSortKey; + } else { + key = QsaMergeRank(intermediate, left_count, right, right_count, rank); + } + output_keys[row_base + start + rank] = key; } + __syncthreads(); } } @@ -881,8 +1030,18 @@ __global__ void QsaSelectKernel(const float* block_scores, const int32_t* topk_i } } +template +__device__ __forceinline__ uint64_t QsaSelectionKey(const float* scores, int index) { + if constexpr (PackedKeys) { + return reinterpret_cast(scores)[index]; + } else { + return topk::PackStableSortKey(scores[index], index); + } +} + // One block per packed query token. Radix-selects a bounded candidate set, then sorts only that // set. The output aliases the corresponding score row after all scores for that row have been read. +template __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, const int32_t* present_state_lengths, const int32_t* overflow_flags, @@ -916,7 +1075,10 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat if (selected == 0) { continue; } - float* scores = block_scores + static_cast(token) * params.state_capacity; + const int64_t score_stride = PackedKeys + ? 2 * ((static_cast(params.state_capacity) + kHierarchicalTopKTileBlocks - 1) / kHierarchicalTopKTileBlocks) * kHierarchicalTopKTileBlocks + : params.state_capacity; + float* scores = block_scores + static_cast(token) * score_stride; if (threadIdx.x == 0) { prefix = 0; @@ -931,7 +1093,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const uint32_t current_prefix = prefix; for (int index = threadIdx.x; index < block_count; index += blockDim.x) { - const uint32_t score_key = static_cast(topk::PackStableSortKey(scores[index], 0) >> 32); + const uint32_t score_key = static_cast(QsaSelectionKey(scores, index) >> 32); if (shift == 24 || (score_key >> (shift + 8)) == (current_prefix >> (shift + 8))) { atomicAdd(&temp.histogram[(score_key >> shift) & 0xffu], 1u); } @@ -946,11 +1108,23 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat } else { prefix |= static_cast(bucket) << shift; remaining = rank; + if constexpr (PackedKeys) { + // A fully selected radix bucket needs no finer threshold. + if (rank == count && prefix != 0) { + --prefix; + remaining = 0; + } + } break; } } } __syncthreads(); + if constexpr (PackedKeys) { + if (remaining == 0) { + break; + } + } } if (threadIdx.x == 0) { @@ -959,7 +1133,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const uint32_t threshold = prefix; for (int index = threadIdx.x; index < block_count; index += blockDim.x) { - const uint64_t key = topk::PackStableSortKey(scores[index], index); + const uint64_t key = QsaSelectionKey(scores, index); if (static_cast(key >> 32) > threshold) { const int slot = atomicAdd(&gathered, 1); selected_keys[slot] = key; @@ -971,10 +1145,11 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat equal_seen = 0; } __syncthreads(); + // Sorted 32-key tiles still encounter equal scores in ascending original-index order. for (int base = 0; base < block_count && equal_seen < remaining; base += blockDim.x) { const int index = base + static_cast(threadIdx.x); const bool equal = index < block_count && - static_cast(topk::PackStableSortKey(scores[index], 0) >> 32) == threshold; + static_cast(QsaSelectionKey(scores, index) >> 32) == threshold; const unsigned int warp_mask = __ballot_sync(0xffffffffu, equal); const int lane = static_cast(threadIdx.x) % kWarpSize; const int warp = static_cast(threadIdx.x) / kWarpSize; @@ -995,7 +1170,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const int equal_rank = equal_seen + equal_warp_offsets[warp] + lane_rank; if (equal && equal_rank < remaining) { - selected_keys[gathered + equal_rank] = topk::PackStableSortKey(scores[index], index); + selected_keys[gathered + equal_rank] = QsaSelectionKey(scores, index); } __syncthreads(); if (threadIdx.x == 0) { @@ -1016,7 +1191,11 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat for (int item = 0; item < kBoundedTopKItemsPerThread; ++item) { const int rank = threadIdx.x * kBoundedTopKItemsPerThread + item; if (rank < selected) { - topk_row[rank] = topk::UnpackStableSortIndex(keys[item]); + if constexpr (PackedKeys) { + reinterpret_cast(scores)[rank] = keys[item]; + } else { + topk_row[rank] = topk::UnpackStableSortIndex(keys[item]); + } } } __syncthreads(); @@ -1211,13 +1390,10 @@ __global__ void CsaUpdateStateKernel(const T* key, const T* gate, const T* key_n // (from the baseline copy in the Launch function below), which is exactly the prior valid // buffer this rejected step must preserve. if (!rejected) { - // Process slots in ascending order with a barrier between slots. When past and present alias, - // the destination slot precedes its source slot, so this is the device-wide equivalent of - // a forward memmove and prevents one thread from overwriting another thread's source. - for (int t = 0; t < present_buffer_length; ++t) { + for (int t = static_cast(threadIdx.x); t < present_buffer_length; t += static_cast(blockDim.x)) { const int virtual_pos = present_buffer_start + t; const int64_t out_base = (static_cast(b) * params.buffer_capacity + t) * width; - for (int c = static_cast(threadIdx.x); c < width; c += static_cast(blockDim.x)) { + for (int c = 0; c < width; ++c) { const float key_value = virtual_pos < old_buf_len ? to_float( @@ -1232,9 +1408,9 @@ __global__ void CsaUpdateStateKernel(const T* key, const T* gate, const T* key_n present_kv_buffer[out_base + c] = from_float(key_value); present_gate_buffer[out_base + c] = from_float(gate_value); } - __syncthreads(); } } + __syncthreads(); } } @@ -1406,7 +1582,7 @@ size_t GetQsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& const size_t rows = static_cast(params.total_tokens); const size_t query_count = rows * params.num_heads * params.head_size; if (UseHierarchicalQsaTopK(params)) { - const size_t key_count = rows * GetHierarchicalTileCount(params) * kHierarchicalTileBlocks; + const size_t key_count = rows * GetHierarchicalTileCount(params) * kHierarchicalTopKTileBlocks; return query_count + 4 * key_count; } return query_count + rows * static_cast(std::max(params.state_capacity, 1)); @@ -1418,21 +1594,50 @@ size_t GetCsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& return rows * params.num_heads * params.head_size + rows * static_cast(std::max(params.state_capacity, 1)); } -template +__global__ void MapQsaValidationStatusKernel( + const int32_t* cumulative_sequence_lengths, + const int32_t* overflow_flags, + int batch_size, + int32_t* validation_status) { + const int batch = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (batch >= batch_size) { + return; + } + const int request_length = cumulative_sequence_lengths[batch + 1] - cumulative_sequence_lengths[batch]; + validation_status[batch] = overflow_flags[batch] == 0 ? 0 : (request_length >= 2 ? 0x99 : 9); +} + +Status LaunchMapQsaValidationStatus( + cudaStream_t stream, + const int32_t* cumulative_sequence_lengths, + const int32_t* overflow_flags, + int batch_size, + int32_t* validation_status) { + if (batch_size <= 0) { + return Status::OK(); + } + constexpr int kStatusThreads = 128; + const int blocks = (batch_size + kStatusThreads - 1) / kStatusThreads; + MapQsaValidationStatusKernel<<>>( + cumulative_sequence_lengths, overflow_flags, batch_size, validation_status); + return CUDA_CALL(cudaGetLastError()); +} + +template Status LaunchQsaPackedSparseAttentionIndexer( cudaStream_t stream, const PackedSparseAttentionIndexerParams& params, const T* query, const T* key, const T* query_norm_weight, const T* key_norm_weight, const T* cos_cache, const T* sin_cache, const int32_t* cumulative_sequence_lengths, - const int32_t* past_sequence_lengths, const int64_t* position_ids, const T* past_key_state, + const int32_t* past_sequence_lengths, const int64_t* position_ids, const TState* past_key_state, const T* past_kv_buffer, const int32_t* past_state_lengths, const int32_t* state_update_capture_count, bool state_update_active, int32_t* selected_indices, - int32_t* selected_counts, T* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths, + int32_t* selected_counts, TState* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths, T* state_update, float* float_workspace, int32_t* overflow_flags) { if (params.batch_size > 0) { const int64_t key_state_elems = static_cast(params.batch_size) * params.state_capacity * params.head_size; if (key_state_elems > 0 && present_key_state != past_key_state) { - ElementwiseCopyKernel<<>>( + ElementwiseCopyKernel<<>>( past_key_state, present_key_state, key_state_elems); } const int64_t buffer_elems = static_cast(params.batch_size) * params.buffer_capacity * params.head_size; @@ -1452,11 +1657,10 @@ Status LaunchQsaPackedSparseAttentionIndexer( } const size_t value_bytes = static_cast(params.head_size) * sizeof(float); - const size_t update_shared_bytes = 2 * value_bytes + kThreads * sizeof(float); - ORT_RETURN_IF_ERROR( - ConfigureDynamicSharedMemory(QsaUpdateStateKernel, update_shared_bytes, "QsaUpdateStateKernel")); + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + QsaUpdateStateKernel, 2 * value_bytes + kThreads * sizeof(float), "QsaUpdateStateKernel")); const int state_blocks = static_cast(std::min(params.batch_size, kSaiMaxGridDimX)); - QsaUpdateStateKernel<<>>( + QsaUpdateStateKernel<<>>( key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, past_kv_buffer, past_state_lengths, state_update_capture_count, state_update_active, present_key_state, present_kv_buffer, present_state_lengths, state_update, overflow_flags, params); @@ -1474,63 +1678,70 @@ Status LaunchQsaPackedSparseAttentionIndexer( const int64_t rotate_rows = static_cast(params.total_tokens) * params.num_heads; const int rotate_blocks = static_cast(std::min(rotate_rows, kSaiMaxGridDimX)); - const size_t rotate_shared_bytes = value_bytes + kThreads * sizeof(float); ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( - PackedRotateQueryKernel, rotate_shared_bytes, "PackedRotateQueryKernel")); - PackedRotateQueryKernel<<>>( + PackedRotateQueryKernel, value_bytes + kThreads * sizeof(float), "PackedRotateQueryKernel")); + PackedRotateQueryKernel<<>>( query, query_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, position_ids, query_rotated, params); const int token_blocks = static_cast(std::min(params.total_tokens, kSaiMaxGridDimX)); if (use_hierarchical_topk) { const int tile_count = GetHierarchicalTileCount(params); - const int key_stride = tile_count * kHierarchicalTileBlocks; + const int key_stride = tile_count * kHierarchicalTopKTileBlocks; const size_t key_count = static_cast(params.total_tokens) * key_stride; uint64_t* merge_input = reinterpret_cast(selection_workspace); uint64_t* merge_output = merge_input + key_count; - if (params.batch_size == 1 && params.total_tokens >= 4) { - const dim3 score_grid(static_cast(std::min(tile_count, 128)), - static_cast((params.total_tokens + 1) / 2)); - QsaPairedScoreTileTopKKernel<<>>( - present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, merge_input, tile_count, params); + const int64_t tile_work = + static_cast(params.total_tokens) * std::min(tile_count, kHierarchicalScoreMaxBlocksPerRow); + QsaScoreTileTopKKernel<<(std::min(tile_work, kSaiMaxGridDimX)), + kHierarchicalTileScoreThreads, 0, stream>>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, merge_input, tile_count, params); + static const bool use_radix_topk = + ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", false); + if (use_radix_topk) { + QsaPartialTopKKernel<<>>( + selection_workspace, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, params); } else { - const dim3 score_grid(static_cast(std::min(tile_count, 128)), - static_cast(params.total_tokens)); - QsaScoreTileTopKKernel<<>>( - present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, merge_input, tile_count, params); - } - int compact_stride = key_stride; - int first_list_width = kHierarchicalTileBlocks; - if (params.total_tokens >= 4) { - const int chunks = (key_stride + 2047) / 2048; - compact_stride = chunks * kBoundedTopKMax; - QsaCompactTileTopKKernel<<(chunks), static_cast(params.total_tokens)), - kHierarchicalMergeThreads, 0, stream>>>( - merge_input, merge_output, key_stride, compact_stride); - std::swap(merge_input, merge_output); - first_list_width = kBoundedTopKMax; - } - for (int list_width = first_list_width; list_width < compact_stride; list_width *= 2) { - const int pairs_per_row = (compact_stride + 2 * list_width - 1) / (2 * list_width); - const int64_t merge_work = static_cast(params.total_tokens) * pairs_per_row; - QsaMergeTileTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), - kHierarchicalMergeThreads, 0, stream>>>( - merge_input, merge_output, compact_stride, list_width, params.total_tokens); - std::swap(merge_input, merge_output); + for (int list_width = kHierarchicalTopKTileBlocks; list_width < key_stride;) { + if (2 * list_width < key_stride) { + const int groups_per_row = (key_stride + 4 * list_width - 1) / (4 * list_width); + const int64_t merge_work = static_cast(params.total_tokens) * groups_per_row; + QsaMergeTwoLevelsTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), + kHierarchicalMergeThreads, 0, stream>>>( + merge_input, merge_output, key_stride, list_width, params.total_tokens); + list_width *= 4; + } else { + const int pairs_per_row = (key_stride + 2 * list_width - 1) / (2 * list_width); + const int64_t merge_work = static_cast(params.total_tokens) * pairs_per_row; + QsaMergeTileTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), + kHierarchicalMergeThreads, 0, stream>>>( + merge_input, merge_output, key_stride, list_width, params.total_tokens); + list_width *= 2; + } + std::swap(merge_input, merge_output); + } } QsaEmitHierarchicalTopKKernel<<>>( merge_input, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, - overflow_flags, compact_stride, selected_indices, selected_counts, params); + overflow_flags, key_stride, selected_indices, selected_counts, params); } else if (use_ragged_prefill) { - QsaRaggedPrefillScoreKernel<<>>( - present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, block_scores, params); + static const bool warp_heads = + ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_WARP_HEADS", false); + if (warp_heads) { + QsaRaggedPrefillScoreKernel<<>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, block_scores, params); + } else { + QsaRaggedPrefillScoreKernel<<>>( + present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, block_scores, params); + } } else if (params.state_capacity > 0) { const int64_t score_work = static_cast(params.total_tokens) * params.state_capacity; const int score_blocks = static_cast(std::min(score_work, kSaiMaxGridDimX)); - QsaBlockScoreKernel<<>>( + QsaBlockScoreKernel<<>>( present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, block_scores, params); } @@ -1539,7 +1750,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( !use_hierarchical_topk && params.block_topk > kSaiFastTopKMax && params.block_topk <= kBoundedTopKMax; int32_t* topk_indices = use_bounded_topk ? reinterpret_cast(block_scores) : nullptr; if (use_bounded_topk) { - QsaPartialTopKKernel<<>>( + QsaPartialTopKKernel<<>>( block_scores, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, overflow_flags, params); } @@ -1596,11 +1807,10 @@ Status LaunchCsaPackedSparseAttentionIndexer( } const size_t value_bytes = static_cast(params.head_size) * sizeof(float); - const size_t update_shared_bytes = value_bytes + kThreads * sizeof(float); - ORT_RETURN_IF_ERROR( - ConfigureDynamicSharedMemory(CsaUpdateStateKernel, update_shared_bytes, "CsaUpdateStateKernel")); const int state_blocks = static_cast(std::min(params.batch_size, kSaiMaxGridDimX)); - CsaUpdateStateKernel<<>>( + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + CsaUpdateStateKernel, value_bytes + kThreads * sizeof(float), "CsaUpdateStateKernel")); + CsaUpdateStateKernel<<>>( key, gate, key_norm_weight, cos_cache, sin_cache, position_bias, cumulative_sequence_lengths, past_sequence_lengths, past_kv_buffer, past_gate_buffer, past_state_lengths, present_key_state, present_kv_buffer, present_gate_buffer, present_state_lengths, overflow_flags, params); @@ -1614,10 +1824,9 @@ Status LaunchCsaPackedSparseAttentionIndexer( const int64_t rotate_rows = static_cast(params.total_tokens) * params.num_heads; const int rotate_blocks = static_cast(std::min(rotate_rows, kSaiMaxGridDimX)); - const size_t rotate_shared_bytes = value_bytes + kThreads * sizeof(float); ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( - PackedRotateQueryKernel, rotate_shared_bytes, "PackedRotateQueryKernel")); - PackedRotateQueryKernel<<>>( + PackedRotateQueryKernel, value_bytes + kThreads * sizeof(float), "PackedRotateQueryKernel")); + PackedRotateQueryKernel<<>>( query, query_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, position_ids, query_rotated, params); @@ -1637,21 +1846,29 @@ Status LaunchCsaPackedSparseAttentionIndexer( return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(T) \ - template Status LaunchQsaPackedSparseAttentionIndexer( \ - cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ - const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const T*, const T*, const int32_t*, \ - const int32_t*, bool, int32_t*, int32_t*, T*, T*, int32_t*, T*, float*, int32_t*); \ +#define INSTANTIATE_QSA(T, TState) \ + template Status LaunchQsaPackedSparseAttentionIndexer( \ + cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ + const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const TState*, const T*, \ + const int32_t*, const int32_t*, bool, int32_t*, int32_t*, TState*, T*, int32_t*, T*, float*, int32_t*); + +#define INSTANTIATE_CSA(T) \ template Status LaunchCsaPackedSparseAttentionIndexer( \ cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ const T*, const T*, const T*, const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const T*, \ const T*, const T*, const int32_t*, int32_t*, int32_t*, T*, T*, T*, int32_t*, float*, int32_t*); -INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(float) -INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(half) -INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(__nv_bfloat16) +INSTANTIATE_QSA(float, float) +INSTANTIATE_QSA(half, half) +INSTANTIATE_QSA(__nv_bfloat16, __nv_bfloat16) +INSTANTIATE_QSA(__nv_bfloat16, float) + +INSTANTIATE_CSA(float) +INSTANTIATE_CSA(half) +INSTANTIATE_CSA(__nv_bfloat16) -#undef INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER +#undef INSTANTIATE_QSA +#undef INSTANTIATE_CSA } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h index 81f15253ffd22..bce3f9474cd36 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h @@ -39,6 +39,8 @@ struct PackedSparseAttentionIndexerParams { // policy_mode = "qsa" int block_topk = 0; // token_budget / compress_ratio int state_update_capacity = 0; + void* mid_kv_buffer = nullptr; // [batch, buffer_capacity, head_size] T + int32_t* mid_state_lengths = nullptr; // [batch, 2] // policy_mode = "csa" int index_topk = 0; @@ -60,7 +62,7 @@ Status LaunchValidatePackedIndexerReuse(cudaStream_t stream, const int32_t* cumu // that request are left exactly as their past_* counterparts) and the select kernels force that // request's selected_indices/selected_counts to the safe empty result (-1 / 0) rather than // selecting against a partially updated state. -template +template Status LaunchQsaPackedSparseAttentionIndexer( cudaStream_t stream, const PackedSparseAttentionIndexerParams& params, @@ -73,20 +75,27 @@ Status LaunchQsaPackedSparseAttentionIndexer( const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, - const T* past_key_state, + const TState* past_key_state, const T* past_kv_buffer, const int32_t* past_state_lengths, const int32_t* state_update_capture_count, bool state_update_active, int32_t* selected_indices, int32_t* selected_counts, - T* present_key_state, + TState* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths, T* state_update, float* float_workspace, int32_t* overflow_flags); +Status LaunchMapQsaValidationStatus( + cudaStream_t stream, + const int32_t* cumulative_sequence_lengths, + const int32_t* overflow_flags, + int batch_size, + int32_t* validation_status); + template Status LaunchCsaPackedSparseAttentionIndexer( cudaStream_t stream, diff --git a/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc index ee4b4d299a40f..673581b577ca0 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc @@ -66,6 +66,7 @@ ONNX_OPERATOR_KERNEL_EX( kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) + .TypeConstraint("TCache", WebGpuSupportedFloatTypes()) .TypeConstraint("M", DataTypeImpl::GetTensorType()) .MayInplace(3, 1) .MayInplace(4, 2) @@ -555,6 +556,12 @@ Status DynamicSparseAttention::ComputeInternal(onnxruntime::webgpu::ComputeConte const Tensor* value = context.Input(2); const Tensor* past_key = context.Input(3); const Tensor* past_value = context.Input(4); + ORT_RETURN_IF_NOT((past_key == nullptr || past_key->GetElementType() == query->GetElementType()) && + (past_value == nullptr || past_value->GetElementType() == query->GetElementType()), + "DynamicSparseAttention: WebGPU cache types must match the query."); + ORT_RETURN_IF_NOT((context.InputCount() <= 18 || context.Input(18) == nullptr) && + (context.InputCount() <= 19 || context.Input(19) == nullptr), + "DynamicSparseAttention: INT8 cache scales are only supported by CUDA."); const Tensor* auxiliary_key = context.Input(5); const Tensor* auxiliary_value = context.Input(6); const Tensor* selected_indices = context.Input(7); diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc index accf0241c1d2e..e25c593a149df 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -389,6 +389,10 @@ Status PackedSparseAttentionIndexerQsaSelectProgram::GenerateShaderCode(ShaderHe const auto& present_key_state = shader.AddInput("present_key_state", ShaderUsage::UseUniform); const auto& rotary_cache = shader.AddInput("rotary_cache", ShaderUsage::UseUniform); const auto& sequence_metadata = shader.AddInput("sequence_metadata", ShaderUsage::UseUniform); + const ShaderVariableHelper* position_ids = nullptr; + if (has_position_ids_) { + position_ids = &shader.AddInput("position_ids", ShaderUsage::UseUniform); + } const auto& present_state_lengths = shader.AddInput("present_state_lengths", ShaderUsage::UseUniform); const auto& overflow_flags = shader.AddInput("overflow_flags", ShaderUsage::UseUniform); const auto& selected_indices = shader.AddOutput("selected_indices", ShaderUsage::UseUniform); @@ -411,12 +415,24 @@ Status PackedSparseAttentionIndexerQsaSelectProgram::GenerateShaderCode(ShaderHe << " let p = u32(position);\n" << " let cr = uniforms.compress_ratio;\n" << " return p / cr + select(0u, 1u, p % cr == cr - 1u);\n" - << "}\n" - << "fn abs_position(token: u32, b: u32) -> i32 {\n" - << " let req_start = " << sequence_metadata.GetByOffset("b") << ";\n" - << " return " << sequence_metadata.GetByOffset("uniforms.batch_size + 1u + b") - << " + (i32(token) - req_start);\n" - << "}\n" + << "}\n"; + if (has_position_ids_) { + shader.AdditionalImplementation() + << "fn abs_position(token: u32, b: u32) -> i32 {\n" + << " let raw = " << position_ids->GetByOffset("token", true) << ";\n" + << " if ((raw.y & 0x80000000u) != 0u) { return -2147483648; }\n" + << " if (raw.y != 0u || raw.x > 2147483647u) { return 2147483647; }\n" + << " return i32(raw.x);\n" + << "}\n"; + } else { + shader.AdditionalImplementation() + << "fn abs_position(token: u32, b: u32) -> i32 {\n" + << " let req_start = " << sequence_metadata.GetByOffset("b") << ";\n" + << " return " << sequence_metadata.GetByOffset("uniforms.batch_size + 1u + b") + << " + (i32(token) - req_start);\n" + << "}\n"; + } + shader.AdditionalImplementation() << "fn normalized_query_value(token: u32, head: u32, d: u32) -> f32 {\n" << " let base = token * uniforms.query_row_stride + head * uniforms.head_size;\n" << " var square_sum = 0.0;\n" @@ -904,9 +920,9 @@ Status PackedSparseAttentionIndexer::ComputeInternal(onnxruntime::webgpu::Comput } const bool position_ids_provided = psai::kPositionIds < context.InputCount() && context.Input(psai::kPositionIds) != nullptr; - ORT_RETURN_IF(position_ids_provided == is_qsa, + ORT_RETURN_IF(!is_qsa && !position_ids_provided, "PackedSparseAttentionIndexer: position_ids is required for " - "policy_mode 'csa' and must be omitted for policy_mode 'qsa'"); + "policy_mode 'csa'"); const bool gate_buffer_provided = psai::kPastGateBuffer < context.InputCount() && context.Input(psai::kPastGateBuffer) != nullptr; ORT_RETURN_IF(gate_buffer_provided != !is_qsa, "PackedSparseAttentionIndexer: past_gate_buffer ", @@ -934,6 +950,7 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont const Tensor* sin_cache = context.Input(psai::kSinCache); const Tensor* cu_seqlens = context.Input(psai::kCumulativeSequenceLengths); const Tensor* past_seqlens = context.Input(psai::kPastSequenceLengths); + const Tensor* position_ids = context.Input(psai::kPositionIds); const Tensor* past_key_state = context.Input(psai::kPastKeyState); const Tensor* past_kv_buffer = context.Input(psai::kPastKvBuffer); const Tensor* past_state_lengths = context.Input(psai::kPastStateLengths); @@ -976,6 +993,9 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont } ORT_RETURN_IF_ERROR(CheckShape(query_norm, "query_norm_weight", {head_size})); ORT_RETURN_IF_ERROR(CheckShape(norm, "key_norm_weight", {head_size})); + if (position_ids != nullptr) { + ORT_RETURN_IF_ERROR(CheckShape(position_ids, "position_ids", {total_tokens})); + } RotaryCacheShape rotary; ORT_RETURN_IF_ERROR(CheckRotaryCache(cos_cache, sin_cache, batch_size, rotary)); @@ -1138,14 +1158,17 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont ORT_RETURN_IF_ERROR(CheckStorageBindingLimit( context, "QSA selection", {&query_metadata, present_key_state, &rotary_cache, &sequence_metadata, present_state_lengths, &overflow_flags, - selected_indices, selected_counts})); - PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched}; - select.CacheHint(rotary.batched) + selected_indices, selected_counts, position_ids})); + PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched, position_ids != nullptr}; + select.CacheHint(rotary.batched, position_ids != nullptr) .SetWorkgroupSize(kWorkgroupSize) .AddInputs({{&query_metadata, ProgramTensorMetadataDependency::Type}, {present_key_state, ProgramTensorMetadataDependency::Type}, {&rotary_cache, ProgramTensorMetadataDependency::Type}, {&sequence_metadata, ProgramTensorMetadataDependency::Type}}); + if (position_ids != nullptr) { + select.AddInput({position_ids, ProgramTensorMetadataDependency::Type}); + } select.AddInputs({{present_state_lengths, ProgramTensorMetadataDependency::Type}, {&overflow_flags, ProgramTensorMetadataDependency::Type}}) .AddOutputs({{selected_indices, ProgramTensorMetadataDependency::Type}, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index fecac2f17f523..f4df1e5dd6fb5 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -483,7 +483,6 @@ void GroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); - const int64_t num_heads = getAttribute(ctx, "num_heads", 0); const int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0); if (num_heads <= 0 || num_heads > std::numeric_limits::max()) { @@ -524,7 +523,8 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex continue; } - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, output_index); + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput( + ctx, ctx.getInputType(input_index) != nullptr ? input_index : 0, output_index); if (hasInputShape(ctx, input_index)) { ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, input_index, output_index); continue; @@ -589,6 +589,12 @@ The main cache uses BNSH layout and ordinary contiguous append semantics. Select 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. 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. )DOC"; ONNX_MS_OPERATOR_SET_SCHEMA( @@ -629,8 +635,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(2, "value", "Current main value [batch, sequence, kv_num_heads * head_size].", "T", OpSchema::Optional) - .Input(3, "past_key", "Main key cache in BNSH layout.", "T", OpSchema::Optional) - .Input(4, "past_value", "Main value cache in BNSH layout.", "T", OpSchema::Optional) + .Input(3, "past_key", "Main key cache in BNSH layout, same type as query or INT8.", "TCache", OpSchema::Optional) + .Input(4, "past_value", "Main value cache in BNSH layout, same type as past_key.", "TCache", OpSchema::Optional) .Input(5, "auxiliary_key", "Read-only auxiliary key sequence in BNSH layout.", "T", OpSchema::Optional) .Input(6, "auxiliary_value", "Read-only auxiliary value sequence in BNSH layout.", "T", @@ -652,13 +658,34 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Input(16, "head_sink", "Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.", "T", OpSchema::Optional) + .Input(17, "committed_length", + "Optional single-request committed token count for cursor validation. When supplied, validation_status " + "is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent.", + "M", OpSchema::Optional) + .Input(18, "key_scale", "Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].", + "S", OpSchema::Optional) + .Input(19, "value_scale", "Positive finite per-channel INT8 value-cache scales [kv_num_heads * head_size].", + "S", OpSchema::Optional) .Output(0, "output", "Attention output [batch, sequence, num_heads * head_size].", "T") - .Output(1, "present_key", "Updated main key cache in BNSH layout.", "T", OpSchema::Optional) - .Output(2, "present_value", "Updated main value cache in BNSH layout.", "T", OpSchema::Optional) + .Output(1, "present_key", "Updated main key cache in BNSH layout.", "TCache", OpSchema::Optional) + .Output(2, "present_value", "Updated main value cache in BNSH layout.", "TCache", OpSchema::Optional) + .Output(3, "validation_status", "Optional int32 status: 0 when valid, a validation code when invalid.", + "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain all floating-point inputs and outputs to one element type.") + .TypeConstraint("TCache", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)", "tensor(int8)"}, + "Cache type must match T unless INT8 per-channel scales are supplied.") + .TypeConstraint("S", {"tensor(float)"}, "INT8 cache scales use float32.") .TypeConstraint("M", {"tensor(int32)"}, "Constrain selection and sequence metadata to int32.") - .TypeAndShapeInferenceFunction(DynamicSparseAttentionTypeAndShapeInference)); + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + DynamicSparseAttentionTypeAndShapeInference(ctx); + if (ctx.getNumOutputs() > 3 && ctx.hasOutput(3)) { + updateOutputElemType(ctx, 3, ONNX_NAMESPACE::TensorProto_DataType_INT32); + ONNX_NAMESPACE::TensorShapeProto status_shape; + status_shape.add_dim()->set_dim_value(1); + updateOutputShape(ctx, 3, status_shape); + } + })); constexpr const char* Attention_ver1_doc = R"DOC( Multi-Head Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT). @@ -2861,6 +2888,15 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference const bool has_capture_count = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateCaptureCount); const bool has_state_update_active = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateActive); const bool has_state_update_output = ctx.hasOutput(psai::kStateUpdate); + const bool has_mid_kv_buffer = ctx.hasOutput(psai::kMidKvBuffer); + const bool has_mid_state_lengths = ctx.hasOutput(psai::kMidStateLengths); + const bool has_validation_status = ctx.hasOutput(psai::kValidationStatus); + if (has_mid_kv_buffer != has_mid_state_lengths) { + fail_shape_inference("PackedSparseAttentionIndexer: mid_kv_buffer and mid_state_lengths must be connected together"); + } + if (!is_qsa && (has_mid_kv_buffer || has_validation_status)) { + fail_shape_inference("PackedSparseAttentionIndexer: mid-state and validation outputs are only supported for qsa"); + } if (!is_qsa && (state_update_capacity > 0 || has_capture_count || has_state_update_active || has_state_update_output)) { fail_shape_inference("PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'"); @@ -2872,9 +2908,9 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } if (ctx.getNumOutputs() < static_cast(psai::kFixedOutputCount) || - ctx.getNumOutputs() > static_cast(has_indexshare ? psai::kOutputCount : psai::kStateUpdate + 1)) { - fail_shape_inference("PackedSparseAttentionIndexer: expected ", psai::kFixedOutputCount, " or ", - psai::kOutputCount, " declared outputs, got ", ctx.getNumOutputs()); + ctx.getNumOutputs() > static_cast(psai::kMaxOutputCount)) { + fail_shape_inference("PackedSparseAttentionIndexer: expected ", psai::kFixedOutputCount, " to ", + psai::kMaxOutputCount, " declared outputs, got ", ctx.getNumOutputs()); } if (ctx.hasOutput(psai::kPresentGateBuffer) == is_qsa) { fail_shape_inference("PackedSparseAttentionIndexer: output ", psai::kPresentGateBuffer, @@ -2885,8 +2921,20 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference updateOutputElemType(ctx, psai::kSelectedCounts, ONNX_NAMESPACE::TensorProto_DataType_INT32); if (has_indexshare) updateOutputElemType(ctx, psai::kIndexShareStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); updateOutputElemType(ctx, psai::kPresentStateLengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); + if (has_mid_state_lengths) { + updateOutputElemType(ctx, psai::kMidStateLengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); + propagateShapeFromInputToOutput(ctx, psai::kPastStateLengths, psai::kMidStateLengths); + } + if (has_validation_status) { + updateOutputElemType(ctx, psai::kValidationStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); + propagateShapeFromInputToOutput(ctx, psai::kPastSequenceLengths, psai::kValidationStatus); + } propagateElemTypeFromInputToOutput(ctx, psai::kPastKeyState, psai::kPresentKeyState); propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kPresentKvBuffer); + if (has_mid_kv_buffer) { + propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); + propagateShapeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); + } if (!is_qsa) { // present_gate_buffer keeps its fixed positional slot (with an empty name) for policy_mode // 'qsa'; only propagate its type/shape when it is actually produced. @@ -3154,6 +3202,8 @@ Unlike SparseAttentionIndexer, this operator: rejected as a deterministic no-op on state rather than truncated or allowed to corrupt memory; * additionally emits selected_counts, the exact number of active (non -1) entries per query, so that no downstream consumer needs to scan selected_indices for its query's true count. + * for QSA, optionally emits the pending-state snapshot after the first request token and a + per-request validation status for overflow or inconsistent state metadata. Both policy_mode values keep the semantics of SparseAttentionIndexer, applied independently to each request's own packed token range and fixed-capacity state slice: @@ -3318,8 +3368,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "past_key_state", "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.", - "T") + "never changes across calls. For QSA, this state may use FLOAT while query/key tensors use " + "FLOAT16 or BFLOAT16.", + "T_KEY") .Input(13, "past_kv_buffer", "Generic fixed-capacity pending-token buffer. Shape is (batch_size, 2 * compress_ratio - 1, " @@ -3373,7 +3424,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Output(2, "present_key_state", "Updated generic key state, with the same fixed shape as past_key_state.", - "T") + "T_KEY") .Output(3, "present_kv_buffer", "Updated generic pending-token buffer, with the same fixed shape as past_kv_buffer.", @@ -3397,9 +3448,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Output(7, "indexshare_status", "Selection status [N], forwarded on reuse and zero on refresh.", "M", OpSchema::Optional) + .Output(8, "mid_kv_buffer", + "Only for policy_mode 'qsa': pending-buffer snapshot after the first token of each request.", + "T", OpSchema::Optional) + .Output(9, "mid_state_lengths", + "Only for policy_mode 'qsa': state lengths after the first token of each request.", + "M", OpSchema::Optional) + .Output(10, "validation_status", + "Only for policy_mode 'qsa': per-request overflow status; two-token rejection is encoded as 0x99.", + "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain floating point tensors to float, float16 and bfloat16.") + .TypeConstraint("T_KEY", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain the generic key-state tensor to float, float16 and bfloat16.") .TypeConstraint("I", {"tensor(int64)"}, "Constrain position ids to 64-bit integer tensors.") .TypeConstraint("M", {"tensor(int32)"}, "Constrain packed metadata, generic state lengths and selected indices/counts to 32-bit " diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 87188812c680e..c0dc15e633e82 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -88,6 +88,10 @@ struct DynamicSparseAttentionCase { std::vector q_norm_weight; std::vector k_norm_weight; std::vector head_sink; + std::vector key_scale; + std::vector value_scale; + std::optional committed_length; + std::optional expected_validation_status; std::vector expected_output; std::vector expected_present_key; @@ -108,7 +112,7 @@ std::vector ToTensorData(const std::vector& values) { } } -template +template void RunDynamicSparseAttentionCase( const DynamicSparseAttentionCase& c, std::unique_ptr execution_provider, @@ -163,12 +167,12 @@ void RunDynamicSparseAttentionCase( ToTensorData(c.value)); } if (c.has_past) { - tester.AddInput("past_key", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.past_key)); - tester.AddInput("past_value", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.past_value)); + tester.AddInput("past_key", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.past_key)); + tester.AddInput("past_value", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.past_value)); } else { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); @@ -229,15 +233,33 @@ void RunDynamicSparseAttentionCase( } else { tester.AddInput("head_sink", {c.num_heads}, ToTensorData(c.head_sink)); } + if (c.committed_length.has_value()) { + tester.AddInput("committed_length", {1}, {*c.committed_length}); + } else { + tester.AddOptionalInputEdge(); + } + if (!c.key_scale.empty() || !c.value_scale.empty()) { + if (c.key_scale.empty()) + tester.AddOptionalInputEdge(); + else + tester.AddInput("key_scale", {int64_t(c.key_scale.size())}, c.key_scale); + if (c.value_scale.empty()) + tester.AddOptionalInputEdge(); + else + tester.AddInput("value_scale", {int64_t(c.value_scale.size())}, c.value_scale); + } tester.AddOutput("output", {c.batch_size, c.sequence_length, hidden_size}, ToTensorData(c.expected_output)); - tester.AddOutput("present_key", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.expected_present_key)); - tester.AddOutput("present_value", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.expected_present_value)); + tester.AddOutput("present_key", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.expected_present_key)); + tester.AddOutput("present_value", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.expected_present_value)); + if (c.expected_validation_status.has_value()) { + tester.AddOutput("validation_status", {1}, {*c.expected_validation_status}); + } tester.SetOutputTolerance(0.005f); std::vector> execution_providers; @@ -276,6 +298,283 @@ DynamicSparseAttentionCase MakeSingleTokenSelectedValueCase(float value = 9.0f) } #ifdef USE_CUDA +DynamicSparseAttentionCase MakeInt8CacheCase() { + DynamicSparseAttentionCase c; + c.sequence_length = 2; + c.num_heads = 4; + c.kv_num_heads = 2; + c.cache_sequence_length = 4; + c.max_selected = 3; + c.query.assign(64, 0.0f); + c.key.resize(32); + c.value.resize(32); + c.key_scale.resize(16); + c.value_scale.resize(16); + c.past_key.resize(64); + c.past_value.resize(64); + for (size_t i = 0; i < 64; ++i) { + c.past_key[i] = float(int(i % 13) - 6); + c.past_value[i] = float(int(i % 17) - 8); + } + for (size_t i = 0; i < 16; ++i) { + c.key_scale[i] = i < 8 ? 0.5f : 0.25f; + c.value_scale[i] = i < 8 ? 0.25f : 0.5f; + } + for (size_t i = 0; i < 32; ++i) { + c.key[i] = i % 3 == 0 ? -300.0f : 0.25f * float(i); + c.value[i] = i % 3 == 0 ? 300.0f : 0.125f * float(i); + } + c.key[0] = -1e20f; + c.value[0] = 1e20f; + c.selected_indices = {0, 2, -1, 1, 2, 3}; + c.selected_counts = {2, 3}; + c.seqlens_k = {3}; + c.total_sequence_length = 4; + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + for (int s = 0; s < 2; ++s) { + for (int h = 0; h < 2; ++h) { + for (int d = 0; d < 8; ++d) { + const int channel = h * 8 + d; + const int src = s * 16 + channel; + const int dst = (h * 4 + 2 + s) * 8 + d; + c.expected_present_key[dst] = std::clamp(std::nearbyint(c.key[src] / c.key_scale[channel]), -128.0f, 127.0f); + c.expected_present_value[dst] = std::clamp(std::nearbyint(c.value[src] / c.value_scale[channel]), -128.0f, 127.0f); + } + } + } + c.expected_output.assign(64, 0.0f); + for (int s = 0; s < 2; ++s) { + for (int h = 0; h < 4; ++h) { + const int kv = h / 2; + for (int d = 0; d < 8; ++d) { + float sum = 0.0f; + for (int i = 0; i < c.selected_counts[s]; ++i) { + const int pos = c.selected_indices[s * 3 + i]; + sum += c.expected_present_value[(kv * 4 + pos) * 8 + d] * c.value_scale[kv * 8 + d]; + } + c.expected_output[(s * 4 + h) * 8 + d] = sum / c.selected_counts[s]; + } + } + } + return c; +} + +TEST(DynamicSparseAttentionTest, Int8CachePerChannelGqaAppendAndAttention_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase(MakeInt8CacheCase(), std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheBFloat16_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase(MakeInt8CacheCase(), std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheSplitAttention_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + auto c = MakeInt8CacheCase(); + c.cache_sequence_length = 1024; + c.total_sequence_length = 1024; + c.seqlens_k = {1023}; + c.max_selected = 512; + c.selected_counts = {512, 512}; + c.selected_indices.resize(1024); + std::iota(c.selected_indices.begin(), c.selected_indices.end(), 0); + c.past_key.assign(2 * 1024 * 8, 0.0f); + c.past_value.assign(2 * 1024 * 8, 3.0f); + auto key = c.past_key; + auto value = c.past_value; + for (int h = 0; h < 2; ++h) { + for (int s = 0; s < 2; ++s) { + for (int d = 0; d < 8; ++d) { + key[(h * 1024 + 1022 + s) * 8 + d] = c.expected_present_key[(h * 4 + 2 + s) * 8 + d]; + value[(h * 1024 + 1022 + s) * 8 + d] = c.expected_present_value[(h * 4 + 2 + s) * 8 + d]; + } + } + } + c.expected_present_key = key; + c.expected_present_value = value; + for (int s = 0; s < 2; ++s) { + for (int h = 0; h < 4; ++h) { + for (int d = 0; d < 8; ++d) { + float sum = 0.0f; + for (int i = 0; i < 512; ++i) { + sum += value[((h / 2) * 1024 + s * 512 + i) * 8 + d] * c.value_scale[(h / 2) * 8 + d]; + } + c.expected_output[(s * 4 + h) * 8 + d] = sum / 512.0f; + } + } + } + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheRejectsMissingOrMalformedScales_CUDA) { + for (int variant = 0; variant < 3; ++variant) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + auto c = MakeInt8CacheCase(); + if (variant == 0) { + c.key_scale.clear(); + c.value_scale.clear(); + } else if (variant == 1) { + c.value_scale.clear(); + } else { + c.key_scale.resize(8); + } + const char* error = variant == 0 ? "INT8 caches require per-channel" + : variant == 1 ? "must be provided together" + : "INT8 cache scales must have shape"; + RunDynamicSparseAttentionCase(c, std::move(cuda_ep), + OpTester::ExpectResult::kExpectFailure, error); + } +} + +DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool auxiliary) { + DynamicSparseAttentionCase c; + c.batch_size = 2; + c.sequence_length = sequence_length; + c.num_heads = 20; + c.kv_num_heads = 2; + c.head_size = 256; + c.cache_sequence_length = 520; + c.max_selected = 512; + c.attention_mode = auxiliary ? "local_plus_selected" : "selected_only"; + c.local_window_size = auxiliary ? 16 : -1; + c.smooth_softmax = 1; + c.scale = 0.0625f; + c.head_sink.assign(c.num_heads, 0.25f); + c.seqlens_k = {519, 518}; + c.total_sequence_length = 520; + const int channels = static_cast(c.kv_num_heads * c.head_size); + const int rows = static_cast(c.batch_size * c.sequence_length); + c.query.resize(rows * c.num_heads * c.head_size); + c.key.resize(rows * channels); + c.value.resize(rows * channels); + c.key_scale.resize(channels); + c.value_scale.resize(channels); + for (int d = 0; d < channels; ++d) { + c.key_scale[d] = d % 2 ? 0.125f : 0.0625f; + c.value_scale[d] = d % 2 ? 0.25f : 0.125f; + } + for (size_t i = 0; i < c.query.size(); ++i) { + c.query[i] = float(int(i % 7) - 3) * 0.125f; + } + for (size_t i = 0; i < c.key.size(); ++i) { + c.key[i] = float(int(i % 11) - 5) * c.key_scale[i % channels]; + c.value[i] = float(int(i % 13) - 6) * c.value_scale[i % channels]; + } + c.past_key.resize(c.batch_size * c.cache_sequence_length * channels); + c.past_value.resize(c.past_key.size()); + for (size_t i = 0; i < c.past_key.size(); ++i) { + c.past_key[i] = float(int(i % 17) - 8); + c.past_value[i] = float(int(i % 19) - 9); + } + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + for (int row = 0; row < rows; ++row) { + const int b = row / sequence_length; + const int pos = c.seqlens_k[b] + 1 - sequence_length + row % sequence_length; + for (int d = 0; d < channels; ++d) { + const int64_t dst = + ((b * c.kv_num_heads + d / c.head_size) * c.cache_sequence_length + pos) * c.head_size + + d % c.head_size; + c.expected_present_key[dst] = c.key[row * channels + d] / c.key_scale[d]; + c.expected_present_value[dst] = c.value[row * channels + d] / c.value_scale[d]; + } + } + if (auxiliary) { + c.selected_kv_source = "auxiliary"; + c.auxiliary_sequence_length = 512; + c.auxiliary_key.resize(c.batch_size * c.auxiliary_sequence_length * channels); + c.auxiliary_value.resize(c.auxiliary_key.size()); + for (size_t i = 0; i < c.auxiliary_key.size(); ++i) { + c.auxiliary_key[i] = float(int(i % 17) - 8) * 0.125f; + c.auxiliary_value[i] = float(int(i % 19) - 9) * 0.25f; + } + } + c.selected_counts.resize(rows); + c.selected_indices.assign(rows * c.max_selected, -1); + c.expected_output.assign(c.query.size(), 0.0f); + for (int row = 0; row < rows; ++row) { + const int b = row / sequence_length; + const int pos = c.seqlens_k[b] + 1 - sequence_length + row % sequence_length; + c.selected_counts[row] = row == 0 ? 0 : 509; + std::iota(c.selected_indices.begin() + row * c.max_selected, + c.selected_indices.begin() + row * c.max_selected + c.selected_counts[row], 0); + for (int h = 0; h < c.num_heads; ++h) { + const int kv = h / (c.num_heads / c.kv_num_heads); + const int64_t out = (row * c.num_heads + h) * c.head_size; + std::vector numerator(c.head_size, 0.0f); + double denominator = std::exp(double(c.head_sink[h])); + const int local_count = auxiliary ? static_cast(c.local_window_size) : 0; + const int local_start = pos - local_count + 1; + for (int i = 0; i < local_count + c.selected_counts[row]; ++i) { + const bool selected = i >= local_count; + const int index = selected ? i - local_count : local_start + i; + if (selected && !auxiliary && index >= local_start) continue; + const bool from_auxiliary = selected && auxiliary; + const int64_t length = from_auxiliary ? c.auxiliary_sequence_length : c.cache_sequence_length; + const int64_t offset = ((b * c.kv_num_heads + kv) * length + index) * c.head_size; + double dot = 0.0; + for (int d = 0; d < c.head_size; ++d) { + const float key = from_auxiliary ? c.auxiliary_key[offset + d] + : c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d]; + dot += c.query[out + d] * key; + } + const double weight = std::exp(dot * c.scale); + denominator += weight; + for (int d = 0; d < c.head_size; ++d) { + const float value = from_auxiliary ? c.auxiliary_value[offset + d] + : c.expected_present_value[offset + d] * c.value_scale[kv * c.head_size + d]; + numerator[d] += static_cast(weight * value); + } + } + for (int d = 0; d < c.head_size; ++d) { + c.expected_output[out + d] = static_cast(numerator[d] / denominator); + } + } + } + return c; +} + +TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecode_CUDA) { + for (int sequence_length : {1, 2}) { + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + + ", auxiliary=" + std::to_string(auxiliary)); + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase( + MakeGroupedInt8CacheCase(sequence_length, auxiliary), std::move(cuda_ep)); + } + } +} + +TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeFloat16_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase( + MakeGroupedInt8CacheCase(2, false), std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheRejectsNonpositiveOrNonfiniteScales_CUDA) { + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DYNAMIC_SPARSE_ATTENTION_STRICT_VALIDATION", "1"}}}; + for (float invalid : {0.0f, -1.0f, std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + auto c = MakeInt8CacheCase(); + c.value_scale[15] = invalid; + RunDynamicSparseAttentionCase( + c, std::move(cuda_ep), OpTester::ExpectResult::kExpectFailure, + "INT8 cache scales must be positive and finite"); + } +} + Status ResolveDynamicSparseAttentionGraph(size_t output_count, int64_t query_width = 8, int64_t num_heads = 1, int64_t kv_num_heads = 1, bool packed_qkv = false, const char* query_width_symbol = nullptr, @@ -450,6 +749,32 @@ TEST(DynamicSparseAttentionTest, SelectedOnlyMainVariableCountsGqaAndCacheAppend RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); } +TEST(DynamicSparseAttentionTest, ValidationStatusReportsCommittedCursor_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available."; + } + + auto c = MakeSingleTokenSelectedOnlyCase(); + c.cache_sequence_length = 2; + c.total_sequence_length = 2; + c.past_key.assign(16, 0.0f); + c.past_value.assign(16, 0.0f); + c.key.assign(8, 0.0f); + c.value.assign(8, 0.0f); + c.expected_present_key.assign(16, 0.0f); + c.expected_present_value.assign(16, 0.0f); + c.committed_length = 0; + c.expected_validation_status = 0; + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + + cuda_ep = DefaultCudaExecutionProvider(); + ASSERT_NE(cuda_ep, nullptr); + c.committed_length = 1; + c.expected_validation_status = 9; + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); +} + TEST(DynamicSparseAttentionTest, LocalPlusSelectedAuxiliaryJointSoftmaxSinkSharedKv_CUDA) { auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) { diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 3ccf50716367c..2b605676a7bb9 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -947,7 +947,7 @@ QsaPackedProblem MakeQsaPackedProblem(QsaPackedProblem problem = {}) { return problem; } -template +template void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedProblem(), ProviderKind provider_kind = ProviderKind::Cuda, QsaPackedResult* actual = nullptr, bool packed_qk = false, const ConfigOptions* webgpu_options = nullptr, int padding = 0) { @@ -963,7 +963,7 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP problem.key_norm_weight = RoundTrip(problem.key_norm_weight); problem.cos_cache = RoundTrip(problem.cos_cache); problem.sin_cache = RoundTrip(problem.sin_cache); - problem.past_key_state = RoundTrip(problem.past_key_state); + problem.past_key_state = RoundTrip(problem.past_key_state); problem.past_kv_buffer = RoundTrip(problem.past_kv_buffer); QsaPackedResult expected; @@ -1024,8 +1024,8 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP } else { test.AddInput("position_ids", {total_tokens}, problem.position_ids); } - test.AddInput("past_key_state", {batch_size, problem.state_capacity, head_size}, - ToElementType(problem.past_key_state)); + test.AddInput("past_key_state", {batch_size, problem.state_capacity, head_size}, + ToElementType(problem.past_key_state)); test.AddInput("past_kv_buffer", {batch_size, buffer_capacity, head_size}, ToElementType(problem.past_kv_buffer)); test.AddOptionalInputEdge(); // past_gate_buffer @@ -1048,8 +1048,8 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP } test.AddOutput("selected_indices", {total_tokens, output_capacity}, expected.selected_indices); test.AddOutput("selected_counts", {total_tokens}, expected.selected_counts); - test.AddOutput("present_key_state", {batch_size, problem.state_capacity, head_size}, - ToElementType(expected.present_key_state), false, 0.0f, tolerance); + test.AddOutput("present_key_state", {batch_size, problem.state_capacity, head_size}, + ToElementType(expected.present_key_state), false, 0.0f, tolerance); test.AddOutput("present_kv_buffer", {batch_size, buffer_capacity, head_size}, ToElementType(expected.present_kv_buffer), false, 0.0f, tolerance); test.AddOptionalOutputEdge(); // present_gate_buffer @@ -1065,7 +1065,7 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP RunOnProvider(test, std::move(provider)); if (actual != nullptr) { const auto& fetches = test.GetFetches(); - actual->present_key_state = CopyTensorToFloat(fetches[2].Get()); + actual->present_key_state = CopyTensorToFloat(fetches[2].Get()); actual->present_kv_buffer = CopyTensorToFloat(fetches[3].Get()); actual->present_state_lengths.assign(fetches[4].Get().DataAsSpan().begin(), fetches[4].Get().DataAsSpan().end()); @@ -1080,6 +1080,10 @@ TEST(PackedSparseAttentionIndexerTest, QsaFloat16) { RunQsaPackedTest TEST(PackedSparseAttentionIndexerTest, QsaBFloat16) { RunQsaPackedTest(2.0e-2f); } +TEST(PackedSparseAttentionIndexerTest, QsaBFloat16FloatState) { + RunQsaPackedTest(2.0e-2f); +} + TEST(PackedSparseAttentionIndexerTest, QsaPackedQueryKey) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::Cuda, nullptr, true); } @@ -1341,6 +1345,119 @@ TEST(PackedSparseAttentionIndexerTest, QsaMalformedStateLengthsAreClamped) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(std::move(problem))); } +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalTileBoundariesPreserveStableTies) { + for (int live_blocks : {0, 31, 32, 33, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049, 2051}) { + SCOPED_TRACE(live_blocks); + QsaPackedProblem problem; + problem.batch_size = 2; + problem.cumulative_sequence_lengths = {0, 1, 2}; + problem.past_sequence_lengths = {live_blocks * 4, std::max(0, live_blocks - 1) * 4}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = 2051; + problem.max_position = std::max(32, live_blocks * 4 + 2); + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 0.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + RunQsaPackedTest(2.0e-3f, problem); + RunQsaPackedTest(2.0e-2f, std::move(problem)); + } +} + +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalPackedRowsWithLargeCapacity) { + for (int tokens : {1, 2, 3, 8, 64}) { + SCOPED_TRACE(tokens); + QsaPackedProblem problem; + problem.batch_size = 1; + problem.cumulative_sequence_lengths = {0, tokens}; + problem.past_sequence_lengths = {4097 * 4 + 3}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = 65536; + problem.max_position = problem.past_sequence_lengths[0] + tokens + 1; + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 0.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + RunQsaPackedTest(2.0e-3f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-2f, std::move(problem), ProviderKind::Cuda); + } +} + +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalMergeLevelsWithPartialGroups) { + for (int capacity : {513, 1025, 2051, 4097}) { + SCOPED_TRACE(capacity); + QsaPackedProblem problem; + problem.batch_size = 2; + problem.cumulative_sequence_lengths = {0, 2, 3}; + problem.past_sequence_lengths = {(capacity - 1) * 4, 33 * 4}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = capacity; + problem.max_position = capacity * 4 + 4; + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 1.0f); + std::fill(problem.query_norm_weight.begin(), problem.query_norm_weight.end(), 1.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + for (int batch = 0; batch < problem.batch_size; ++batch) { + for (int block = 0; block < capacity; ++block) { + const float value = static_cast((block * 73) % capacity + 1) / capacity; + std::fill_n(problem.past_key_state.begin() + + (static_cast(batch) * capacity + block) * problem.head_size, + problem.head_size, value); + } + } + RunQsaPackedTest(2.0e-3f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-2f, std::move(problem), ProviderKind::Cuda); + } +} + +TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalScoringReusesKeysAcrossHeads) { + for (int live_blocks : {4095, 4096, 4097, 8191}) { + SCOPED_TRACE(live_blocks); + QsaPackedProblem problem; + problem.batch_size = 2; + problem.cumulative_sequence_lengths = {0, 2, 3}; + problem.past_sequence_lengths = {live_blocks * 4, 33 * 4}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = 8193; + problem.max_position = live_blocks * 4 + 4; + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 1.0f); + std::fill(problem.query_norm_weight.begin(), problem.query_norm_weight.end(), 1.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + for (int batch = 0; batch < problem.batch_size; ++batch) { + for (int block = 0; block < problem.state_capacity; ++block) { + const float value = static_cast((block * 73) % problem.state_capacity + 1) / + problem.state_capacity; + std::fill_n(problem.past_key_state.begin() + + (static_cast(batch) * problem.state_capacity + block) * problem.head_size, + problem.head_size, value); + } + } + RunQsaPackedTest(1.0e-5f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-3f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-2f, problem, ProviderKind::Cuda); + RunQsaPackedTest(2.0e-2f, std::move(problem), ProviderKind::Cuda); + } +} + #ifdef USE_WEBGPU TEST(PackedSparseAttentionIndexerWebGpuTest, QsaFloat) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::WebGpu); From 7be8d68244b362c898113c5bb13c3da5ee7545a4 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:28:21 +0000 Subject: [PATCH 02/23] Align QSA position overrides and CUDA-only outputs across providers --- .../packed_sparse_attention_indexer.md | 31 ++++++++++++++----- .../bert/packed_sparse_attention_indexer.cc | 10 ++++++ .../bert/packed_sparse_attention_indexer.h | 6 ++-- .../core/graph/contrib_ops/bert_defs.cc | 7 ++--- ...packed_sparse_attention_indexer_op_test.cc | 24 +++++++++++--- 5 files changed, 60 insertions(+), 18 deletions(-) diff --git a/docs/contrib_ops/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md index bd6877cf32189..9c589f98a3a63 100644 --- a/docs/contrib_ops/packed_sparse_attention_indexer.md +++ b/docs/contrib_ops/packed_sparse_attention_indexer.md @@ -71,8 +71,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 | optional qsa / required csa | +| 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,11 +86,21 @@ 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 | `mid_kv_buffer` | same shape as `past_kv_buffer` | T | optional CUDA qsa snapshot | +| 9 | `mid_state_lengths` | same shape as `past_state_lengths` | int32 | optional CUDA qsa snapshot | +| 10 | `validation_status` | `(batch_size)` | int32 | optional CUDA qsa status | + +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. +`mid_kv_buffer` and `mid_state_lengths` must be connected together and capture the state +after each request's first token. `validation_status` is zero on success; rejected +requests report 9, or `0x99` for requests containing at least two tokens. `capacity` is `token_budget + compress_ratio - 1` for `qsa` and `index_topk` for `csa`, exactly the same formula (`SelectedCapacity`) used by `SparseAttentionIndexer`. @@ -196,9 +206,9 @@ 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"`. Optional QSA `position_ids` override query rotary positions and causal +visibility. Callers must keep these positions consistent with prepared block keys and the cache +indices consumed by attention. ## 6. Policy `csa` @@ -216,7 +226,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`; `qsa` optionally accepts an explicit position override). 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`. @@ -281,6 +291,13 @@ 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. +- CUDA has two experimental, default-off dispatch controls. Set them before the first eligible + call; each value is cached per process. `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1` uses live-key + radix selection for eligible short QSA calls. `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1` uses + warp-local scoring for eligible packed prefill calls with more than 64 tokens, four heads, + head size 128 and compression ratio four. Both retain causal visibility, stable tie ordering, + selected counts and state-update semantics. Broad full-model quality and performance + qualification remains outstanding; neither option is a general tuning recommendation. - 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/webgpu/bert/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc index e25c593a149df..8d514a180af15 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -36,6 +36,7 @@ ONNX_OPERATOR_KERNEL_EX( kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) + .TypeConstraint("T_KEY", WebGpuSupportedFloatTypes()) .TypeConstraint("I", DataTypeImpl::GetTensorType()) .TypeConstraint("M", DataTypeImpl::GetTensorType()) .MayInplace(psai::kPastKeyState, psai::kPresentKeyState) @@ -911,6 +912,15 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo& i Status PackedSparseAttentionIndexer::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { ORT_RETURN_IF(context.InputCount() > psai::kIndexShareMode && context.Input(psai::kIndexShareMode), "PackedSparseAttentionIndexer IndexShare requires the CUDA provider"); + for (int index : {psai::kMidKvBuffer, psai::kMidStateLengths, psai::kValidationStatus}) { + Tensor* output = context.OutputCount() > index ? context.Output(index, TensorShape({0})) : nullptr; + ORT_RETURN_IF(output != nullptr, + "PackedSparseAttentionIndexer: mid-state and validation outputs require the CUDA provider"); + } + const Tensor* query = context.Input(psai::kQuery); + const Tensor* key_state = context.Input(psai::kPastKeyState); + ORT_RETURN_IF(query != nullptr && key_state != nullptr && query->DataType() != key_state->DataType(), + "PackedSparseAttentionIndexer: mixed input/state types require the CUDA provider"); const bool is_qsa = policy_ == psai::Policy::kQsa; constexpr int kCsaOnlyInputs[] = {psai::kGate, psai::kPositionBias, psai::kHeadWeights}; for (int index : kCsaOnlyInputs) { diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h index 69f590b73a2bd..353cc0413a813 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h @@ -65,9 +65,10 @@ class PackedSparseAttentionIndexerQsaUpdateProgram final class PackedSparseAttentionIndexerQsaSelectProgram final : public Program { public: - explicit PackedSparseAttentionIndexerQsaSelectProgram(bool cos_cache_batched) + PackedSparseAttentionIndexerQsaSelectProgram(bool cos_cache_batched, bool has_position_ids) : Program{"PackedSparseAttentionIndexerQsaSelect"}, - cos_cache_batched_{cos_cache_batched} {} + cos_cache_batched_{cos_cache_batched}, + has_position_ids_{has_position_ids} {} Status GenerateShaderCode(ShaderHelper& shader) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"total_tokens", ProgramUniformVariableDataType::Uint32}, @@ -87,6 +88,7 @@ class PackedSparseAttentionIndexerQsaSelectProgram final private: bool cos_cache_batched_; + bool has_position_ids_; }; // One invocation per request: closes every new compression window (softmax-gated pool -> RMSNorm diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index f4df1e5dd6fb5..eb1cc02c454c4 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2880,10 +2880,9 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference : " is required when policy_mode is 'csa'"); } } - if (PackedSparseAttentionIndexerHasInput(ctx, psai::kPositionIds) == is_qsa) { + if (!is_qsa && !PackedSparseAttentionIndexerHasInput(ctx, psai::kPositionIds)) { fail_shape_inference("PackedSparseAttentionIndexer: input ", psai::kPositionIds, - is_qsa ? " (position_ids) must be omitted when policy_mode is 'qsa'" - : " (position_ids) is required when policy_mode is 'csa'"); + " (position_ids) is required when policy_mode is 'csa'"); } const bool has_capture_count = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateCaptureCount); const bool has_state_update_active = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateActive); @@ -3218,7 +3217,7 @@ request's own packed token range and fixed-capacity state slice: 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. + past_sequence_lengths and cumulative_sequence_lengths unless optional position_ids are supplied. policy_mode = "csa" ("compressed sparse attention" block indexer) Applies the same window-plan arithmetic as SparseAttentionIndexer (overlap/leftover/new window diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 2b605676a7bb9..63a0f5bead6e7 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -124,7 +124,7 @@ struct GraphOptions { int64_t BufferCapacity(int64_t compress_ratio) { return 2 * compress_ratio - 1; } // Builds a fixed-input node; csa-only slots are left empty for policy_mode "qsa", as the schema -// requires. position_ids (slot 11) is omitted for "qsa" and forced on for "csa". +// requires. position_ids (slot 11) is optional for "qsa" and forced on for "csa". void AddNode(ModelTestBuilder& builder, const GraphOptions& options) { const bool is_csa = options.policy_mode == psai::kPolicyModeCsa; const int64_t width = is_csa ? 2 * options.head_size : options.head_size; @@ -477,11 +477,13 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsQsaWithCsaInput) { "must be omitted when policy_mode is 'qsa'"); } -TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsQsaWithPositionIds) { +TEST(PackedSparseAttentionIndexerShapeInferenceTest, QsaAcceptsPositionIds) { GraphOptions options; options.add_position_ids = true; - ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, - "must be omitted when policy_mode is 'qsa'"); + std::unique_ptr model; + ASSERT_STATUS_OK(BuildAndResolve([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, model)); + const auto& node = *model->MainGraph().Nodes().begin(); + EXPECT_TRUE(node.InputDefs()[psai::kPositionIds]->Exists()); } TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsCsaMissingPositionIds) { @@ -534,7 +536,7 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsWrongOutputCount) { GraphOptions options; options.output_count = 4; ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, - "output size 4 not in range [min=6, max=8]"); + "output size 4 not in range [min=6, max=11]"); } TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsStateUpdateCapacityWithoutCaptureCount) { @@ -1345,6 +1347,18 @@ TEST(PackedSparseAttentionIndexerTest, QsaMalformedStateLengthsAreClamped) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(std::move(problem))); } +TEST(PackedSparseAttentionIndexerTest, QsaExplicitPositionIds) { + auto problem = MakeQsaPackedProblem(); + problem.position_ids.resize(static_cast(problem.TotalTokens())); + for (int batch = 0; batch < problem.batch_size; ++batch) { + const int start = problem.cumulative_sequence_lengths[batch]; + for (int token = start; token < problem.cumulative_sequence_lengths[batch + 1]; ++token) { + problem.position_ids[static_cast(token)] = problem.past_sequence_lengths[batch] + token - start + 1; + } + } + RunQsaPackedTest(1.0e-5f, std::move(problem), ProviderKind::Cuda); +} + TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalTileBoundariesPreserveStableTies) { for (int live_blocks : {0, 31, 32, 33, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049, 2051}) { SCOPED_TRACE(live_blocks); From b538d1fc287b9b502cf6371a07eedc359a28e083 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:28:45 +0000 Subject: [PATCH 03/23] Update generated QSA position override documentation --- docs/ContribOperators.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 70be4dc1cba47..f7391f2bd1d24 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5139,7 +5139,7 @@ This version of the operator has been available since version 1 of the 'com.micr 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. + past_sequence_lengths and cumulative_sequence_lengths unless optional position_ids are supplied. policy_mode = "csa" ("compressed sparse attention" block indexer) Applies the same window-plan arithmetic as SparseAttentionIndexer (overlap/leftover/new window From 538bd9827bb2800aadb2d53b445524e464decfa8 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 21:03:22 +0000 Subject: [PATCH 04/23] Fix sparse attention workspace and optional interface handling --- docs/ContribOperators.md | 85 +++++++++++-------- docs/OperatorKernels.md | 4 +- .../bert/dynamic_sparse_attention_impl.cu | 6 +- .../sparse/packed_sparse_attention_indexer.cc | 5 ++ .../webgpu/bert/dynamic_sparse_attention.cc | 5 ++ .../bert/packed_sparse_attention_indexer.cc | 9 +- .../core/graph/contrib_ops/bert_defs.cc | 15 ++-- .../dynamic_sparse_attention_op_test.cc | 78 ++++++++++++++--- ...packed_sparse_attention_indexer_op_test.cc | 56 ++++++++++-- 9 files changed, 197 insertions(+), 66 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index f7391f2bd1d24..438478ccb4099 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1722,12 +1722,12 @@ 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. - - CUDA also supports INT8 main caches with float32 per-channel key/value scales. Appended rows are - quantized with round-to-nearest-even and saturation to [-128, 127]; attention reads the quantized - representation, including current rows. Auxiliary caches retain the query's floating-point type. - Scales must be positive and finite; requested validation reports invalid scales as status 10, - and CUDA strict validation raises an error. WebGPU supports only query-typed floating caches. + + Main caches may instead use INT8 with positive finite float32 key_scale and value_scale vectors, one scale per + KV-head channel. 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 @@ -1736,35 +1736,35 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
attention_mode : string
+
attention_mode : string (default is selected_only)
One of 'selected_only' or 'local_plus_selected'.
-
auxiliary_kv_shared : int
+
auxiliary_kv_shared : int (default is 0)
Use auxiliary_key as both key and value when auxiliary_value is omitted.
-
do_rotary : int
+
do_rotary : int (default is 0)
Whether to apply rotary embedding to Q and newly appended K.
-
is_causal : int
+
is_causal : int (default is 1)
Must be 1. DynamicSparseAttention version 1 supports causal attention only.
-
kv_num_heads : int (required)
+
kv_num_heads : int
Number of main and auxiliary KV heads.
-
local_window_size : int
+
local_window_size : int (default is -1)
Number of causally visible main-cache entries in local_plus_selected mode.
-
num_heads : int (required)
+
num_heads : int
Number of query heads.
-
qk_norm_epsilon : float
+
qk_norm_epsilon : float (default is 0.0)
Epsilon for optional per-head Q/K RMS normalization.
-
rotary_interleaved : int
+
rotary_interleaved : int (default is 0)
Whether rotary pairs use interleaved layout.
-
rotary_offset : int
+
rotary_offset : int (default is 0)
First head channel covered by rotary embedding. Must be nonnegative, a multiple of 8, and leave room for the rotary dimensions within head_size.
scale : float
Scaling factor applied to QK. Defaults to 1/sqrt(head_size).
-
selected_kv_source : string
+
selected_kv_source : string (default is main)
Source addressed by selected indices: 'main' or 'auxiliary'.
-
smooth_softmax : int
+
smooth_softmax : int (default is 0)
Add a zero-valued sink logit when no explicit head_sink is supplied.
-#### Inputs (11 - 20) +#### Inputs (5 - 20)
query : T
@@ -1774,9 +1774,9 @@ This version of the operator has been available since version 1 of the 'com.micr
value (optional) : T
Current main value [batch, sequence, kv_num_heads * head_size].
past_key (optional) : TCache
-
Main key cache in BNSH layout, query-typed or INT8.
+
Main key cache in BNSH layout, same type as query or INT8.
past_value (optional) : TCache
-
Main value cache in BNSH layout, matching past_key.
+
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
@@ -1802,11 +1802,11 @@ This version of the operator has been available since version 1 of the 'com.micr
head_sink (optional) : T
Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.
committed_length (optional) : M
-
Optional single-request committed cursor. Inconsistency with sequence metadata produces validation status 9.
+
Optional single-request committed token count for cursor validation. When supplied, validation_status is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent. 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]. Both scales are required for INT8 caches.
+
Positive finite per-channel INT8 value-cache scales [kv_num_heads * head_size].
#### Outputs (1 - 4) @@ -1819,7 +1819,7 @@ This version of the operator has been available since version 1 of the 'com.micr
present_value (optional) : TCache
Updated main value cache in BNSH layout.
validation_status (optional) : M
-
One int32 validation status: zero for valid inputs; 9 for cursor mismatch; 10 for invalid INT8 scales.
+
Optional CUDA-only int32 status: 0 when valid, a validation code when invalid.
#### Type Constraints @@ -1828,7 +1828,7 @@ 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 matches T unless INT8 per-channel scales are supplied.
+
Cache type must match T unless INT8 per-channel scales are supplied.
S : tensor(float)
INT8 cache scales use float32.
M : tensor(int32)
@@ -5125,6 +5125,8 @@ This version of the operator has been available since version 1 of the 'com.micr rejected as a deterministic no-op on state rather than truncated or allowed to corrupt memory; * additionally emits selected_counts, the exact number of active (non -1) entries per query, so that no downstream consumer needs to scan selected_indices for its query's true count. + * for QSA, optionally emits the pending-state snapshot after the first request token and a + per-request validation status for overflow or inconsistent state metadata. Both policy_mode values keep the semantics of SparseAttentionIndexer, applied independently to each request's own packed token range and fixed-capacity state slice: @@ -5182,31 +5184,31 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
compress_ratio : int (required)
+
compress_ratio : int
Number of consecutive tokens folded into one compressed/pooled entry. Must be > 0 and 2 * compress_ratio - 1 must not exceed INT_MAX.
-
epsilon : float
+
epsilon : float (default is 0.0)
Epsilon of the RMS normalization applied to queries and compressed keys. Default is 1e-6.
head_weight_scale : float
Only for policy_mode 'csa': scale applied to head_weights. Default is 1/sqrt(num_heads). Must be omitted when policy_mode is 'qsa'.
index_topk : int
Only for policy_mode 'csa': number of compressed entries selected per query. Must be > 0. Must be omitted when policy_mode is 'qsa'.
-
indexshare_state_policy : string
+
indexshare_state_policy : string (default is default)
Reuse state policy: default holds QSA and advances CSA; hold or advance overrides it.
-
max_output_entries : int
+
max_output_entries : int (default is 0)
Fixed padded output width for the optional IndexShare interface.
-
policy_mode : string (required)
+
policy_mode : string
Indexer policy. Must be exactly 'qsa' (token indexer) or 'csa' (compressed block indexer).
scale : float
Scale applied to the per-head ReLU scores. Default is 1/sqrt(head_size).
-
state_capacity : int (required)
+
state_capacity : int
Fixed capacity (number of entries) of past_key_state / present_key_state. Must be > 0.
-
state_update_capacity : int
+
state_update_capacity : int (default is 0)
Only for policy_mode 'qsa': maximum number of leading token transitions captured per request. Must be in [0, 8]. Default is 0.
token_budget : int
Only for policy_mode 'qsa': maximum number of tokens selected from complete blocks. Must be > 0 and divisible by compress_ratio. Must be omitted when policy_mode is 'csa'.
-#### Inputs (16 - 22) +#### Inputs (10 - 22)
query : T
@@ -5233,8 +5235,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. For QSA, this state may use FLOAT while query/key tensors use FLOAT16 or BFLOAT16.
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
@@ -5255,14 +5257,14 @@ This version of the operator has been available since version 1 of the 'com.micr
Merged device status [N].
-#### Outputs (6 - 8) +#### Outputs (5 - 11)
selected_indices : M
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.
@@ -5274,6 +5276,12 @@ This version of the operator has been available since version 1 of the 'com.micr
Only for policy_mode 'qsa': compact transition payloads with shape (batch_size, state_update_capacity, head_size). A token that completes a compression block stores the prepared block representative; any other captured token stores its raw key. Inactive and unused slots are zero.
indexshare_status (optional) : M
Selection status [N], forwarded on reuse and zero on refresh.
+
mid_kv_buffer (optional) : T
+
Only for CUDA policy_mode 'qsa' refresh: pending-buffer snapshot after the first token of each request. Not supported in IndexShare reuse mode.
+
mid_state_lengths (optional) : M
+
Only for CUDA policy_mode 'qsa' refresh: state lengths after the first token of each request. Not supported in IndexShare reuse mode.
+
validation_status (optional) : M
+
Only for CUDA policy_mode 'qsa' refresh: per-request overflow status; two-token rejection is encoded as 0x99. Not supported in IndexShare reuse mode.
#### Type Constraints @@ -5281,6 +5289,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)
@@ -8789,3 +8799,4 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
+ diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 669e3083f6b15..0aec82e268565 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1098,7 +1098,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)| @@ -1136,7 +1136,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**|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* mid_kv_buffer:**T**
*out* mid_state_lengths:**M**
*out* validation_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/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index ccdfdf76bb7ee..6bd20ba933cb2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -1750,11 +1750,13 @@ size_t GetDynamicSparseAttentionWorkspaceSize( const size_t query_blocks = static_cast(parameters.batch_size) * parameters.sequence_length * parameters.num_heads; + const bool grouped_decode = + UseGroupedDecodeAttention(parameters, element_size, max_shared_memory_per_block); const size_t split_count = - UseGroupedDecodeAttention(parameters, element_size, max_shared_memory_per_block) + grouped_decode ? GetGroupedDecodeSplitCount(parameters) : GetAttentionSplitCount(parameters); - if (split_count == 1) { + if (split_count == 1 && !grouped_decode) { return 0; } return SafeInt(query_blocks) * split_count * diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index eaedf62a4ffca..30cec6e80742c 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -646,6 +646,11 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* cont template Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { + for (int index : {psai::kMidKvBuffer, psai::kMidStateLengths, psai::kValidationStatus}) { + Tensor* output = context->OutputCount() > index ? context->Output(index, TensorShape({0})) : nullptr; + ORT_RETURN_IF(output != nullptr, + "PackedSparseAttentionIndexer: outputs 8-10 are not supported in IndexShare reuse mode"); + } const Tensor* query = context->Input(psai::kQuery); const Tensor* lengths = context->Input(psai::kPastSequenceLengths); const Tensor* indices = context->Input(psai::kMergedIndices); diff --git a/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc index 673581b577ca0..3d796a25fe71c 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/dynamic_sparse_attention.cc @@ -551,6 +551,11 @@ DynamicSparseAttention::DynamicSparseAttention(const OpKernelInfo& info) : WebGp } Status DynamicSparseAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + ORT_RETURN_IF(context.InputCount() > 17 && context.Input(17) != nullptr, + "DynamicSparseAttention: committed_length is only supported by CUDA."); + Tensor* validation_status = context.OutputCount() > 3 ? context.Output(3, TensorShape({0})) : nullptr; + ORT_RETURN_IF(validation_status != nullptr, + "DynamicSparseAttention: validation_status is only supported by CUDA."); const Tensor* query = context.Input(0); const Tensor* key = context.Input(1); const Tensor* value = context.Input(2); diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc index 8d514a180af15..3695eb7487c29 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -70,13 +70,17 @@ Status CheckShape(const Tensor* tensor, const char* name, std::initializer_list< Status CheckStorageBindingLimit(const onnxruntime::webgpu::ComputeContext& context, const char* program_name, - std::initializer_list tensors) { + std::initializer_list tensors, + const Tensor* optional_tensor = nullptr) { const auto& limits = context.DeviceLimits(); uint32_t binding_count = 0; for (const Tensor* tensor : tensors) { ORT_RETURN_IF(tensor == nullptr, "PackedSparseAttentionIndexer: missing tensor for ", program_name); binding_count += StorageBindingSegments(tensor, limits.maxStorageBufferBindingSize); } + if (optional_tensor != nullptr) { + binding_count += StorageBindingSegments(optional_tensor, limits.maxStorageBufferBindingSize); + } if (binding_count > limits.maxStorageBuffersPerShaderStage) { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "PackedSparseAttentionIndexer (WebGPU): ", program_name, " requires ", binding_count, @@ -1168,7 +1172,8 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont ORT_RETURN_IF_ERROR(CheckStorageBindingLimit( context, "QSA selection", {&query_metadata, present_key_state, &rotary_cache, &sequence_metadata, present_state_lengths, &overflow_flags, - selected_indices, selected_counts, position_ids})); + selected_indices, selected_counts}, + position_ids)); PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched, position_ids != nullptr}; select.CacheHint(rotary.batched, position_ids != nullptr) .SetWorkgroupSize(kWorkgroupSize) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index eb1cc02c454c4..1f2bb69cb1053 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -660,7 +660,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "T", OpSchema::Optional) .Input(17, "committed_length", "Optional single-request committed token count for cursor validation. When supplied, validation_status " - "is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent.", + "is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent. CUDA only.", "M", OpSchema::Optional) .Input(18, "key_scale", "Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].", "S", OpSchema::Optional) @@ -669,7 +669,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Output(0, "output", "Attention output [batch, sequence, num_heads * head_size].", "T") .Output(1, "present_key", "Updated main key cache in BNSH layout.", "TCache", OpSchema::Optional) .Output(2, "present_value", "Updated main value cache in BNSH layout.", "TCache", OpSchema::Optional) - .Output(3, "validation_status", "Optional int32 status: 0 when valid, a validation code when invalid.", + .Output(3, "validation_status", "Optional CUDA-only int32 status: 0 when valid, a validation code when invalid.", "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain all floating-point inputs and outputs to one element type.") @@ -2940,7 +2940,7 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference propagateElemTypeFromInputToOutput(ctx, psai::kPastGateBuffer, psai::kPresentGateBuffer); } if (has_state_update_output) { - propagateElemTypeFromInputToOutput(ctx, psai::kPastKeyState, psai::kStateUpdate); + propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kStateUpdate); } const auto* query_shape = PackedSparseAttentionIndexerShape(ctx, psai::kQuery, 2); @@ -3448,13 +3448,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Output(7, "indexshare_status", "Selection status [N], forwarded on reuse and zero on refresh.", "M", OpSchema::Optional) .Output(8, "mid_kv_buffer", - "Only for policy_mode 'qsa': pending-buffer snapshot after the first token of each request.", + "Only for CUDA policy_mode 'qsa' refresh: pending-buffer snapshot after the first token of each request. " + "Not supported in IndexShare reuse mode.", "T", OpSchema::Optional) .Output(9, "mid_state_lengths", - "Only for policy_mode 'qsa': state lengths after the first token of each request.", + "Only for CUDA policy_mode 'qsa' refresh: state lengths after the first token of each request. " + "Not supported in IndexShare reuse mode.", "M", OpSchema::Optional) .Output(10, "validation_status", - "Only for policy_mode 'qsa': per-request overflow status; two-token rejection is encoded as 0x99.", + "Only for CUDA policy_mode 'qsa' refresh: per-request overflow status; two-token rejection is encoded as 0x99. " + "Not supported in IndexShare reuse mode.", "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index c0dc15e633e82..5a3ad2f1dd8a1 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -106,7 +106,7 @@ std::vector ToTensorData(const std::vector& values) { std::vector converted; converted.reserve(values.size()); for (float value : values) { - converted.emplace_back(value); + converted.emplace_back(static_cast(value)); } return converted; } @@ -432,12 +432,14 @@ TEST(DynamicSparseAttentionTest, Int8CacheRejectsMissingOrMalformedScales_CUDA) } } -DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool auxiliary) { +DynamicSparseAttentionCase MakeGroupedCacheCase(int sequence_length, bool auxiliary, + bool floating_cache = false, int batch_size = 2, + int num_heads = 20, int kv_num_heads = 2) { DynamicSparseAttentionCase c; - c.batch_size = 2; + c.batch_size = batch_size; c.sequence_length = sequence_length; - c.num_heads = 20; - c.kv_num_heads = 2; + c.num_heads = num_heads; + c.kv_num_heads = kv_num_heads; c.head_size = 256; c.cache_sequence_length = 520; c.max_selected = 512; @@ -446,7 +448,8 @@ DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool au c.smooth_softmax = 1; c.scale = 0.0625f; c.head_sink.assign(c.num_heads, 0.25f); - c.seqlens_k = {519, 518}; + c.seqlens_k.resize(batch_size); + for (int b = 0; b < batch_size; ++b) c.seqlens_k[b] = 519 - b % 2; c.total_sequence_length = 520; const int channels = static_cast(c.kv_num_heads * c.head_size); const int rows = static_cast(c.batch_size * c.sequence_length); @@ -456,8 +459,8 @@ DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool au c.key_scale.resize(channels); c.value_scale.resize(channels); for (int d = 0; d < channels; ++d) { - c.key_scale[d] = d % 2 ? 0.125f : 0.0625f; - c.value_scale[d] = d % 2 ? 0.25f : 0.125f; + c.key_scale[d] = floating_cache ? 1.0f : (d % 2 ? 0.125f : 0.0625f); + c.value_scale[d] = floating_cache ? 1.0f : (d % 2 ? 0.25f : 0.125f); } for (size_t i = 0; i < c.query.size(); ++i) { c.query[i] = float(int(i % 7) - 3) * 0.125f; @@ -505,7 +508,7 @@ DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool au std::iota(c.selected_indices.begin() + row * c.max_selected, c.selected_indices.begin() + row * c.max_selected + c.selected_counts[row], 0); for (int h = 0; h < c.num_heads; ++h) { - const int kv = h / (c.num_heads / c.kv_num_heads); + const int64_t kv = h / (c.num_heads / c.kv_num_heads); const int64_t out = (row * c.num_heads + h) * c.head_size; std::vector numerator(c.head_size, 0.0f); double denominator = std::exp(double(c.head_sink[h])); @@ -537,9 +540,46 @@ DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool au } } } + if (floating_cache) { + c.key_scale.clear(); + c.value_scale.clear(); + } return c; } +TEST(DynamicSparseAttentionTest, FloatingCacheGroupedReferenceParity_CUDA) { + for (int sequence_length : {1, 3}) { + for (bool auxiliary : {false, true}) { + for (bool bf16 : {false, true}) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP() << "CUDA EP not available."; + SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + + ", auxiliary=" + std::to_string(auxiliary) + ", bf16=" + std::to_string(bf16)); + auto c = MakeGroupedCacheCase(sequence_length, auxiliary, true); + // Ten heads per KV head exercise partial groups of both six (prefill) and eight (decode). + if (bf16) { + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + } else { + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + } + } + } + } +} + +TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) { + for (bool bf16 : {false, true}) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP() << "CUDA EP not available."; + auto c = MakeGroupedCacheCase(1, false, true, 128, 4, 1); + if (bf16) { + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + } else { + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + } + } +} + TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecode_CUDA) { for (int sequence_length : {1, 2}) { for (bool auxiliary : {false, true}) { @@ -548,7 +588,7 @@ TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecode_CUDA) { auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) GTEST_SKIP(); RunDynamicSparseAttentionCase( - MakeGroupedInt8CacheCase(sequence_length, auxiliary), std::move(cuda_ep)); + MakeGroupedCacheCase(sequence_length, auxiliary), std::move(cuda_ep)); } } } @@ -557,7 +597,7 @@ TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeFloat16_CUDA) { auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) GTEST_SKIP(); RunDynamicSparseAttentionCase( - MakeGroupedInt8CacheCase(2, false), std::move(cuda_ep)); + MakeGroupedCacheCase(2, false), std::move(cuda_ep)); } TEST(DynamicSparseAttentionTest, Int8CacheRejectsNonpositiveOrNonfiniteScales_CUDA) { @@ -1894,6 +1934,22 @@ TEST(DynamicSparseAttentionTest, RejectsUnknownModeAndSource_CUDA) { #ifdef USE_WEBGPU +TEST(DynamicSparseAttentionTest, RejectsCudaValidationOptions_WebGPU) { + for (bool committed : {false, true}) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) GTEST_SKIP() << "WebGPU EP not available."; + auto c = MakeSingleTokenSelectedOnlyCase(); + if (committed) { + c.committed_length = 0; + } else { + c.expected_validation_status = 0; + } + RunDynamicSparseAttentionCase( + c, std::move(webgpu_ep), OpTester::ExpectResult::kExpectFailure, + committed ? "committed_length is only supported by CUDA" : "validation_status is only supported by CUDA"); + } +} + TEST(DynamicSparseAttentionTest, TokenDecodeWithAliasedCache_WebGPU) { auto webgpu_ep = DefaultWebGpuExecutionProvider(); if (!webgpu_ep) { diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 63a0f5bead6e7..03a17be5c73c3 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -599,10 +599,12 @@ std::unique_ptr CreateProvider(ProviderKind provider_kind, #endif } -void RunOnProvider(OpTester& test, std::unique_ptr provider) { +void RunOnProvider(OpTester& test, std::unique_ptr provider, + OpTester::ExpectResult expected_result = OpTester::ExpectResult::kExpectSuccess, + const std::string& expected_error = "") { std::vector> providers; providers.push_back(std::move(provider)); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + test.Run(expected_result, expected_error, {}, nullptr, &providers); } std::vector MakeWave(size_t count, float phase, float step) { @@ -952,7 +954,8 @@ QsaPackedProblem MakeQsaPackedProblem(QsaPackedProblem problem = {}) { template void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedProblem(), ProviderKind provider_kind = ProviderKind::Cuda, QsaPackedResult* actual = nullptr, - bool packed_qk = false, const ConfigOptions* webgpu_options = nullptr, int padding = 0) { + bool packed_qk = false, const ConfigOptions* webgpu_options = nullptr, int padding = 0, + const std::string& expected_error = "") { auto provider = CreateProvider(provider_kind, webgpu_options); if (provider == nullptr) { GTEST_SKIP() << (provider_kind == ProviderKind::Cuda ? "CUDA" : "WebGPU") @@ -1064,7 +1067,9 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP if (problem.state_update_capacity == 0) test.AddOptionalOutputEdge(); test.AddOutput("status", {total_tokens}, std::vector(static_cast(total_tokens), 0)); } - RunOnProvider(test, std::move(provider)); + RunOnProvider(test, std::move(provider), + expected_error.empty() ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + expected_error); if (actual != nullptr) { const auto& fetches = test.GetFetches(); actual->present_key_state = CopyTensorToFloat(fetches[2].Get()); @@ -1095,7 +1100,7 @@ TEST(PackedSparseAttentionIndexerTest, IndexShareRefreshPreservesBudgetAndPadsRo } TEST(PackedSparseAttentionIndexerTest, IndexShareReusePackedQkAndStatus) { - for (int invalid = 0; invalid < 7; ++invalid) { + for (int invalid = 0; invalid < 9; ++invalid) { SCOPED_TRACE(invalid); std::vector cumulative{0, 1, 2}; std::vector past_lengths{5, 5}; @@ -1143,7 +1148,17 @@ TEST(PackedSparseAttentionIndexerTest, IndexShareReusePackedQkAndStatus) { test.AddOutput("present_lengths", {2, 2}, state_lengths); test.AddOptionalOutputEdge(); test.AddOutput("status", {2}, {invalid == 0 ? 0 : 1, 2}); - RunOnProvider(test, std::move(provider)); + if (invalid == 7) { + test.AddOutput("mid_buffer", {2, 3, 2}, std::vector(12)); + test.AddOutput("mid_lengths", {2, 2}, std::vector(4)); + } else if (invalid == 8) { + test.AddOptionalOutputEdge(); + test.AddOptionalOutputEdge(); + test.AddOutput("validation", {2}, {0, 0}); + } + RunOnProvider(test, std::move(provider), + invalid >= 7 ? OpTester::ExpectResult::kExpectFailure : OpTester::ExpectResult::kExpectSuccess, + invalid >= 7 ? "outputs 8-10 are not supported in IndexShare reuse mode" : ""); } } @@ -1202,6 +1217,17 @@ TEST(PackedSparseAttentionIndexerTest, QsaStateUpdateCapturesRawAndBlockTransiti RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(std::move(problem)), ProviderKind::Cuda); } +TEST(PackedSparseAttentionIndexerTest, QsaBFloat16FloatStateCapture) { + QsaPackedProblem problem; + problem.batch_size = 2; + problem.compress_ratio = 4; + problem.cumulative_sequence_lengths = {0, 8, 13}; + problem.past_sequence_lengths = {0, 3}; + problem.state_update_capacity = 8; + problem.state_update_capture_count = {8, 3}; + RunQsaPackedTest(2.0e-2f, MakeQsaPackedProblem(std::move(problem)), ProviderKind::Cuda); +} + TEST(PackedSparseAttentionIndexerTest, QsaStateUpdateActiveZeroClearsAllSlots) { QsaPackedProblem problem; problem.batch_size = 2; @@ -1514,6 +1540,24 @@ TEST(PackedSparseAttentionIndexerWebGpuTest, QsaPortableStorageBindingLimit) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::WebGpu, nullptr, false, &config_options); } +TEST(PackedSparseAttentionIndexerWebGpuTest, QsaExplicitPositionsStorageBindingLimit) { + auto problem = MakeQsaPackedProblem(); + problem.position_ids.resize(static_cast(problem.TotalTokens())); + for (int b = 0; b < problem.batch_size; ++b) { + const int start = problem.cumulative_sequence_lengths[b]; + const int end = problem.cumulative_sequence_lengths[b + 1]; + for (int token = start; token < end; ++token) { + problem.position_ids[token] = problem.past_sequence_lengths[b] + token - start; + } + } + for (const std::string limit : {"8", "9"}) { + ConfigOptions config_options; + ASSERT_STATUS_OK(config_options.AddConfigEntry(webgpu::options::kMaxStorageBuffersPerShaderStage, limit)); + RunQsaPackedTest(1.0e-5f, problem, ProviderKind::WebGpu, nullptr, false, &config_options, 0, + limit == "8" ? "QSA selection requires 9 storage-buffer bindings" : ""); + } +} + TEST(PackedSparseAttentionIndexerWebGpuTest, QsaAllEmptyBatchPreservesState) { QsaPackedProblem problem; problem.batch_size = 1; From 9356555892c9779438cc08d533486daaf84c2317 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 22:10:04 +0000 Subject: [PATCH 05/23] Fix WebGPU binding-limit test and clarify mixed key-state support --- docs/ContribOperators.md | 2 +- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 4 ++-- .../contrib_ops/packed_sparse_attention_indexer_op_test.cc | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 438478ccb4099..d1e34d9a039e9 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5236,7 +5236,7 @@ This version of the operator has been available since version 1 of the 'com.micr
position_ids (optional) : I
Only for policy_mode 'csa': absolute position of every packed query, with shape (total_tokens).
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. For QSA, this state may use FLOAT while query/key tensors use FLOAT16 or BFLOAT16.
+
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
diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 1f2bb69cb1053..7d14d4eeab284 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -3367,8 +3367,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "past_key_state", "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. For QSA, this state may use FLOAT while query/key tensors use " - "FLOAT16 or BFLOAT16.", + "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.", "T_KEY") .Input(13, "past_kv_buffer", diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 03a17be5c73c3..5db1dd84bc2c8 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -1552,7 +1552,7 @@ TEST(PackedSparseAttentionIndexerWebGpuTest, QsaExplicitPositionsStorageBindingL } for (const std::string limit : {"8", "9"}) { ConfigOptions config_options; - ASSERT_STATUS_OK(config_options.AddConfigEntry(webgpu::options::kMaxStorageBuffersPerShaderStage, limit)); + ASSERT_STATUS_OK(config_options.AddConfigEntry(webgpu::options::kMaxStorageBuffersPerShaderStage, limit.c_str())); RunQsaPackedTest(1.0e-5f, problem, ProviderKind::WebGpu, nullptr, false, &config_options, 0, limit == "8" ? "QSA selection requires 9 storage-buffer bindings" : ""); } From d21946e6cb37621b9a2e5f623ea56e635b3e7bdf Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 21 Oct 2026 22:13:53 +0000 Subject: [PATCH 06/23] Add compact QSA selection and specialized INT8 attention dispatch Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../packed_sparse_attention_indexer.md | 28 ++++ .../bert/dynamic_sparse_attention_impl.cu | 142 +++++++++++++++--- .../packed_sparse_attention_indexer_impl.cu | 137 +++++++++++++---- .../contrib_ops/cuda/sparse/qsa_topk.cuh | 62 ++++++++ .../dynamic_sparse_attention_op_test.cc | 44 +++++- ...packed_sparse_attention_indexer_op_test.cc | 22 +++ 6 files changed, 377 insertions(+), 58 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh diff --git a/docs/contrib_ops/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md index 9c589f98a3a63..953eb7204efa5 100644 --- a/docs/contrib_ops/packed_sparse_attention_indexer.md +++ b/docs/contrib_ops/packed_sparse_attention_indexer.md @@ -1,5 +1,33 @@ # PackedSparseAttentionIndexer — Operator Documentation +## Experimental CUDA decode specializations + +Set these process-level options before the first operator invocation. All three +default to `0`; only `0` and `1` are accepted. Existing dispatch remains available +when an option is disabled or its geometry is unsupported. + +- `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1`: compact QSA chunks of 2048 keys + to 512 keys with device-side 512/1024/2048 sorting probes. Up to five chunks + use a fused final merge/emit; larger selections use the existing pairwise + merge. Stable tie ordering, causal tails, capacity stride/padding, and + overflow handling are preserved. This option is mutually exclusive with + `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1`. +- `ORT_DSA_INT8_VECTOR_STAGING=1`: the companion `DynamicSparseAttention` + grouped INT8 path uses eight-byte cache loads and independent per-channel + key/value scales. Its existing guards still require head size 256, FP16/BF16 + queries, one or two tokens, GQA ratio at least four, more than 256 candidates, + and sufficient resources. Auxiliary floating-point caches retain their own + unscaled load path. +- `ORT_DSA_SHARED_SCALE_MERGE=1`: grouped attention merges with at most 32 + splits compute the exponential factor once per split and share it across + output channels. A nonzero partial sum, not the maximum-score sentinel, + determines whether a split is active. Active factors that underflow to zero + remain in the accumulation. + +These options change kernels, not graph schemas or cache layouts. Qualify +accuracy and performance on the target runtime, hardware, and workload before +enabling them in a deployment. + 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. 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 6bd20ba933cb2..13297e1417c07 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -14,6 +14,7 @@ #include #include "core/common/safeint.h" +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cuda_common.h" namespace onnxruntime { @@ -892,7 +893,16 @@ __global__ void GroupedPrefillDynamicSparseAttentionKernel( } } -template +__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, @@ -1084,7 +1094,32 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( ((static_cast(b) * kv_num_heads + kv_head) * source_length + candidate_index) * head_size; - if constexpr (kQuantizedCache) { + 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] = @@ -1507,7 +1542,7 @@ __global__ void MergeDynamicSparseAttentionKernel(const float* partial_max, } } -template +template __global__ void MergeGroupedDecodeDynamicSparseAttentionKernel(const float* partial_max, const float* partial_sum, const float* partial_output, @@ -1525,6 +1560,51 @@ __global__ void MergeGroupedDecodeDynamicSparseAttentionKernel(const float* part 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; @@ -1973,19 +2053,31 @@ Status LaunchDynamicSparseAttention( const size_t grouped_shared_bytes = GetGroupedAttentionSharedBytes(parameters, data.quantized_cache ? sizeof(float) : sizeof(T)); if (data.quantized_cache) { - GroupedDecodeDynamicSparseAttentionKernel - <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( - data.prepared_query, data.present_key, data.present_value, - data.auxiliary_key, data.auxiliary_value, data.selected_indices, - data.selected_counts, data.seqlens_k, data.head_sink, - partial_max, partial_sum, partial_output, - static_cast(grouped_blocks), split_count, parameters.sequence_length, - parameters.num_heads, parameters.kv_num_heads, parameters.head_size, - parameters.cache_capacity, parameters.auxiliary_sequence_length, - parameters.max_selected, parameters.local_window_size, parameters.scale, - local_plus_selected, - parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, - parameters.use_smooth_softmax, data.key_scale, data.value_scale); + static const int vectorized_staging = + ParseEnvironmentVariableWithDefault("ORT_DSA_INT8_VECTOR_STAGING", 0); + ORT_ENFORCE(vectorized_staging == 0 || vectorized_staging == 1, + "ORT_DSA_INT8_VECTOR_STAGING must be 0 or 1."); + const auto launch = [&](auto vectorized) { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, data.key_scale, data.value_scale); + }; + if (vectorized_staging) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } } else if (parameters.cache_capacity <= 131072) { GroupedDecodeDynamicSparseAttentionKernel <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( @@ -2018,9 +2110,21 @@ Status LaunchDynamicSparseAttention( CUDA_RETURN_IF_ERROR(cudaGetLastError()); if (split_count <= 32) { - MergeGroupedDecodeDynamicSparseAttentionKernel<<(query_blocks), threads, 0, stream>>>( - partial_max, partial_sum, partial_output, data.output, - static_cast(query_blocks), split_count, parameters.head_size); + static const int shared_scales = + ParseEnvironmentVariableWithDefault("ORT_DSA_SHARED_SCALE_MERGE", 0); + ORT_ENFORCE(shared_scales == 0 || shared_scales == 1, + "ORT_DSA_SHARED_SCALE_MERGE must be 0 or 1."); + const auto merge = [&](auto specialized) { + MergeGroupedDecodeDynamicSparseAttentionKernel + <<(query_blocks), threads, 0, stream>>>( + partial_max, partial_sum, partial_output, data.output, + static_cast(query_blocks), split_count, parameters.head_size); + }; + if (shared_scales) { + merge(std::true_type{}); + } else { + merge(std::false_type{}); + } } else { const size_t merge_shared_bytes = static_cast(threads) * sizeof(float); MergeDynamicSparseAttentionKernel<<(query_blocks), threads, merge_shared_bytes, stream>>>( diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index b61b2d8c80b47..b1c293ffdbe1a 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -26,6 +26,7 @@ // See docs/contrib_ops/cuda/packed_sparse_attention_indexer.md for the full operator contract. #include "contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h" +#include "contrib_ops/cuda/sparse/qsa_topk.cuh" #include #include @@ -53,6 +54,7 @@ namespace { constexpr int kThreads = 128; constexpr int kWarpSize = 32; constexpr int kBoundedTopKMax = 512; +static_assert(kBoundedTopKMax == qsa_topk::kTopK); constexpr int kBoundedTopKItemsPerThread = kBoundedTopKMax / kThreads; constexpr int kDistributedTopKMaxRows = 64; constexpr int kDistributedTopKMinBlocks = 2048; @@ -876,6 +878,48 @@ __global__ void QsaMergeTwoLevelsTopKKernel(const uint64_t* input_keys, uint64_t } } +__device__ __forceinline__ void QsaEmitTopKRow(const uint64_t* top_keys, int token, + const int32_t* cumulative_sequence_lengths, + const int32_t* past_sequence_lengths, + const int64_t* position_ids, + const int32_t* present_state_lengths, + const int32_t* overflow_flags, + int32_t* selected_indices, int32_t* selected_counts, + PackedSparseAttentionIndexerParams params) { + int32_t* output = selected_indices + static_cast(token) * params.output_capacity; + const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); + const int64_t abs_position = + params.has_position_ids + ? position_ids[token] + : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]); + const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength]; + const int64_t causal_count = SaiCausalThreshold(abs_position, params.compress_ratio); + const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); + const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; + const int selected = min(params.block_topk, block_count); + for (int rank = threadIdx.x; rank < selected; rank += blockDim.x) { + const int selected_block = topk::UnpackStableSortIndex(top_keys[rank]); + for (int t = 0; t < params.compress_ratio; ++t) { + output[rank * params.compress_ratio + t] = selected_block * params.compress_ratio + t; + } + } + const int64_t block_start = static_cast(block_count) * params.compress_ratio; + const int64_t natural_tail = abs_position >= block_start ? abs_position - block_start + 1 : 0; + const int remaining_capacity = params.capacity - selected * params.compress_ratio; + const int64_t tail_count_64 = natural_tail < remaining_capacity ? natural_tail : remaining_capacity; + const int tail_count = overflow_flags[batch] == 0 ? static_cast(tail_count_64 < 0 ? 0 : tail_count_64) : 0; + for (int t = threadIdx.x; t < tail_count; t += blockDim.x) { + output[selected * params.compress_ratio + t] = static_cast(block_start + t); + } + for (int position = selected * params.compress_ratio + tail_count + threadIdx.x; + position < params.output_capacity; position += blockDim.x) { + output[position] = -1; + } + if (threadIdx.x == 0) { + selected_counts[token] = selected * params.compress_ratio + tail_count; + } +} + __global__ void QsaEmitHierarchicalTopKKernel(const uint64_t* tile_keys, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, @@ -885,38 +929,37 @@ __global__ void QsaEmitHierarchicalTopKKernel(const uint64_t* tile_keys, int32_t* selected_indices, int32_t* selected_counts, PackedSparseAttentionIndexerParams params) { for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) { - int32_t* output = selected_indices + static_cast(token) * params.output_capacity; - const int batch = PackedBatchOfToken(cumulative_sequence_lengths, params.batch_size, token); - const int64_t abs_position = - params.has_position_ids - ? position_ids[token] - : static_cast(past_sequence_lengths[batch]) + (token - cumulative_sequence_lengths[batch]); - const int key_len_after = present_state_lengths[batch * 2 + psai::kKeyStateLength]; - const int64_t causal_count = SaiCausalThreshold(abs_position, params.compress_ratio); - const int64_t visible_64 = causal_count < key_len_after ? causal_count : static_cast(key_len_after); - const int block_count = overflow_flags[batch] == 0 ? static_cast(visible_64 < 0 ? 0 : visible_64) : 0; - const int selected = min(params.block_topk, block_count); - for (int rank = threadIdx.x; rank < selected; rank += blockDim.x) { - const int selected_block = topk::UnpackStableSortIndex(tile_keys[static_cast(token) * key_stride + rank]); - for (int t = 0; t < params.compress_ratio; ++t) { - output[rank * params.compress_ratio + t] = selected_block * params.compress_ratio + t; - } - } - const int64_t block_start = static_cast(block_count) * params.compress_ratio; - const int64_t natural_tail = abs_position >= block_start ? abs_position - block_start + 1 : 0; - const int remaining_capacity = params.capacity - selected * params.compress_ratio; - const int64_t tail_count_64 = natural_tail < remaining_capacity ? natural_tail : remaining_capacity; - const int tail_count = overflow_flags[batch] == 0 ? static_cast(tail_count_64 < 0 ? 0 : tail_count_64) : 0; - for (int t = threadIdx.x; t < tail_count; t += blockDim.x) { - output[selected * params.compress_ratio + t] = static_cast(block_start + t); - } - for (int position = selected * params.compress_ratio + tail_count + threadIdx.x; - position < params.output_capacity; position += blockDim.x) { - output[position] = -1; - } - if (threadIdx.x == 0) { - selected_counts[token] = selected * params.compress_ratio + tail_count; + QsaEmitTopKRow(tile_keys + static_cast(token) * key_stride, token, + cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, selected_indices, selected_counts, params); + } +} + +__global__ void QsaFinalMergeAndEmitKernel(const uint64_t* compacted_keys, + const int32_t* cumulative_sequence_lengths, + const int32_t* past_sequence_lengths, + const int64_t* position_ids, + const int32_t* present_state_lengths, + const int32_t* overflow_flags, int compact_stride, int chunks, + int32_t* selected_indices, int32_t* selected_counts, + PackedSparseAttentionIndexerParams params) { + __shared__ uint64_t merged[kBoundedTopKMax]; + const int tid = static_cast(threadIdx.x); + for (int token = static_cast(blockIdx.x); token < params.total_tokens; token += static_cast(gridDim.x)) { + const uint64_t* row = compacted_keys + static_cast(token) * compact_stride; + merged[tid] = row[tid]; + __syncthreads(); + for (int chunk = 1; chunk < chunks; ++chunk) { + const uint64_t* next = row + chunk * kBoundedTopKMax; + if (next[0] == topk::kPaddingSortKey) break; + const uint64_t key = QsaMergeRank(merged, kBoundedTopKMax, next, kBoundedTopKMax, tid); + __syncthreads(); + merged[tid] = key; + __syncthreads(); } + QsaEmitTopKRow(merged, token, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, selected_indices, selected_counts, params); + __syncthreads(); } } @@ -1699,7 +1742,37 @@ Status LaunchQsaPackedSparseAttentionIndexer( present_state_lengths, overflow_flags, merge_input, tile_count, params); static const bool use_radix_topk = ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", false); - if (use_radix_topk) { + static const int compact_topk = + ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK", 0); + ORT_ENFORCE(compact_topk == 0 || compact_topk == 1, + "ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK must be 0 or 1."); + ORT_ENFORCE(!compact_topk || !use_radix_topk, + "Compact and radix packed sparse indexer top-K options cannot both be enabled."); + if (compact_topk) { + const int chunks = (key_stride + qsa_topk::kChunkSize - 1) / qsa_topk::kChunkSize; + const int compact_stride = chunks * kBoundedTopKMax; + qsa_topk::QsaCompactTileTopKKernel<<>>( + merge_input, merge_output, key_stride, compact_stride); + std::swap(merge_input, merge_output); + if (chunks <= 5) { + QsaFinalMergeAndEmitKernel<<>>( + merge_input, cumulative_sequence_lengths, past_sequence_lengths, position_ids, + present_state_lengths, overflow_flags, compact_stride, chunks, selected_indices, selected_counts, params); + } else { + for (int list_width = kBoundedTopKMax; list_width < compact_stride; list_width *= 2) { + const int pairs_per_row = (compact_stride + 2 * list_width - 1) / (2 * list_width); + const int64_t merge_work = static_cast(params.total_tokens) * pairs_per_row; + QsaMergeTileTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), + kHierarchicalMergeThreads, 0, stream>>>( + merge_input, merge_output, compact_stride, list_width, params.total_tokens); + std::swap(merge_input, merge_output); + } + QsaEmitHierarchicalTopKKernel<<>>( + merge_input, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, + overflow_flags, compact_stride, selected_indices, selected_counts, params); + } + return CUDA_CALL(cudaGetLastError()); + } else if (use_radix_topk) { QsaPartialTopKKernel<<>>( selection_workspace, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, overflow_flags, params); diff --git a/onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh b/onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh new file mode 100644 index 0000000000000..a178940858f97 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include "core/providers/cuda/cu_inc/topk_warp_sort.cuh" + +namespace onnxruntime::contrib::cuda::qsa_topk { + +constexpr int kChunkSize = 2048; +constexpr int kTopK = 512; +constexpr int kThreads = 512; + +namespace topk = onnxruntime::cuda::topk; + +template +using ChunkSort = cub::BlockMergeSort; + +template +__device__ __forceinline__ void SortChunk(const uint64_t* input, uint64_t* output, int input_count, + typename ChunkSort::TempStorage& temp) { + uint64_t keys[ItemsPerThread]; +#pragma unroll + for (int item = 0; item < ItemsPerThread; ++item) { + const int index = threadIdx.x * ItemsPerThread + item; + keys[item] = index < input_count ? input[index] : topk::kPaddingSortKey; + } + ChunkSort(temp).Sort(keys, topk::Greater()); +#pragma unroll + for (int item = 0; item < ItemsPerThread; ++item) { + const int rank = threadIdx.x * ItemsPerThread + item; + if (rank < kTopK) output[rank] = keys[item]; + } +} + +__global__ void QsaCompactTileTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, + int input_stride, int output_stride) { + __shared__ union { + ChunkSort<1>::TempStorage small; + ChunkSort<2>::TempStorage medium; + ChunkSort<4>::TempStorage large; + } temp; + const int chunk_start = static_cast(blockIdx.x) * kChunkSize; + const uint64_t* input = input_keys + static_cast(blockIdx.y) * input_stride + chunk_start; + uint64_t* output = output_keys + static_cast(blockIdx.y) * output_stride + blockIdx.x * kTopK; + const int count = min(kChunkSize, input_stride - chunk_start); + + // Scoring emits sorted 32-key tiles in a live prefix followed by zero padding. + // Tile-aligned probes adapt to device context growth without graph recapture. + if (input[0] == topk::kPaddingSortKey) { + output[threadIdx.x] = topk::kPaddingSortKey; + } else if (count > 1024 && input[1024] != topk::kPaddingSortKey) { + SortChunk<4>(input, output, count, temp.large); + } else if (count > 512 && input[512] != topk::kPaddingSortKey) { + SortChunk<2>(input, output, min(count, 1024), temp.medium); + } else { + SortChunk<1>(input, output, min(count, 512), temp.small); + } +} + +} // namespace onnxruntime::contrib::cuda::qsa_topk diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 5a3ad2f1dd8a1..334ff9e765d5b 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -434,7 +434,8 @@ TEST(DynamicSparseAttentionTest, Int8CacheRejectsMissingOrMalformedScales_CUDA) DynamicSparseAttentionCase MakeGroupedCacheCase(int sequence_length, bool auxiliary, bool floating_cache = false, int batch_size = 2, - int num_heads = 20, int kv_num_heads = 2) { + int num_heads = 20, int kv_num_heads = 2, + bool varied_channels = false, int selected_count = 509) { DynamicSparseAttentionCase c; c.batch_size = batch_size; c.sequence_length = sequence_length; @@ -459,11 +460,11 @@ DynamicSparseAttentionCase MakeGroupedCacheCase(int sequence_length, bool auxili c.key_scale.resize(channels); c.value_scale.resize(channels); for (int d = 0; d < channels; ++d) { - c.key_scale[d] = floating_cache ? 1.0f : (d % 2 ? 0.125f : 0.0625f); - c.value_scale[d] = floating_cache ? 1.0f : (d % 2 ? 0.25f : 0.125f); + c.key_scale[d] = floating_cache ? 1.0f : (varied_channels ? (d % 8 + 1) * 0.03125f : (d % 2 ? 0.125f : 0.0625f)); + c.value_scale[d] = floating_cache ? 1.0f : (varied_channels ? (8 - d % 8) * 0.0625f : (d % 2 ? 0.25f : 0.125f)); } for (size_t i = 0; i < c.query.size(); ++i) { - c.query[i] = float(int(i % 7) - 3) * 0.125f; + c.query[i] = float(int(i % 7) - 3) * (varied_channels ? 0.00390625f : 0.125f); } for (size_t i = 0; i < c.key.size(); ++i) { c.key[i] = float(int(i % 11) - 5) * c.key_scale[i % channels]; @@ -471,9 +472,10 @@ DynamicSparseAttentionCase MakeGroupedCacheCase(int sequence_length, bool auxili } c.past_key.resize(c.batch_size * c.cache_sequence_length * channels); c.past_value.resize(c.past_key.size()); + constexpr int8_t kSignedBytes[] = {-128, -127, -65, -1, 0, 1, 64, 127}; for (size_t i = 0; i < c.past_key.size(); ++i) { - c.past_key[i] = float(int(i % 17) - 8); - c.past_value[i] = float(int(i % 19) - 9); + c.past_key[i] = varied_channels ? float(kSignedBytes[i % 8]) : float(int(i % 17) - 8); + c.past_value[i] = varied_channels ? float(kSignedBytes[7 - i % 8]) : float(int(i % 19) - 9); } c.expected_present_key = c.past_key; c.expected_present_value = c.past_value; @@ -504,7 +506,7 @@ DynamicSparseAttentionCase MakeGroupedCacheCase(int sequence_length, bool auxili for (int row = 0; row < rows; ++row) { const int b = row / sequence_length; const int pos = c.seqlens_k[b] + 1 - sequence_length + row % sequence_length; - c.selected_counts[row] = row == 0 ? 0 : 509; + c.selected_counts[row] = row == 0 ? 0 : selected_count; std::iota(c.selected_indices.begin() + row * c.max_selected, c.selected_indices.begin() + row * c.max_selected + c.selected_counts[row], 0); for (int h = 0; h < c.num_heads; ++h) { @@ -600,6 +602,34 @@ TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeFloat16_CUDA) { MakeGroupedCacheCase(2, false), std::move(cuda_ep)); } +TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeVectorStagingChannels_CUDA) { + for (int sequence_length : {1, 2}) { + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + + ", auxiliary=" + std::to_string(auxiliary)); + const auto c = MakeGroupedCacheCase(sequence_length, auxiliary, false, 2, 20, 2, true); + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + cuda_ep = DefaultCudaExecutionProvider(); + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + } + } +} + +TEST(DynamicSparseAttentionTest, GroupedMergeActiveSplitBoundaries_CUDA) { + for (int count : {1, 31, 32, 33, 255, 256, 257, 511}) { + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("selected_count=" + std::to_string(count) + + ", auxiliary=" + std::to_string(auxiliary)); + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase( + MakeGroupedCacheCase(2, auxiliary, false, 2, 20, 2, true, count), std::move(cuda_ep)); + } + } +} + TEST(DynamicSparseAttentionTest, Int8CacheRejectsNonpositiveOrNonfiniteScales_CUDA) { ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{{"ORT_DYNAMIC_SPARSE_ATTENTION_STRICT_VALIDATION", "1"}}}; diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 5db1dd84bc2c8..6ba1a06bbb5c9 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -1344,6 +1344,28 @@ TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalPairedVerification) { } } +TEST(PackedSparseAttentionIndexerTest, QsaCompactChunkAndFallbackBoundaries) { + for (int live_blocks : {511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049, 8193, 10240, 10241}) { + SCOPED_TRACE(live_blocks); + QsaPackedProblem problem; + problem.batch_size = 2; + problem.cumulative_sequence_lengths = {0, 2, 3}; + problem.past_sequence_lengths = {live_blocks * 4, 0}; + problem.head_size = 128; + problem.num_heads = 4; + problem.rotary_width = 64; + problem.compress_ratio = 4; + problem.token_budget = 2048; + problem.state_capacity = std::max(2051, live_blocks + 1); + problem.max_position = live_blocks * 4 + 4; + problem = MakeQsaPackedProblem(std::move(problem)); + std::fill(problem.query.begin(), problem.query.end(), 0.0f); + std::fill(problem.cos_cache.begin(), problem.cos_cache.end(), 1.0f); + std::fill(problem.sin_cache.begin(), problem.sin_cache.end(), 0.0f); + RunQsaPackedTest(2.0e-2f, std::move(problem), ProviderKind::Cuda); + } +} + TEST(PackedSparseAttentionIndexerTest, QsaRaggedPrefill) { QsaPackedProblem problem; problem.batch_size = 1; From b71498b2660189a4d3638de2095cf1d74e691ce7 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 17:32:10 +0000 Subject: [PATCH 07/23] Apply lintrunner formatting to PR changes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/sparse/packed_sparse_attention_indexer_common.h | 8 ++++---- .../cuda/bert/dynamic_sparse_attention_impl.cu | 2 +- .../packed_sparse_attention_indexer_op_test.cc | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h index 0bb41846aac38..8e95378a3f4f4 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h @@ -71,10 +71,10 @@ enum OutputIndex : int { kPresentStateLengths = 5, // [batch_size, 2], int32 kStateUpdate = 6, // qsa only, optional: [batch_size, state_update_capacity, head_size] kIndexShareStatus = 7, - kStateStatus = 8, // optional: [batch_size], int32, see StateStatusCode - kMidKvBuffer = 9, // qsa only, optional: snapshot after the first request token - kMidStateLengths = 10, // qsa only, optional: lengths after the first request token - kValidationStatus = 11, // qsa only, optional: per-request validation status + kStateStatus = 8, // optional: [batch_size], int32, see StateStatusCode + kMidKvBuffer = 9, // qsa only, optional: snapshot after the first request token + kMidStateLengths = 10, // qsa only, optional: lengths after the first request token + kValidationStatus = 11, // qsa only, optional: per-request validation status kOutputCount = 9, kMaxOutputCount = 12, }; 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 4c8553950e4e7..05265581e336a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -2074,7 +2074,7 @@ Status LaunchDynamicSparseAttention( parameters.cache_capacity, parameters.max_selected, parameters.scale); } } else if (grouped_decode_threads_supported && - UseGroupedDecodeAttention(parameters, sizeof(T), max_shared_memory_per_block, data.quantized_cache)) { + UseGroupedDecodeAttention(parameters, sizeof(T), max_shared_memory_per_block, data.quantized_cache)) { const int split_count = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); const int64_t partial_blocks = query_blocks * split_count; const int64_t grouped_blocks = diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 606ac20a145cc..af40c5f2c2979 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -1114,10 +1114,10 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP ASSERT_EQ(problem.cumulative_sequence_lengths[batch + 1] - problem.cumulative_sequence_lengths[batch], 1); } test.AddOutput("mid_kv_buffer", {batch_size, buffer_capacity, head_size}, - ToElementType(expected.present_kv_buffer), false, 0.0f, tolerance); + ToElementType(expected.present_kv_buffer), false, 0.0f, tolerance); test.AddOutput("mid_state_lengths", {batch_size, 2}, expected.present_state_lengths); test.AddOutput("validation_status", {batch_size}, - std::vector(static_cast(batch_size), 0)); + std::vector(static_cast(batch_size), 0)); } RunOnProvider(test, std::move(provider), expected_error.empty() ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, From c0c4e7dd2045c6039e838137a743193956240e5a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 22 Oct 2026 23:30:30 +0000 Subject: [PATCH 08/23] Enable default-on sparse attention optimizations, S<=8 grouped decode, and fuse validation kernels Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/dynamic_sparse_attention.cc | 8 +- .../bert/dynamic_sparse_attention_impl.cu | 154 ++++++++---------- .../cuda/bert/dynamic_sparse_attention_impl.h | 11 +- .../sparse/packed_sparse_attention_indexer.cc | 12 +- .../packed_sparse_attention_indexer_impl.cu | 35 +--- .../packed_sparse_attention_indexer_impl.h | 8 +- 6 files changed, 86 insertions(+), 142 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index f7a65fbad3663..9555b9f919cd6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -256,12 +256,8 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons ORT_RETURN_IF_ERROR(ValidateDynamicSparseAttentionOnDevice( stream, data.selected_indices, data.selected_counts, data.seqlens_k, data.position_ids, parameters, validation_error, validation_workspace.get(), - GetDeviceProp().sharedMemPerBlock, strict_validation_, data.key_scale, data.value_scale)); - if (committed_length != nullptr) { - ORT_RETURN_IF_ERROR(ValidateDynamicSparseAttentionCursor( - stream, data.seqlens_k, committed_length->Data(), - parameters.sequence_length, parameters.cache_capacity, validation_error)); - } + GetDeviceProp().sharedMemPerBlock, strict_validation_, data.key_scale, data.value_scale, + committed_length == nullptr ? nullptr : committed_length->Data())); } const size_t attention_workspace_elements = 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 05265581e336a..13a7140f3b859 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -18,8 +18,6 @@ #include "core/platform/env_var_utils.h" #include "core/providers/cuda/cuda_common.h" -#include "core/platform/env_var_utils.h" - namespace onnxruntime { namespace contrib { namespace cuda { @@ -77,7 +75,11 @@ __device__ __forceinline__ void MmaM16N8K16(float (&output)[4], __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 == 9) { + atomicCAS(error_flag, 9, static_cast(error)); + } } __global__ void ValidateInputsKernel(const int32_t* selected_indices, @@ -99,7 +101,8 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, int32_t* error_flag, const float* key_scale, const float* value_scale, - int cache_channels) { + int cache_channels, + const int32_t* committed_length) { const int row = static_cast(blockIdx.x); const int b = row / sequence_length; const int s = row - b * sequence_length; @@ -225,6 +228,14 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, } return; } + + if (row == 0 && threadIdx.x == 0 && committed_length != nullptr) { + const int64_t committed = *committed_length; + const int64_t total = committed + sequence_length; + if (committed < 0 || total > main_capacity || seqlens_k[0] != total - 1) { + atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, 9); + } + } } template @@ -1732,8 +1743,8 @@ bool UseGroupedDecodeAttention(const DynamicSparseAttentionParameters& parameter size_t element_size, size_t max_shared_memory_per_block, bool quantized_cache = false) { - return (parameters.sequence_length == 1 || - (quantized_cache && parameters.sequence_length == 2)) && + 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) && @@ -1777,23 +1788,6 @@ bool UseSharedValidationHashTable(const DynamicSparseAttentionParameters& parame shared_bytes < max_shared_memory_per_block; } -__global__ void ValidateCommittedLengthKernel( - const int32_t* seqlens_k, - const int32_t* committed_length, - int32_t sequence_length, - int32_t cache_capacity, - int32_t* status) { - if (blockIdx.x != 0 || threadIdx.x != 0) { - return; - } - const int64_t committed = *committed_length; - const int64_t total = committed + sequence_length; - if (*status == kDynamicSparseAttentionValidationOk && - (committed < 0 || total > cache_capacity || *seqlens_k != total - 1)) { - *status = 9; - } -} - } // namespace size_t GetDynamicSparseAttentionValidationWorkspaceSize( @@ -1824,17 +1818,19 @@ size_t GetDynamicSparseAttentionWorkspaceSize( return SafeInt(query_blocks) * split_count * (static_cast(parameters.head_size) + 2); } - 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 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 bool grouped_decode = - UseGroupedDecodeAttention(parameters, element_size, max_shared_memory_per_block); const size_t split_count = grouped_decode ? GetGroupedDecodeSplitCount(parameters) @@ -1858,7 +1854,8 @@ Status ValidateDynamicSparseAttentionOnDevice( size_t max_shared_memory_per_block, bool copy_result_to_host, const float* key_scale, - const float* value_scale) { + const float* value_scale, + const int32_t* committed_length) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(error_flag, 0, sizeof(int32_t), stream)); const int64_t row_count = static_cast(parameters.batch_size) * parameters.sequence_length; @@ -1887,7 +1884,8 @@ Status ValidateDynamicSparseAttentionOnDevice( parameters.rotary_max_position, parameters.do_rotary, parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, validation_bitmap, validation_hash_entries, use_shared_hash_table, error_flag, - key_scale, value_scale, parameters.kv_num_heads * parameters.head_size); + key_scale, value_scale, parameters.kv_num_heads * parameters.head_size, + committed_length); CUDA_RETURN_IF_ERROR(cudaGetLastError()); if (!copy_result_to_host) { @@ -1898,7 +1896,7 @@ Status ValidateDynamicSparseAttentionOnDevice( CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( &host_error, error_flag, sizeof(host_error), cudaMemcpyDeviceToHost, stream)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); - if (host_error == kDynamicSparseAttentionValidationOk) { + if (host_error == kDynamicSparseAttentionValidationOk || host_error == 9) { return Status::OK(); } @@ -1938,18 +1936,6 @@ Status ValidateDynamicSparseAttentionOnDevice( "DynamicSparseAttention: ", message, "."); } -Status ValidateDynamicSparseAttentionCursor( - cudaStream_t stream, - const int32_t* seqlens_k, - const int32_t* committed_length, - int32_t sequence_length, - int32_t cache_capacity, - int32_t* status) { - ValidateCommittedLengthKernel<<<1, 1, 0, stream>>>( - seqlens_k, committed_length, sequence_length, cache_capacity, status); - return CUDA_CALL(cudaGetLastError()); -} - template Status LaunchDynamicSparseAttention( cudaStream_t stream, @@ -2035,46 +2021,37 @@ Status LaunchDynamicSparseAttention( const size_t fused_shared_bytes = (static_cast(candidate_capacity) + static_cast(fused_threads)) * sizeof(float) + static_cast(parameters.head_size) * sizeof(T); - bool use_prefill_kernel = false; - if constexpr (std::is_same_v) { - const bool enabled = ParseEnvironmentVariableWithDefault("ORT_DSA_PREFILL_KERNEL", 1) != 0; - // max_shared_memory_per_block is the default 48 KB limit; the kernel opts in to more and requires SM80+. - int device = 0; - int cc_major = 0; - int optin_shared = 0; - CUDA_RETURN_IF_ERROR(cudaGetDevice(&device)); - CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device)); - CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); - constexpr int kMaxGridDimYZ = 65535; - use_prefill_kernel = - enabled && !data.quantized_cache && cc_major >= 8 && parameters.sequence_length >= 8 && - parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && - parameters.selected_kv_source == DynamicSparseAttentionKvSource::kMain && - !parameters.use_smooth_softmax && parameters.max_selected > 0 && - parameters.kv_num_heads > 0 && parameters.kv_num_heads <= kMaxGridDimYZ && - parameters.num_heads % parameters.kv_num_heads == 0 && - parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && - static_cast(optin_shared) >= dsa_prefill::kSharedBytes; - } constexpr int kGroupedDecodeThreads = 32 * kGroupedDecodeQueryHeadsPerBlock; constexpr int kGroupedPrefillThreads = 32 * kGroupedPrefillQueryHeadsPerBlock; const bool grouped_decode_threads_supported = kGroupedDecodeThreads <= max_threads_per_block; const bool grouped_prefill_threads_supported = kGroupedPrefillThreads <= max_threads_per_block; - if (use_prefill_kernel) { - if constexpr (std::is_same_v) { - CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( - dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - static_cast(dsa_prefill::kSharedBytes))); - const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), - static_cast(parameters.kv_num_heads)); - dsa_prefill::DynamicSparseAttentionPrefillKernel<<>>( - data.prepared_query, data.present_key, data.present_value, data.selected_indices, data.selected_counts, - data.seqlens_k, data.output, parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, - parameters.cache_capacity, parameters.max_selected, parameters.scale); + const bool use_grouped_decode = + grouped_decode_threads_supported && + UseGroupedDecodeAttention(parameters, sizeof(T), max_shared_memory_per_block, data.quantized_cache); + bool use_prefill_kernel = false; + if constexpr (std::is_same_v) { + if (!use_grouped_decode) { + const bool enabled = ParseEnvironmentVariableWithDefault("ORT_DSA_PREFILL_KERNEL", 1) != 0; + // max_shared_memory_per_block is the default 48 KB limit; the kernel opts in to more and requires SM80+. + int device = 0; + int cc_major = 0; + int optin_shared = 0; + CUDA_RETURN_IF_ERROR(cudaGetDevice(&device)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + constexpr int kMaxGridDimYZ = 65535; + use_prefill_kernel = + enabled && !data.quantized_cache && cc_major >= 8 && parameters.sequence_length >= 8 && + parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kMain && + !parameters.use_smooth_softmax && parameters.max_selected > 0 && + parameters.kv_num_heads > 0 && parameters.kv_num_heads <= kMaxGridDimYZ && + parameters.num_heads % parameters.kv_num_heads == 0 && + parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && + static_cast(optin_shared) >= dsa_prefill::kSharedBytes; } - } else if (grouped_decode_threads_supported && - UseGroupedDecodeAttention(parameters, sizeof(T), max_shared_memory_per_block, data.quantized_cache)) { + } + if (use_grouped_decode) { const int split_count = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); const int64_t partial_blocks = query_blocks * split_count; const int64_t grouped_blocks = @@ -2091,7 +2068,7 @@ Status LaunchDynamicSparseAttention( GetGroupedAttentionSharedBytes(parameters, data.quantized_cache ? sizeof(float) : sizeof(T)); if (data.quantized_cache) { static const int vectorized_staging = - ParseEnvironmentVariableWithDefault("ORT_DSA_INT8_VECTOR_STAGING", 0); + ParseEnvironmentVariableWithDefault("ORT_DSA_INT8_VECTOR_STAGING", 1); ORT_ENFORCE(vectorized_staging == 0 || vectorized_staging == 1, "ORT_DSA_INT8_VECTOR_STAGING must be 0 or 1."); const auto launch = [&](auto vectorized) { @@ -2148,7 +2125,7 @@ Status LaunchDynamicSparseAttention( if (split_count <= 32) { static const int shared_scales = - ParseEnvironmentVariableWithDefault("ORT_DSA_SHARED_SCALE_MERGE", 0); + ParseEnvironmentVariableWithDefault("ORT_DSA_SHARED_SCALE_MERGE", 1); ORT_ENFORCE(shared_scales == 0 || shared_scales == 1, "ORT_DSA_SHARED_SCALE_MERGE must be 0 or 1."); const auto merge = [&](auto specialized) { @@ -2168,6 +2145,19 @@ Status LaunchDynamicSparseAttention( partial_max, partial_sum, partial_output, data.output, static_cast(query_blocks), split_count, parameters.head_size); } + } else if (use_prefill_kernel) { + if constexpr (std::is_same_v) { + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kSharedBytes))); + const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), + static_cast(parameters.kv_num_heads)); + dsa_prefill::DynamicSparseAttentionPrefillKernel<<>>( + data.prepared_query, data.present_key, data.present_value, data.selected_indices, data.selected_counts, + data.seqlens_k, data.output, parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, + parameters.cache_capacity, parameters.max_selected, parameters.scale); + } } else if (!data.quantized_cache && grouped_prefill_threads_supported && UseGroupedPrefillAttention(parameters, sizeof(T), max_shared_memory_per_block)) { const int64_t grouped_blocks = diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index 2367f77b89630..6050386bd9c71 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -75,15 +75,8 @@ Status ValidateDynamicSparseAttentionOnDevice( size_t max_shared_memory_per_block, bool copy_result_to_host, const float* key_scale = nullptr, - const float* value_scale = nullptr); - -Status ValidateDynamicSparseAttentionCursor( - cudaStream_t stream, - const int32_t* seqlens_k, - const int32_t* committed_length, - int32_t sequence_length, - int32_t cache_capacity, - int32_t* status); + const float* value_scale = nullptr, + const int32_t* committed_length = nullptr); template Status LaunchDynamicSparseAttention( diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index 8840b1573cc08..2cbe321adc9fe 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -415,6 +415,8 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* cont mid_kv_buffer == nullptr ? nullptr : reinterpret_cast(mid_kv_buffer->MutableData()); params.mid_state_lengths = mid_state_lengths == nullptr ? nullptr : mid_state_lengths->MutableData(); + params.validation_status = + validation_status == nullptr ? nullptr : validation_status->MutableData(); auto float_workspace = GetScratchBuffer(GetQsaPackedWorkspaceFloatCount(params), GetComputeStream(context)); Tensor* state_status = GetStateStatusOutput(context, batch_size); @@ -426,7 +428,7 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* cont const CudaT* key_data = has_separate_key ? reinterpret_cast(key->Data()) : query_data + query_width; - ORT_RETURN_IF_ERROR((LaunchQsaPackedSparseAttentionIndexer( + return LaunchQsaPackedSparseAttentionIndexer( Stream(context), params, query_data, key_data, @@ -449,13 +451,7 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* cont present_state_lengths->MutableData(), state_update != nullptr ? reinterpret_cast(state_update->MutableData()) : nullptr, float_workspace.get(), - overflow_flags))); - if (validation_status != nullptr) { - return LaunchMapQsaValidationStatus( - Stream(context), cumulative_sequence_lengths->Data(), overflow_flags, - static_cast(batch_size), validation_status->MutableData()); - } - return Status::OK(); + overflow_flags); } template diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index eba7d53fb141e..a0d142526aecb 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -200,6 +200,10 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con overflow_flags[b] = invalid_metadata ? psai::kStateStatusInvalidMetadata : rejected ? psai::kStateStatusCapacityOverflow : psai::kStateStatusOk; + if (params.validation_status != nullptr) { + const int request_length = req_end - req_start; + params.validation_status[b] = !rejected ? 0 : (request_length >= 2 ? 0x99 : 9); + } } if (params.mid_kv_buffer != nullptr && params.mid_state_lengths != nullptr) { @@ -1643,35 +1647,6 @@ size_t GetCsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& return rows * params.num_heads * params.head_size + rows * static_cast(std::max(params.state_capacity, 1)); } -__global__ void MapQsaValidationStatusKernel( - const int32_t* cumulative_sequence_lengths, - const int32_t* overflow_flags, - int batch_size, - int32_t* validation_status) { - const int batch = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (batch >= batch_size) { - return; - } - const int request_length = cumulative_sequence_lengths[batch + 1] - cumulative_sequence_lengths[batch]; - validation_status[batch] = overflow_flags[batch] == 0 ? 0 : (request_length >= 2 ? 0x99 : 9); -} - -Status LaunchMapQsaValidationStatus( - cudaStream_t stream, - const int32_t* cumulative_sequence_lengths, - const int32_t* overflow_flags, - int batch_size, - int32_t* validation_status) { - if (batch_size <= 0) { - return Status::OK(); - } - constexpr int kStatusThreads = 128; - const int blocks = (batch_size + kStatusThreads - 1) / kStatusThreads; - MapQsaValidationStatusKernel<<>>( - cumulative_sequence_lengths, overflow_flags, batch_size, validation_status); - return CUDA_CALL(cudaGetLastError()); -} - template Status LaunchQsaPackedSparseAttentionIndexer( cudaStream_t stream, const PackedSparseAttentionIndexerParams& params, const T* query, const T* key, @@ -1749,7 +1724,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( static const bool use_radix_topk = ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", false); static const int compact_topk = - ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK", 0); + ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK", 1); ORT_ENFORCE(compact_topk == 0 || compact_topk == 1, "ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK must be 0 or 1."); ORT_ENFORCE(!compact_topk || !use_radix_topk, diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h index a33d0750c12bf..6552a0bbe6de9 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h @@ -41,6 +41,7 @@ struct PackedSparseAttentionIndexerParams { int state_update_capacity = 0; void* mid_kv_buffer = nullptr; // [batch, buffer_capacity, head_size] T int32_t* mid_state_lengths = nullptr; // [batch, 2] + int32_t* validation_status = nullptr; // [batch] // policy_mode = "csa" int index_topk = 0; @@ -90,13 +91,6 @@ Status LaunchQsaPackedSparseAttentionIndexer( float* float_workspace, int32_t* overflow_flags); -Status LaunchMapQsaValidationStatus( - cudaStream_t stream, - const int32_t* cumulative_sequence_lengths, - const int32_t* overflow_flags, - int batch_size, - int32_t* validation_status); - template Status LaunchCsaPackedSparseAttentionIndexer( cudaStream_t stream, From 6548c21cde55be23344bd05095972147098a231d Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:27:28 +0000 Subject: [PATCH 09/23] [CUDA] Fold #33546 and S=4 prefix-state outputs (mid_kv_buffer2/3, mid_state_lengths2/3) + max_splits=32 into PR #33472 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 10 +++++++++- docs/OperatorKernels.md | 2 +- .../cuda/dynamic_sparse_attention.md | 9 +++++++++ .../packed_sparse_attention_indexer_common.h | 6 +++++- .../bert/dynamic_sparse_attention_impl.cu | 3 ++- .../sparse/packed_sparse_attention_indexer.cc | 18 +++++++++++++++++ .../packed_sparse_attention_indexer_impl.cu | 14 ++++++++----- .../packed_sparse_attention_indexer_impl.h | 4 ++++ .../core/graph/contrib_ops/bert_defs.cc | 20 +++++++++++++++++++ .../dynamic_sparse_attention_op_test.cc | 4 ++-- 10 files changed, 79 insertions(+), 11 deletions(-) create mode 100644 docs/contrib_ops/cuda/dynamic_sparse_attention.md diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 97197c4993f02..024ae399e0d71 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5326,7 +5326,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Merged device status [N].
-#### Outputs (5 - 12) +#### Outputs (5 - 16)
selected_indices : M
@@ -5353,6 +5353,14 @@ This version of the operator has been available since version 1 of the 'com.micr
Only for CUDA policy_mode 'qsa' refresh: state lengths after the first token of each request. Not supported in IndexShare reuse mode.
validation_status (optional) : M
Only for CUDA policy_mode 'qsa' refresh: per-request overflow status; two-token rejection is encoded as 0x99. Not supported in IndexShare reuse mode.
+
mid_kv_buffer2 (optional) : T
+
CUDA QSA pending-buffer snapshot after at most two request tokens.
+
mid_state_lengths2 (optional) : M
+
CUDA QSA state lengths after at most two request tokens.
+
mid_kv_buffer3 (optional) : T
+
CUDA QSA pending-buffer snapshot after at most three request tokens.
+
mid_state_lengths3 (optional) : M
+
CUDA QSA state lengths after at most three request tokens.
#### Type Constraints diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index ed6f15d60c9ab..064bee07e72b7 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1150,7 +1150,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_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**
*out* mid_kv_buffer:**T**
*out* mid_state_lengths:**M**
*out* validation_status:**M**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**T_KEY** = 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**
*out* mid_kv_buffer:**T**
*out* mid_state_lengths:**M**
*out* validation_status:**M**
*out* mid_kv_buffer2:**T**
*out* mid_state_lengths2:**M**
*out* mid_kv_buffer3:**T**
*out* mid_state_lengths3:**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..0cbccd1bc25c7 --- /dev/null +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -0,0 +1,9 @@ +# 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}`). diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h index 8e95378a3f4f4..038c72be89811 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h @@ -75,8 +75,12 @@ enum OutputIndex : int { kMidKvBuffer = 9, // qsa only, optional: snapshot after the first request token kMidStateLengths = 10, // qsa only, optional: lengths after the first request token kValidationStatus = 11, // qsa only, optional: per-request validation status + kMidKvBuffer2 = 12, + kMidStateLengths2 = 13, + kMidKvBuffer3 = 14, + kMidStateLengths3 = 15, kOutputCount = 9, - kMaxOutputCount = 12, + kMaxOutputCount = 16, }; // Values of the optional state_status output. They follow the status convention of 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 13a7140f3b859..dfc0c575a72da 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -1764,7 +1764,8 @@ int GetGroupedDecodeSplitCount(const DynamicSparseAttentionParameters& parameter const int64_t occupancy_splits = std::max(1, (target_blocks + grouped_head_blocks - 1) / grouped_head_blocks); - return static_cast(std::min(candidate_splits, occupancy_splits)); + 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) { diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index 2cbe321adc9fe..e6493889d275e 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -392,6 +392,22 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* cont } PackedSparseAttentionIndexerParams params; + Tensor* mid_buffers[2]{}; + Tensor* mid_lengths[2]{}; + for (int prefix = 0; prefix < 2; ++prefix) { + const int buffer_index = psai::kMidKvBuffer2 + prefix * 2; + const int length_index = buffer_index + 1; + if (context->OutputCount() > buffer_index) + mid_buffers[prefix] = context->Output(buffer_index, TensorShape({batch_size, buffer_capacity, head_size})); + if (context->OutputCount() > length_index) + mid_lengths[prefix] = context->Output(length_index, TensorShape({batch_size, psai::kStateLengthColumns})); + ORT_RETURN_IF((mid_buffers[prefix] == nullptr) != (mid_lengths[prefix] == nullptr), + "PackedSparseAttentionIndexer: every prefix buffer requires its matching lengths"); + } + params.mid_kv_buffer2 = mid_buffers[0] ? mid_buffers[0]->MutableData() : nullptr; + params.mid_state_lengths2 = mid_lengths[0] ? mid_lengths[0]->MutableData() : nullptr; + params.mid_kv_buffer3 = mid_buffers[1] ? mid_buffers[1]->MutableData() : nullptr; + params.mid_state_lengths3 = mid_lengths[1] ? mid_lengths[1]->MutableData() : nullptr; params.batch_size = static_cast(batch_size); params.total_tokens = static_cast(total_tokens); params.num_heads = static_cast(num_heads); @@ -654,6 +670,8 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* cont template Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { + ORT_RETURN_IF(context->OutputCount() > psai::kMidKvBuffer2, + "PackedSparseAttentionIndexer: additional prefix outputs require refresh, not IndexShare reuse"); for (int index : {psai::kMidKvBuffer, psai::kMidStateLengths, psai::kValidationStatus}) { Tensor* output = context->OutputCount() > index ? context->Output(index, TensorShape({0})) : nullptr; ORT_RETURN_IF(output != nullptr, diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index a0d142526aecb..6578cfbf1216c 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -206,16 +206,20 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con } } - if (params.mid_kv_buffer != nullptr && params.mid_state_lengths != nullptr) { - const int mid_req = min(req_len, 1); + void* snapshots[3] = {params.mid_kv_buffer, params.mid_kv_buffer2, params.mid_kv_buffer3}; + int32_t* length_snapshots[3] = {params.mid_state_lengths, params.mid_state_lengths2, params.mid_state_lengths3}; +#pragma unroll + for (int prefix = 0; prefix < 3; ++prefix) { + if (snapshots[prefix] == nullptr || length_snapshots[prefix] == nullptr) continue; + const int mid_req = min(req_len, prefix + 1); const int mid_pending = old_buf_len + mid_req; const int mid_blocks = rejected ? 0 : mid_pending / params.compress_ratio; const int mid_buf_len = rejected ? old_buf_len : mid_pending % params.compress_ratio; if (threadIdx.x == 0) { - params.mid_state_lengths[b * 2 + psai::kKeyStateLength] = old_key_len + mid_blocks; - params.mid_state_lengths[b * 2 + psai::kBufferLength] = mid_buf_len; + length_snapshots[prefix][b * 2 + psai::kKeyStateLength] = old_key_len + mid_blocks; + length_snapshots[prefix][b * 2 + psai::kBufferLength] = mid_buf_len; } - T* mid_buffer = static_cast(params.mid_kv_buffer); + T* mid_buffer = static_cast(snapshots[prefix]); const int mid_total = params.buffer_capacity * params.head_size; for (int i = static_cast(threadIdx.x); i < mid_total; i += static_cast(blockDim.x)) { const int t = i / params.head_size; diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h index 6552a0bbe6de9..748b0d450fda8 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h @@ -41,6 +41,10 @@ struct PackedSparseAttentionIndexerParams { int state_update_capacity = 0; void* mid_kv_buffer = nullptr; // [batch, buffer_capacity, head_size] T int32_t* mid_state_lengths = nullptr; // [batch, 2] + void* mid_kv_buffer2 = nullptr; // [batch, buffer_capacity, head_size] T + int32_t* mid_state_lengths2 = nullptr; // [batch, 2] + void* mid_kv_buffer3 = nullptr; // [batch, buffer_capacity, head_size] T + int32_t* mid_state_lengths3 = nullptr; // [batch, 2] int32_t* validation_status = nullptr; // [batch] // policy_mode = "csa" diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 5a08f90baff53..289020132e463 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2953,6 +2953,18 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); propagateShapeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); } + for (int prefix = 0; prefix < 2; ++prefix) { + const int buffer = psai::kMidKvBuffer2 + 2 * prefix; + const int lengths = buffer + 1; + if (ctx.hasOutput(buffer) != ctx.hasOutput(lengths)) + fail_shape_inference("PackedSparseAttentionIndexer: prefix buffer and lengths must be connected together"); + if (!ctx.hasOutput(buffer)) continue; + if (!is_qsa) fail_shape_inference("PackedSparseAttentionIndexer: prefix outputs require policy_mode='qsa'"); + propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, buffer); + propagateShapeFromInputToOutput(ctx, psai::kPastKvBuffer, buffer); + updateOutputElemType(ctx, lengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); + propagateShapeFromInputToOutput(ctx, psai::kPastStateLengths, lengths); + } if (!is_qsa) { // present_gate_buffer keeps its fixed positional slot (with an empty name) for policy_mode // 'qsa'; only propagate its type/shape when it is actually produced. @@ -3497,6 +3509,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Only for CUDA policy_mode 'qsa' refresh: per-request overflow status; two-token rejection is encoded as 0x99. " "Not supported in IndexShare reuse mode.", "M", OpSchema::Optional) + .Output(12, "mid_kv_buffer2", "CUDA QSA pending-buffer snapshot after at most two request tokens.", + "T", OpSchema::Optional) + .Output(13, "mid_state_lengths2", "CUDA QSA state lengths after at most two request tokens.", + "M", OpSchema::Optional) + .Output(14, "mid_kv_buffer3", "CUDA QSA pending-buffer snapshot after at most three request tokens.", + "T", OpSchema::Optional) + .Output(15, "mid_state_lengths3", "CUDA QSA state lengths after at most three request tokens.", + "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain floating point tensors to float, float16 and bfloat16.") diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index cf678e02942b6..94f6081d8ad9e 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -583,7 +583,7 @@ TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) } TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecode_CUDA) { - for (int sequence_length : {1, 2}) { + for (int sequence_length : {1, 2, 3, 4}) { for (bool auxiliary : {false, true}) { SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + ", auxiliary=" + std::to_string(auxiliary)); @@ -603,7 +603,7 @@ TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeFloat16_CUDA) { } TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeVectorStagingChannels_CUDA) { - for (int sequence_length : {1, 2}) { + for (int sequence_length : {1, 2, 3, 4}) { for (bool auxiliary : {false, true}) { SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + ", auxiliary=" + std::to_string(auxiliary)); From 3b52694c9bc8c404a9e2aba9b39c853364fb5139 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:41:09 +0000 Subject: [PATCH 10/23] Apply lintrunner formatting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../sparse/packed_sparse_attention_indexer_impl.h | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h index 748b0d450fda8..f31099b28327e 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h @@ -39,13 +39,13 @@ struct PackedSparseAttentionIndexerParams { // policy_mode = "qsa" int block_topk = 0; // token_budget / compress_ratio int state_update_capacity = 0; - void* mid_kv_buffer = nullptr; // [batch, buffer_capacity, head_size] T - int32_t* mid_state_lengths = nullptr; // [batch, 2] - void* mid_kv_buffer2 = nullptr; // [batch, buffer_capacity, head_size] T - int32_t* mid_state_lengths2 = nullptr; // [batch, 2] - void* mid_kv_buffer3 = nullptr; // [batch, buffer_capacity, head_size] T - int32_t* mid_state_lengths3 = nullptr; // [batch, 2] - int32_t* validation_status = nullptr; // [batch] + void* mid_kv_buffer = nullptr; // [batch, buffer_capacity, head_size] T + int32_t* mid_state_lengths = nullptr; // [batch, 2] + void* mid_kv_buffer2 = nullptr; // [batch, buffer_capacity, head_size] T + int32_t* mid_state_lengths2 = nullptr; // [batch, 2] + void* mid_kv_buffer3 = nullptr; // [batch, buffer_capacity, head_size] T + int32_t* mid_state_lengths3 = nullptr; // [batch, 2] + int32_t* validation_status = nullptr; // [batch] // policy_mode = "csa" int index_topk = 0; From 5ed039c7df034c437dcdba295cccf8e0c4f4dd3f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 00:36:05 +0000 Subject: [PATCH 11/23] [CUDA] Revert unused QSA position_ids and WebGPU shader changes in PackedSparseAttentionIndexer Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 4 +- .../sparse/packed_sparse_attention_indexer.cc | 9 +-- .../bert/packed_sparse_attention_indexer.cc | 58 +++++-------------- .../bert/packed_sparse_attention_indexer.h | 6 +- .../core/graph/contrib_ops/bert_defs.cc | 9 +-- ...packed_sparse_attention_indexer_op_test.cc | 38 +----------- 6 files changed, 30 insertions(+), 94 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 279d718855ede..b58a830322183 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5300,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 unless optional position_ids are supplied. + 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 diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index e6493889d275e..87744046054d0 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -222,9 +222,9 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* } const bool position_ids_provided = psai::kPositionIds < context->InputCount() && context->Input(psai::kPositionIds) != nullptr; - ORT_RETURN_IF(!is_qsa && !position_ids_provided, - "PackedSparseAttentionIndexer: input ", psai::kPositionIds, - " (position_ids) is required for policy_mode 'csa'"); + ORT_RETURN_IF(position_ids_provided == is_qsa, + "PackedSparseAttentionIndexer: position_ids is required for " + "policy_mode 'csa' and must be omitted for policy_mode 'qsa'"); const bool gate_buffer_provided = psai::kPastGateBuffer < context->InputCount() && context->Input(psai::kPastGateBuffer) != nullptr; ORT_RETURN_IF(gate_buffer_provided != !is_qsa, "PackedSparseAttentionIndexer: input ", psai::kPastGateBuffer, @@ -316,9 +316,6 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* cont } ORT_RETURN_IF_ERROR(CheckShape(query_norm_weight, "query_norm_weight", {head_size})); ORT_RETURN_IF_ERROR(CheckShape(key_norm_weight, "key_norm_weight", {head_size})); - if (position_ids != nullptr) { - ORT_RETURN_IF_ERROR(CheckShape(position_ids, "position_ids", {total_tokens})); - } RotaryCacheShape rotary; ORT_RETURN_IF_ERROR(CheckRotaryCache(cos_cache, sin_cache, batch_size, rotary)); diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc index 41268aa3d9850..10e65756e3cfe 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -70,17 +70,13 @@ Status CheckShape(const Tensor* tensor, const char* name, std::initializer_list< Status CheckStorageBindingLimit(const onnxruntime::webgpu::ComputeContext& context, const char* program_name, - std::initializer_list tensors, - const Tensor* optional_tensor = nullptr) { + std::initializer_list tensors) { const auto& limits = context.DeviceLimits(); uint32_t binding_count = 0; for (const Tensor* tensor : tensors) { ORT_RETURN_IF(tensor == nullptr, "PackedSparseAttentionIndexer: missing tensor for ", program_name); binding_count += StorageBindingSegments(tensor, limits.maxStorageBufferBindingSize); } - if (optional_tensor != nullptr) { - binding_count += StorageBindingSegments(optional_tensor, limits.maxStorageBufferBindingSize); - } if (binding_count > limits.maxStorageBuffersPerShaderStage) { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "PackedSparseAttentionIndexer (WebGPU): ", program_name, " requires ", binding_count, @@ -394,10 +390,6 @@ Status PackedSparseAttentionIndexerQsaSelectProgram::GenerateShaderCode(ShaderHe const auto& present_key_state = shader.AddInput("present_key_state", ShaderUsage::UseUniform); const auto& rotary_cache = shader.AddInput("rotary_cache", ShaderUsage::UseUniform); const auto& sequence_metadata = shader.AddInput("sequence_metadata", ShaderUsage::UseUniform); - const ShaderVariableHelper* position_ids = nullptr; - if (has_position_ids_) { - position_ids = &shader.AddInput("position_ids", ShaderUsage::UseUniform); - } const auto& present_state_lengths = shader.AddInput("present_state_lengths", ShaderUsage::UseUniform); const auto& overflow_flags = shader.AddInput("overflow_flags", ShaderUsage::UseUniform); const auto& selected_indices = shader.AddOutput("selected_indices", ShaderUsage::UseUniform); @@ -420,24 +412,12 @@ Status PackedSparseAttentionIndexerQsaSelectProgram::GenerateShaderCode(ShaderHe << " let p = u32(position);\n" << " let cr = uniforms.compress_ratio;\n" << " return p / cr + select(0u, 1u, p % cr == cr - 1u);\n" - << "}\n"; - if (has_position_ids_) { - shader.AdditionalImplementation() - << "fn abs_position(token: u32, b: u32) -> i32 {\n" - << " let raw = " << position_ids->GetByOffset("token", true) << ";\n" - << " if ((raw.y & 0x80000000u) != 0u) { return -2147483648; }\n" - << " if (raw.y != 0u || raw.x > 2147483647u) { return 2147483647; }\n" - << " return i32(raw.x);\n" - << "}\n"; - } else { - shader.AdditionalImplementation() - << "fn abs_position(token: u32, b: u32) -> i32 {\n" - << " let req_start = " << sequence_metadata.GetByOffset("b") << ";\n" - << " return " << sequence_metadata.GetByOffset("uniforms.batch_size + 1u + b") - << " + (i32(token) - req_start);\n" - << "}\n"; - } - shader.AdditionalImplementation() + << "}\n" + << "fn abs_position(token: u32, b: u32) -> i32 {\n" + << " let req_start = " << sequence_metadata.GetByOffset("b") << ";\n" + << " return " << sequence_metadata.GetByOffset("uniforms.batch_size + 1u + b") + << " + (i32(token) - req_start);\n" + << "}\n" << "fn normalized_query_value(token: u32, head: u32, d: u32) -> f32 {\n" << " let base = token * uniforms.query_row_stride + head * uniforms.head_size;\n" << " var square_sum = 0.0;\n" @@ -937,9 +917,9 @@ Status PackedSparseAttentionIndexer::ComputeInternal(onnxruntime::webgpu::Comput } const bool position_ids_provided = psai::kPositionIds < context.InputCount() && context.Input(psai::kPositionIds) != nullptr; - ORT_RETURN_IF(!is_qsa && !position_ids_provided, + ORT_RETURN_IF(position_ids_provided == is_qsa, "PackedSparseAttentionIndexer: position_ids is required for " - "policy_mode 'csa'"); + "policy_mode 'csa' and must be omitted for policy_mode 'qsa'"); const bool gate_buffer_provided = psai::kPastGateBuffer < context.InputCount() && context.Input(psai::kPastGateBuffer) != nullptr; ORT_RETURN_IF(gate_buffer_provided != !is_qsa, "PackedSparseAttentionIndexer: past_gate_buffer ", @@ -967,7 +947,6 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont const Tensor* sin_cache = context.Input(psai::kSinCache); const Tensor* cu_seqlens = context.Input(psai::kCumulativeSequenceLengths); const Tensor* past_seqlens = context.Input(psai::kPastSequenceLengths); - const Tensor* position_ids = context.Input(psai::kPositionIds); const Tensor* past_key_state = context.Input(psai::kPastKeyState); const Tensor* past_kv_buffer = context.Input(psai::kPastKvBuffer); const Tensor* past_state_lengths = context.Input(psai::kPastStateLengths); @@ -1010,9 +989,6 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont } ORT_RETURN_IF_ERROR(CheckShape(query_norm, "query_norm_weight", {head_size})); ORT_RETURN_IF_ERROR(CheckShape(norm, "key_norm_weight", {head_size})); - if (position_ids != nullptr) { - ORT_RETURN_IF_ERROR(CheckShape(position_ids, "position_ids", {total_tokens})); - } RotaryCacheShape rotary; ORT_RETURN_IF_ERROR(CheckRotaryCache(cos_cache, sin_cache, batch_size, rotary)); @@ -1175,20 +1151,16 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont ORT_RETURN_IF_ERROR(CheckStorageBindingLimit( context, "QSA selection", {&query_metadata, present_key_state, &rotary_cache, &sequence_metadata, present_state_lengths, &overflow_flags, - selected_indices, selected_counts}, - position_ids)); - PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched, position_ids != nullptr}; - select.CacheHint(rotary.batched, position_ids != nullptr) + selected_indices, selected_counts})); + PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched}; + select.CacheHint(rotary.batched) .SetWorkgroupSize(kWorkgroupSize) .AddInputs({{&query_metadata, ProgramTensorMetadataDependency::Type}, {present_key_state, ProgramTensorMetadataDependency::Type}, {&rotary_cache, ProgramTensorMetadataDependency::Type}, - {&sequence_metadata, ProgramTensorMetadataDependency::Type}}); - if (position_ids != nullptr) { - select.AddInput({position_ids, ProgramTensorMetadataDependency::Type}); - } - select.AddInputs({{present_state_lengths, ProgramTensorMetadataDependency::Type}, - {&overflow_flags, ProgramTensorMetadataDependency::Type}}) + {&sequence_metadata, ProgramTensorMetadataDependency::Type}, + {present_state_lengths, ProgramTensorMetadataDependency::Type}, + {&overflow_flags, ProgramTensorMetadataDependency::Type}}) .AddOutputs({{selected_indices, ProgramTensorMetadataDependency::Type}, {selected_counts, ProgramTensorMetadataDependency::Type}}) .SetDispatchGroupSize(ToUint32(total_tokens)) diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h index 353cc0413a813..69f590b73a2bd 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.h @@ -65,10 +65,9 @@ class PackedSparseAttentionIndexerQsaUpdateProgram final class PackedSparseAttentionIndexerQsaSelectProgram final : public Program { public: - PackedSparseAttentionIndexerQsaSelectProgram(bool cos_cache_batched, bool has_position_ids) + explicit PackedSparseAttentionIndexerQsaSelectProgram(bool cos_cache_batched) : Program{"PackedSparseAttentionIndexerQsaSelect"}, - cos_cache_batched_{cos_cache_batched}, - has_position_ids_{has_position_ids} {} + cos_cache_batched_{cos_cache_batched} {} Status GenerateShaderCode(ShaderHelper& shader) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( {"total_tokens", ProgramUniformVariableDataType::Uint32}, @@ -88,7 +87,6 @@ class PackedSparseAttentionIndexerQsaSelectProgram final private: bool cos_cache_batched_; - bool has_position_ids_; }; // One invocation per request: closes every new compression window (softmax-gated pool -> RMSNorm diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 4babb09090985..f5bd16af06cfb 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2897,9 +2897,10 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference : " is required when policy_mode is 'csa'"); } } - if (!is_qsa && !PackedSparseAttentionIndexerHasInput(ctx, psai::kPositionIds)) { + if (PackedSparseAttentionIndexerHasInput(ctx, psai::kPositionIds) == is_qsa) { fail_shape_inference("PackedSparseAttentionIndexer: input ", psai::kPositionIds, - " (position_ids) is required when policy_mode is 'csa'"); + is_qsa ? " (position_ids) must be omitted when policy_mode is 'qsa'" + : " (position_ids) is required when policy_mode is 'csa'"); } const bool has_capture_count = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateCaptureCount); const bool has_state_update_active = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateActive); @@ -3256,8 +3257,8 @@ request's own packed token range and fixed-capacity state slice: 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 unless optional position_ids are supplied. + 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 diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index af40c5f2c2979..6ff7c30859c59 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -501,13 +501,11 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsQsaWithCsaInput) { "must be omitted when policy_mode is 'qsa'"); } -TEST(PackedSparseAttentionIndexerShapeInferenceTest, QsaAcceptsPositionIds) { +TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsQsaWithPositionIds) { GraphOptions options; options.add_position_ids = true; - std::unique_ptr model; - ASSERT_STATUS_OK(BuildAndResolve([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, model)); - const auto& node = *model->MainGraph().Nodes().begin(); - EXPECT_TRUE(node.InputDefs()[psai::kPositionIds]->Exists()); + ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, + "must be omitted when policy_mode is 'qsa'"); } TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsCsaMissingPositionIds) { @@ -1466,18 +1464,6 @@ TEST(PackedSparseAttentionIndexerTest, QsaMalformedStateLengthsAreClamped) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(std::move(problem))); } -TEST(PackedSparseAttentionIndexerTest, QsaExplicitPositionIds) { - auto problem = MakeQsaPackedProblem(); - problem.position_ids.resize(static_cast(problem.TotalTokens())); - for (int batch = 0; batch < problem.batch_size; ++batch) { - const int start = problem.cumulative_sequence_lengths[batch]; - for (int token = start; token < problem.cumulative_sequence_lengths[batch + 1]; ++token) { - problem.position_ids[static_cast(token)] = problem.past_sequence_lengths[batch] + token - start + 1; - } - } - RunQsaPackedTest(1.0e-5f, std::move(problem), ProviderKind::Cuda); -} - TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalTileBoundariesPreserveStableTies) { for (int live_blocks : {0, 31, 32, 33, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049, 2051}) { SCOPED_TRACE(live_blocks); @@ -1772,24 +1758,6 @@ TEST(PackedSparseAttentionIndexerWebGpuTest, QsaPortableStorageBindingLimit) { RunQsaPackedTest(1.0e-5f, MakeQsaPackedProblem(), ProviderKind::WebGpu, nullptr, false, &config_options); } -TEST(PackedSparseAttentionIndexerWebGpuTest, QsaExplicitPositionsStorageBindingLimit) { - auto problem = MakeQsaPackedProblem(); - problem.position_ids.resize(static_cast(problem.TotalTokens())); - for (int b = 0; b < problem.batch_size; ++b) { - const int start = problem.cumulative_sequence_lengths[b]; - const int end = problem.cumulative_sequence_lengths[b + 1]; - for (int token = start; token < end; ++token) { - problem.position_ids[token] = problem.past_sequence_lengths[b] + token - start; - } - } - for (const std::string limit : {"8", "9"}) { - ConfigOptions config_options; - ASSERT_STATUS_OK(config_options.AddConfigEntry(webgpu::options::kMaxStorageBuffersPerShaderStage, limit.c_str())); - RunQsaPackedTest(1.0e-5f, problem, ProviderKind::WebGpu, nullptr, false, &config_options, 0, - limit == "8" ? "QSA selection requires 9 storage-buffer bindings" : ""); - } -} - TEST(PackedSparseAttentionIndexerWebGpuTest, QsaAllEmptyBatchPreservesState) { QsaPackedProblem problem; problem.batch_size = 1; From 6511e8a33a67971645f576aa68773a37bef0cf97 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 04:10:02 +0000 Subject: [PATCH 12/23] fix(cuda): fix shared-provider include order, Windows macro clash, test range, and docs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 47 ++++++++++--------- .../bert/dynamic_sparse_attention_impl.cu | 5 +- .../sparse/packed_sparse_attention_indexer.cc | 4 +- .../contrib_ops/cuda/sparse/qsa_topk.cuh | 12 ++--- ...packed_sparse_attention_indexer_op_test.cc | 2 +- 5 files changed, 37 insertions(+), 33 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index b58a830322183..5eb6f6df1d880 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1789,7 +1789,7 @@ 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. 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, @@ -1803,35 +1803,35 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
attention_mode : string (default is selected_only)
+
attention_mode : string
One of 'selected_only' or 'local_plus_selected'.
-
auxiliary_kv_shared : int (default is 0)
+
auxiliary_kv_shared : int
Use auxiliary_key as both key and value when auxiliary_value is omitted.
-
do_rotary : int (default is 0)
+
do_rotary : int
Whether to apply rotary embedding to Q and newly appended K.
-
is_causal : int (default is 1)
+
is_causal : int
Must be 1. DynamicSparseAttention version 1 supports causal attention only.
-
kv_num_heads : int
+
kv_num_heads : int (required)
Number of main and auxiliary KV heads.
-
local_window_size : int (default is -1)
+
local_window_size : int
Number of causally visible main-cache entries in local_plus_selected mode.
-
num_heads : int
+
num_heads : int (required)
Number of query heads.
-
qk_norm_epsilon : float (default is 0.0)
+
qk_norm_epsilon : float
Epsilon for optional per-head Q/K RMS normalization.
-
rotary_interleaved : int (default is 0)
+
rotary_interleaved : int
Whether rotary pairs use interleaved layout.
-
rotary_offset : int (default is 0)
+
rotary_offset : int
First head channel covered by rotary embedding. Must be nonnegative, a multiple of 8, and leave room for the rotary dimensions within head_size.
scale : float
Scaling factor applied to QK. Defaults to 1/sqrt(head_size).
-
selected_kv_source : string (default is main)
+
selected_kv_source : string
Source addressed by selected indices: 'main' or 'auxiliary'.
-
smooth_softmax : int (default is 0)
+
smooth_softmax : int
Add a zero-valued sink logit when no explicit head_sink is supplied.
-#### Inputs (5 - 20) +#### Inputs (11 - 20)
query : T
@@ -5344,31 +5344,31 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
compress_ratio : int
+
compress_ratio : int (required)
Number of consecutive tokens folded into one compressed/pooled entry. Must be > 0 and 2 * compress_ratio - 1 must not exceed INT_MAX.
-
epsilon : float (default is 0.0)
+
epsilon : float
Epsilon of the RMS normalization applied to queries and compressed keys. Default is 1e-6.
head_weight_scale : float
Only for policy_mode 'csa': scale applied to head_weights. Default is 1/sqrt(num_heads). Must be omitted when policy_mode is 'qsa'.
index_topk : int
Only for policy_mode 'csa': number of compressed entries selected per query. Must be > 0. Must be omitted when policy_mode is 'qsa'.
-
indexshare_state_policy : string (default is default)
+
indexshare_state_policy : string
Reuse state policy: default holds QSA and advances CSA; hold or advance overrides it.
-
max_output_entries : int (default is 0)
+
max_output_entries : int
Fixed padded output width for the optional IndexShare interface.
-
policy_mode : string
+
policy_mode : string (required)
Indexer policy. Must be exactly 'qsa' (token indexer) or 'csa' (compressed block indexer).
scale : float
Scale applied to the per-head ReLU scores. Default is 1/sqrt(head_size).
-
state_capacity : int
+
state_capacity : int (required)
Fixed capacity (number of entries) of past_key_state / present_key_state. Must be > 0.
-
state_update_capacity : int (default is 0)
+
state_update_capacity : int
Only for policy_mode 'qsa': maximum number of leading token transitions captured per request. Must be in [0, 8]. Default is 0.
token_budget : int
Only for policy_mode 'qsa': maximum number of tokens selected from complete blocks. Must be > 0 and divisible by compress_ratio. Must be omitted when policy_mode is 'csa'.
-#### Inputs (10 - 22) +#### Inputs (16 - 22)
query : T
@@ -5417,7 +5417,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Merged device status [N].
-#### Outputs (5 - 16) +#### Outputs (6 - 16)
selected_indices : M
@@ -8985,3 +8985,4 @@ No versioning maintained for experimental ops.
Constrain input and output types to float32 tensors.
+ 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 dfc0c575a72da..9e6673ef94fae 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -15,9 +15,12 @@ #include "contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh" #include "core/common/safeint.h" -#include "core/platform/env_var_utils.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 { namespace contrib { namespace cuda { diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index 87744046054d0..ce211a88ee46c 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -223,8 +223,8 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* const bool position_ids_provided = psai::kPositionIds < context->InputCount() && context->Input(psai::kPositionIds) != nullptr; ORT_RETURN_IF(position_ids_provided == is_qsa, - "PackedSparseAttentionIndexer: position_ids is required for " - "policy_mode 'csa' and must be omitted for policy_mode 'qsa'"); + "PackedSparseAttentionIndexer: position_ids is required for " + "policy_mode 'csa' and must be omitted for policy_mode 'qsa'"); const bool gate_buffer_provided = psai::kPastGateBuffer < context->InputCount() && context->Input(psai::kPastGateBuffer) != nullptr; ORT_RETURN_IF(gate_buffer_provided != !is_qsa, "PackedSparseAttentionIndexer: input ", psai::kPastGateBuffer, diff --git a/onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh b/onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh index a178940858f97..6b8133c78b54c 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh +++ b/onnxruntime/contrib_ops/cuda/sparse/qsa_topk.cuh @@ -37,9 +37,9 @@ __device__ __forceinline__ void SortChunk(const uint64_t* input, uint64_t* outpu __global__ void QsaCompactTileTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, int input_stride, int output_stride) { __shared__ union { - ChunkSort<1>::TempStorage small; - ChunkSort<2>::TempStorage medium; - ChunkSort<4>::TempStorage large; + ChunkSort<1>::TempStorage s1; + ChunkSort<2>::TempStorage s2; + ChunkSort<4>::TempStorage s4; } temp; const int chunk_start = static_cast(blockIdx.x) * kChunkSize; const uint64_t* input = input_keys + static_cast(blockIdx.y) * input_stride + chunk_start; @@ -51,11 +51,11 @@ __global__ void QsaCompactTileTopKKernel(const uint64_t* input_keys, uint64_t* o if (input[0] == topk::kPaddingSortKey) { output[threadIdx.x] = topk::kPaddingSortKey; } else if (count > 1024 && input[1024] != topk::kPaddingSortKey) { - SortChunk<4>(input, output, count, temp.large); + SortChunk<4>(input, output, count, temp.s4); } else if (count > 512 && input[512] != topk::kPaddingSortKey) { - SortChunk<2>(input, output, min(count, 1024), temp.medium); + SortChunk<2>(input, output, min(count, 1024), temp.s2); } else { - SortChunk<1>(input, output, min(count, 512), temp.small); + SortChunk<1>(input, output, min(count, 512), temp.s1); } } diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 6ff7c30859c59..921f846112201 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -558,7 +558,7 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsWrongOutputCount) { GraphOptions options; options.output_count = 4; ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, - "output size 4 not in range [min=6, max=12]"); + "output size 4 not in range [min=6, max=16]"); } TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsStateUpdateCapacityWithoutCaptureCount) { From a523af45d70f4b0b830fded4d2529aa59489abb9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 21:06:27 +0000 Subject: [PATCH 13/23] Optimize GroupedDecodeDynamicSparseAttentionKernel for 12 query heads per KV head - Specialize GroupedDecodeDynamicSparseAttentionKernel with 12 warps (384 threads, head_tiles=1) when num_heads/kv_num_heads == 12, halving HBM sparse KV reads and eliminating idle warps in tile 1 (-37.4% kernel time at S=8). - Raise kTargetGroupedQuantizedDecodeBlocks to 512 and trim unused mma shared memory when !(T==half && !kQuantizedCache). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../bert/dynamic_sparse_attention_impl.cu | 126 ++++++++++++------ 1 file changed, 87 insertions(+), 39 deletions(-) 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 9e6673ef94fae..2fce55999d618 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -37,7 +37,7 @@ constexpr int kGroupedDecodeQueryHeadsPerBlock = 8; constexpr int kGroupedDecodeCandidatesPerSplit = 64; constexpr int kTargetGroupedDecodeBlocks = 128; constexpr int kGroupedQuantizedDecodeCandidatesPerSplit = 32; -constexpr int kTargetGroupedQuantizedDecodeBlocks = 256; +constexpr int kTargetGroupedQuantizedDecodeBlocks = 512; template __device__ __forceinline__ float ToFloat(T value) { @@ -995,8 +995,9 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( __shared__ int candidate_indices[kGroupedPrefillCandidatesPerTile]; __shared__ bool candidate_valid[kGroupedPrefillCandidatesPerTile]; __shared__ bool candidate_sources[kGroupedPrefillCandidatesPerTile]; - __shared__ half mma_query[kGroupedDecodeQueryHeadsPerBlock * 256]; - __shared__ float mma_scores[kGroupedDecodeQueryHeadsPerBlock * 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]{}; @@ -1755,6 +1756,11 @@ bool UseGroupedDecodeAttention(const DynamicSparseAttentionParameters& parameter 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; @@ -1763,7 +1769,7 @@ int GetGroupedDecodeSplitCount(const DynamicSparseAttentionParameters& parameter (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, kGroupedDecodeQueryHeadsPerBlock); + GetGroupedHeadTiles(parameters, GetGroupedDecodeHeadsPerBlock(parameters)); const int64_t occupancy_splits = std::max(1, (target_blocks + grouped_head_blocks - 1) / grouped_head_blocks); @@ -2025,9 +2031,10 @@ Status LaunchDynamicSparseAttention( const size_t fused_shared_bytes = (static_cast(candidate_capacity) + static_cast(fused_threads)) * sizeof(float) + static_cast(parameters.head_size) * sizeof(T); - constexpr int kGroupedDecodeThreads = 32 * kGroupedDecodeQueryHeadsPerBlock; + const int grouped_decode_heads_per_block = GetGroupedDecodeHeadsPerBlock(parameters); + const int grouped_decode_threads = 32 * grouped_decode_heads_per_block; constexpr int kGroupedPrefillThreads = 32 * kGroupedPrefillQueryHeadsPerBlock; - const bool grouped_decode_threads_supported = kGroupedDecodeThreads <= max_threads_per_block; + const bool grouped_decode_threads_supported = grouped_decode_threads <= max_threads_per_block; const bool grouped_prefill_threads_supported = kGroupedPrefillThreads <= max_threads_per_block; const bool use_grouped_decode = grouped_decode_threads_supported && @@ -2060,7 +2067,7 @@ Status LaunchDynamicSparseAttention( const int64_t partial_blocks = query_blocks * split_count; const int64_t grouped_blocks = static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads * - GetGroupedHeadTiles(parameters, kGroupedDecodeQueryHeadsPerBlock) * split_count; + GetGroupedHeadTiles(parameters, grouped_decode_heads_per_block) * split_count; ORT_RETURN_IF_ERROR(CheckBlockCount(partial_blocks, "grouped decode partials")); ORT_RETURN_IF_ERROR(CheckBlockCount(grouped_blocks, "grouped decode attention")); ORT_RETURN_IF_NOT(data.attention_workspace != nullptr, @@ -2075,10 +2082,11 @@ Status LaunchDynamicSparseAttention( ParseEnvironmentVariableWithDefault("ORT_DSA_INT8_VECTOR_STAGING", 1); ORT_ENFORCE(vectorized_staging == 0 || vectorized_staging == 1, "ORT_DSA_INT8_VECTOR_STAGING must be 0 or 1."); - const auto launch = [&](auto vectorized) { - GroupedDecodeDynamicSparseAttentionKernel - <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( + <<(grouped_blocks), 32 * kHeads, grouped_shared_bytes, stream>>>( data.prepared_query, data.present_key, data.present_value, data.auxiliary_key, data.auxiliary_value, data.selected_indices, data.selected_counts, data.seqlens_k, data.head_sink, @@ -2091,39 +2099,79 @@ Status LaunchDynamicSparseAttention( parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, parameters.use_smooth_softmax, data.key_scale, data.value_scale); }; - if (vectorized_staging) { - launch(std::true_type{}); + if (grouped_decode_heads_per_block == 12) { + if (vectorized_staging) { + launch(std::integral_constant{}, std::true_type{}); + } else { + launch(std::integral_constant{}, std::false_type{}); + } } else { - launch(std::false_type{}); + if (vectorized_staging) { + launch(std::integral_constant{}, std::true_type{}); + } else { + launch(std::integral_constant{}, std::false_type{}); + } } } else if (parameters.cache_capacity <= 131072) { - GroupedDecodeDynamicSparseAttentionKernel - <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( - data.prepared_query, data.present_key, data.present_value, - data.auxiliary_key, data.auxiliary_value, data.selected_indices, - data.selected_counts, data.seqlens_k, data.head_sink, - partial_max, partial_sum, partial_output, - static_cast(grouped_blocks), split_count, parameters.sequence_length, - parameters.num_heads, parameters.kv_num_heads, parameters.head_size, - parameters.cache_capacity, parameters.auxiliary_sequence_length, - parameters.max_selected, parameters.local_window_size, parameters.scale, - local_plus_selected, - parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, - parameters.use_smooth_softmax, nullptr, nullptr); + if (grouped_decode_heads_per_block == 12) { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), 32 * 12, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, nullptr, nullptr); + } else { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), grouped_decode_threads, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, nullptr, nullptr); + } } else { - GroupedDecodeDynamicSparseAttentionKernel - <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( - data.prepared_query, data.present_key, data.present_value, - data.auxiliary_key, data.auxiliary_value, data.selected_indices, - data.selected_counts, data.seqlens_k, data.head_sink, - partial_max, partial_sum, partial_output, - static_cast(grouped_blocks), split_count, parameters.sequence_length, - parameters.num_heads, parameters.kv_num_heads, parameters.head_size, - parameters.cache_capacity, parameters.auxiliary_sequence_length, - parameters.max_selected, parameters.local_window_size, parameters.scale, - local_plus_selected, - parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, - parameters.use_smooth_softmax, nullptr, nullptr); + if (grouped_decode_heads_per_block == 12) { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), 32 * 12, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, nullptr, nullptr); + } else { + GroupedDecodeDynamicSparseAttentionKernel + <<(grouped_blocks), grouped_decode_threads, grouped_shared_bytes, stream>>>( + data.prepared_query, data.present_key, data.present_value, + data.auxiliary_key, data.auxiliary_value, data.selected_indices, + data.selected_counts, data.seqlens_k, data.head_sink, + partial_max, partial_sum, partial_output, + static_cast(grouped_blocks), split_count, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.cache_capacity, parameters.auxiliary_sequence_length, + parameters.max_selected, parameters.local_window_size, parameters.scale, + local_plus_selected, + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, + parameters.use_smooth_softmax, nullptr, nullptr); + } } CUDA_RETURN_IF_ERROR(cudaGetLastError()); From 61c5d87f03c18341a4dafe1b9d5ffccacd67f211 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 05:01:24 +0000 Subject: [PATCH 14/23] Remove legacy mid-state snapshot outputs from PackedSparseAttentionIndexer - Remove unused outputs 9..15 (mid_kv_buffer*, mid_state_lengths*, validation_status) from PackedSparseAttentionIndexer and rely on main's state_update (output 6) and state_status (output 8). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 18 +----- docs/OperatorKernels.md | 2 +- .../packed_sparse_attention_indexer.md | 11 ---- .../packed_sparse_attention_indexer_common.h | 10 +-- .../sparse/packed_sparse_attention_indexer.cc | 43 ------------- .../packed_sparse_attention_indexer_impl.cu | 38 ------------ .../packed_sparse_attention_indexer_impl.h | 7 --- .../bert/packed_sparse_attention_indexer.cc | 5 -- .../core/graph/contrib_ops/bert_defs.cc | 61 +------------------ ...packed_sparse_attention_indexer_op_test.cc | 48 ++------------- 10 files changed, 11 insertions(+), 232 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 5eb6f6df1d880..acb492a5c8174 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5285,8 +5285,6 @@ This version of the operator has been available since version 1 of the 'com.micr (the optional state_status output reports which requests were rejected and why); * additionally emits selected_counts, the exact number of active (non -1) entries per query, so that no downstream consumer needs to scan selected_indices for its query's true count. - * for QSA, optionally emits the pending-state snapshot after the first request token and a - per-request validation status for overflow or inconsistent state metadata. Both policy_mode values keep the semantics of SparseAttentionIndexer, applied independently to each request's own packed token range and fixed-capacity state slice: @@ -5417,7 +5415,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Merged device status [N].
-#### Outputs (6 - 16) +#### Outputs (6 - 9)
selected_indices : M
@@ -5438,20 +5436,6 @@ This version of the operator has been available since version 1 of the 'com.micr
Selection status [N], forwarded on reuse and zero on refresh.
state_status (optional) : M
Device-side per-request state update status with shape (batch_size). 0: the request's state update was applied (or intentionally held by IndexShare reuse); 1: invalid device metadata (cumulative_sequence_lengths, past_sequence_lengths or past_state_lengths are inconsistent); 2: the update would exceed state_capacity. A non-zero request leaves its state payloads unchanged, clamps out-of-range past_state_lengths into present_state_lengths, and emits an empty selection (selected_counts is 0 and selected_indices is -1). Intended for CUDA graph replay, where the host checks it after the replay instead of synchronizing inside the operator. Currently supported by the CUDA provider only.
-
mid_kv_buffer (optional) : T
-
Only for CUDA policy_mode 'qsa' refresh: pending-buffer snapshot after the first token of each request. Not supported in IndexShare reuse mode.
-
mid_state_lengths (optional) : M
-
Only for CUDA policy_mode 'qsa' refresh: state lengths after the first token of each request. Not supported in IndexShare reuse mode.
-
validation_status (optional) : M
-
Only for CUDA policy_mode 'qsa' refresh: per-request overflow status; two-token rejection is encoded as 0x99. Not supported in IndexShare reuse mode.
-
mid_kv_buffer2 (optional) : T
-
CUDA QSA pending-buffer snapshot after at most two request tokens.
-
mid_state_lengths2 (optional) : M
-
CUDA QSA state lengths after at most two request tokens.
-
mid_kv_buffer3 (optional) : T
-
CUDA QSA pending-buffer snapshot after at most three request tokens.
-
mid_state_lengths3 (optional) : M
-
CUDA QSA state lengths after at most three request tokens.
#### Type Constraints diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 920cad5330595..cd0be24517121 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -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_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**
*out* mid_kv_buffer:**T**
*out* mid_state_lengths:**M**
*out* validation_status:**M**
*out* mid_kv_buffer2:**T**
*out* mid_state_lengths2:**M**
*out* mid_kv_buffer3:**T**
*out* mid_state_lengths3:**M**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**T_KEY** = 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/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md index 6825fb40355d0..300453da899eb 100644 --- a/docs/contrib_ops/packed_sparse_attention_indexer.md +++ b/docs/contrib_ops/packed_sparse_attention_indexer.md @@ -121,20 +121,9 @@ empty output name for `qsa`; `state_update` is declared only when QSA capture is | 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 | -| 9 | `mid_kv_buffer` | same shape as `past_kv_buffer` | T | optional CUDA qsa snapshot | -| 10 | `mid_state_lengths` | same shape as `past_state_lengths` | int32 | optional CUDA qsa snapshot | -| 11 | `validation_status` | `(batch_size)` | int32 | optional CUDA qsa status | 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. -`mid_kv_buffer` and `mid_state_lengths` must be connected together and capture the state -after each request's first token. `validation_status` is zero on success; rejected -requests report 9, or `0x99` for requests containing at least two tokens. - -Output 8 retains main's `state_status` contract. Graphs built against the earlier -experimental snapshot layout must insert an empty output at slot 8 before the -snapshot outputs, which now occupy slots 9-11. Existing graphs using only the -mainline outputs do not require changes. `capacity` is `token_budget + compress_ratio - 1` for `qsa` and `index_topk` for `csa`, exactly the same formula (`SelectedCapacity`) used by `SparseAttentionIndexer`. diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h index 038c72be89811..81bb28e99ab9a 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h @@ -71,16 +71,8 @@ enum OutputIndex : int { kPresentStateLengths = 5, // [batch_size, 2], int32 kStateUpdate = 6, // qsa only, optional: [batch_size, state_update_capacity, head_size] kIndexShareStatus = 7, - kStateStatus = 8, // optional: [batch_size], int32, see StateStatusCode - kMidKvBuffer = 9, // qsa only, optional: snapshot after the first request token - kMidStateLengths = 10, // qsa only, optional: lengths after the first request token - kValidationStatus = 11, // qsa only, optional: per-request validation status - kMidKvBuffer2 = 12, - kMidStateLengths2 = 13, - kMidKvBuffer3 = 14, - kMidStateLengths3 = 15, + kStateStatus = 8, // optional: [batch_size], int32, see StateStatusCode kOutputCount = 9, - kMaxOutputCount = 16, }; // Values of the optional state_status output. They follow the status convention of diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc index ce211a88ee46c..d6008de18268f 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -365,46 +365,16 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* cont ? context->Output( psai::kStateUpdate, TensorShape({batch_size, state_update_capacity_, head_size})) : nullptr; - Tensor* mid_kv_buffer = - context->OutputCount() > psai::kMidKvBuffer - ? context->Output(psai::kMidKvBuffer, TensorShape({batch_size, buffer_capacity, head_size})) - : nullptr; - Tensor* mid_state_lengths = - context->OutputCount() > psai::kMidStateLengths - ? context->Output(psai::kMidStateLengths, TensorShape({batch_size, psai::kStateLengthColumns})) - : nullptr; - Tensor* validation_status = - context->OutputCount() > psai::kValidationStatus - ? context->Output(psai::kValidationStatus, TensorShape({batch_size})) - : nullptr; ORT_RETURN_IF(selected_indices == nullptr || selected_counts == nullptr || present_key_state == nullptr || present_kv_buffer == nullptr || present_state_lengths == nullptr, "PackedSparseAttentionIndexer: policy_mode 'qsa' requires selected_indices, selected_counts, " "present_key_state, present_kv_buffer and present_state_lengths outputs"); - ORT_RETURN_IF((mid_kv_buffer == nullptr) != (mid_state_lengths == nullptr), - "PackedSparseAttentionIndexer: mid_kv_buffer and mid_state_lengths must be connected together"); if (state_update != nullptr && state_update->SizeInBytes() > 0) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(state_update->MutableDataRaw(), 0, state_update->SizeInBytes(), Stream(context))); } PackedSparseAttentionIndexerParams params; - Tensor* mid_buffers[2]{}; - Tensor* mid_lengths[2]{}; - for (int prefix = 0; prefix < 2; ++prefix) { - const int buffer_index = psai::kMidKvBuffer2 + prefix * 2; - const int length_index = buffer_index + 1; - if (context->OutputCount() > buffer_index) - mid_buffers[prefix] = context->Output(buffer_index, TensorShape({batch_size, buffer_capacity, head_size})); - if (context->OutputCount() > length_index) - mid_lengths[prefix] = context->Output(length_index, TensorShape({batch_size, psai::kStateLengthColumns})); - ORT_RETURN_IF((mid_buffers[prefix] == nullptr) != (mid_lengths[prefix] == nullptr), - "PackedSparseAttentionIndexer: every prefix buffer requires its matching lengths"); - } - params.mid_kv_buffer2 = mid_buffers[0] ? mid_buffers[0]->MutableData() : nullptr; - params.mid_state_lengths2 = mid_lengths[0] ? mid_lengths[0]->MutableData() : nullptr; - params.mid_kv_buffer3 = mid_buffers[1] ? mid_buffers[1]->MutableData() : nullptr; - params.mid_state_lengths3 = mid_lengths[1] ? mid_lengths[1]->MutableData() : nullptr; params.batch_size = static_cast(batch_size); params.total_tokens = static_cast(total_tokens); params.num_heads = static_cast(num_heads); @@ -424,12 +394,6 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* cont params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size)); params.block_topk = static_cast(token_budget_ / compress_ratio_); params.state_update_capacity = static_cast(state_update_capacity_); - params.mid_kv_buffer = - mid_kv_buffer == nullptr ? nullptr : reinterpret_cast(mid_kv_buffer->MutableData()); - params.mid_state_lengths = - mid_state_lengths == nullptr ? nullptr : mid_state_lengths->MutableData(); - params.validation_status = - validation_status == nullptr ? nullptr : validation_status->MutableData(); auto float_workspace = GetScratchBuffer(GetQsaPackedWorkspaceFloatCount(params), GetComputeStream(context)); Tensor* state_status = GetStateStatusOutput(context, batch_size); @@ -667,13 +631,6 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* cont template Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { - ORT_RETURN_IF(context->OutputCount() > psai::kMidKvBuffer2, - "PackedSparseAttentionIndexer: additional prefix outputs require refresh, not IndexShare reuse"); - for (int index : {psai::kMidKvBuffer, psai::kMidStateLengths, psai::kValidationStatus}) { - Tensor* output = context->OutputCount() > index ? context->Output(index, TensorShape({0})) : nullptr; - ORT_RETURN_IF(output != nullptr, - "PackedSparseAttentionIndexer: outputs 9-11 are not supported in IndexShare reuse mode"); - } const Tensor* query = context->Input(psai::kQuery); const Tensor* lengths = context->Input(psai::kPastSequenceLengths); const Tensor* indices = context->Input(psai::kMergedIndices); diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 6578cfbf1216c..c8e0800fba77c 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -200,44 +200,6 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con overflow_flags[b] = invalid_metadata ? psai::kStateStatusInvalidMetadata : rejected ? psai::kStateStatusCapacityOverflow : psai::kStateStatusOk; - if (params.validation_status != nullptr) { - const int request_length = req_end - req_start; - params.validation_status[b] = !rejected ? 0 : (request_length >= 2 ? 0x99 : 9); - } - } - - void* snapshots[3] = {params.mid_kv_buffer, params.mid_kv_buffer2, params.mid_kv_buffer3}; - int32_t* length_snapshots[3] = {params.mid_state_lengths, params.mid_state_lengths2, params.mid_state_lengths3}; -#pragma unroll - for (int prefix = 0; prefix < 3; ++prefix) { - if (snapshots[prefix] == nullptr || length_snapshots[prefix] == nullptr) continue; - const int mid_req = min(req_len, prefix + 1); - const int mid_pending = old_buf_len + mid_req; - const int mid_blocks = rejected ? 0 : mid_pending / params.compress_ratio; - const int mid_buf_len = rejected ? old_buf_len : mid_pending % params.compress_ratio; - if (threadIdx.x == 0) { - length_snapshots[prefix][b * 2 + psai::kKeyStateLength] = old_key_len + mid_blocks; - length_snapshots[prefix][b * 2 + psai::kBufferLength] = mid_buf_len; - } - T* mid_buffer = static_cast(snapshots[prefix]); - const int mid_total = params.buffer_capacity * params.head_size; - for (int i = static_cast(threadIdx.x); i < mid_total; i += static_cast(blockDim.x)) { - const int t = i / params.head_size; - const int d = i - t * params.head_size; - const int virtual_pos = mid_blocks * params.compress_ratio + t; - const int64_t slot = static_cast(b) * params.buffer_capacity * params.head_size + i; - if (t >= mid_buf_len || rejected) { - mid_buffer[slot] = past_kv_buffer[slot]; - } else { - mid_buffer[slot] = - virtual_pos < old_buf_len - ? past_kv_buffer[(static_cast(b) * params.buffer_capacity + virtual_pos) * - params.head_size + - d] - : key[(static_cast(req_start) + (virtual_pos - old_buf_len)) * params.key_row_stride + - d]; - } - } } for (int t = static_cast(threadIdx.x); t < capture_count; t += static_cast(blockDim.x)) { diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h index f31099b28327e..0d2f49ada9f88 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h @@ -39,13 +39,6 @@ struct PackedSparseAttentionIndexerParams { // policy_mode = "qsa" int block_topk = 0; // token_budget / compress_ratio int state_update_capacity = 0; - void* mid_kv_buffer = nullptr; // [batch, buffer_capacity, head_size] T - int32_t* mid_state_lengths = nullptr; // [batch, 2] - void* mid_kv_buffer2 = nullptr; // [batch, buffer_capacity, head_size] T - int32_t* mid_state_lengths2 = nullptr; // [batch, 2] - void* mid_kv_buffer3 = nullptr; // [batch, buffer_capacity, head_size] T - int32_t* mid_state_lengths3 = nullptr; // [batch, 2] - int32_t* validation_status = nullptr; // [batch] // policy_mode = "csa" int index_topk = 0; diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc index 10e65756e3cfe..99d36d449f504 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -896,11 +896,6 @@ PackedSparseAttentionIndexer::PackedSparseAttentionIndexer(const OpKernelInfo& i Status PackedSparseAttentionIndexer::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { ORT_RETURN_IF(context.InputCount() > psai::kIndexShareMode && context.Input(psai::kIndexShareMode), "PackedSparseAttentionIndexer IndexShare requires the CUDA provider"); - for (int index : {psai::kMidKvBuffer, psai::kMidStateLengths, psai::kValidationStatus}) { - Tensor* output = context.OutputCount() > index ? context.Output(index, TensorShape({0})) : nullptr; - ORT_RETURN_IF(output != nullptr, - "PackedSparseAttentionIndexer: mid-state and validation outputs require the CUDA provider"); - } const Tensor* query = context.Input(psai::kQuery); const Tensor* key_state = context.Input(psai::kPastKeyState); ORT_RETURN_IF(query != nullptr && key_state != nullptr && query->DataType() != key_state->DataType(), diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index f5bd16af06cfb..768f88cde070f 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2905,15 +2905,6 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference const bool has_capture_count = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateCaptureCount); const bool has_state_update_active = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateActive); const bool has_state_update_output = ctx.hasOutput(psai::kStateUpdate); - const bool has_mid_kv_buffer = ctx.hasOutput(psai::kMidKvBuffer); - const bool has_mid_state_lengths = ctx.hasOutput(psai::kMidStateLengths); - const bool has_validation_status = ctx.hasOutput(psai::kValidationStatus); - if (has_mid_kv_buffer != has_mid_state_lengths) { - fail_shape_inference("PackedSparseAttentionIndexer: mid_kv_buffer and mid_state_lengths must be connected together"); - } - if (!is_qsa && (has_mid_kv_buffer || has_validation_status)) { - fail_shape_inference("PackedSparseAttentionIndexer: mid-state and validation outputs are only supported for qsa"); - } if (!is_qsa && (state_update_capacity > 0 || has_capture_count || has_state_update_active || has_state_update_output)) { fail_shape_inference("PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'"); @@ -2925,9 +2916,9 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } if (ctx.getNumOutputs() < static_cast(psai::kFixedOutputCount) || - ctx.getNumOutputs() > static_cast(psai::kMaxOutputCount)) { - fail_shape_inference("PackedSparseAttentionIndexer: expected ", psai::kFixedOutputCount, " to ", - psai::kMaxOutputCount, " declared outputs, got ", ctx.getNumOutputs()); + ctx.getNumOutputs() > static_cast(psai::kOutputCount)) { + fail_shape_inference("PackedSparseAttentionIndexer: expected ", psai::kFixedOutputCount, " or ", + psai::kOutputCount, " declared outputs, got ", ctx.getNumOutputs()); } if (ctx.hasOutput(psai::kPresentGateBuffer) == is_qsa) { fail_shape_inference("PackedSparseAttentionIndexer: output ", psai::kPresentGateBuffer, @@ -2938,36 +2929,12 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference updateOutputElemType(ctx, psai::kSelectedCounts, ONNX_NAMESPACE::TensorProto_DataType_INT32); if (has_indexshare) updateOutputElemType(ctx, psai::kIndexShareStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); updateOutputElemType(ctx, psai::kPresentStateLengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); - if (has_mid_state_lengths) { - updateOutputElemType(ctx, psai::kMidStateLengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); - propagateShapeFromInputToOutput(ctx, psai::kPastStateLengths, psai::kMidStateLengths); - } - if (has_validation_status) { - updateOutputElemType(ctx, psai::kValidationStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); - propagateShapeFromInputToOutput(ctx, psai::kPastSequenceLengths, psai::kValidationStatus); - } const bool has_state_status_output = ctx.hasOutput(psai::kStateStatus); if (has_state_status_output) { updateOutputElemType(ctx, psai::kStateStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); } propagateElemTypeFromInputToOutput(ctx, psai::kPastKeyState, psai::kPresentKeyState); propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kPresentKvBuffer); - if (has_mid_kv_buffer) { - propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); - propagateShapeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); - } - for (int prefix = 0; prefix < 2; ++prefix) { - const int buffer = psai::kMidKvBuffer2 + 2 * prefix; - const int lengths = buffer + 1; - if (ctx.hasOutput(buffer) != ctx.hasOutput(lengths)) - fail_shape_inference("PackedSparseAttentionIndexer: prefix buffer and lengths must be connected together"); - if (!ctx.hasOutput(buffer)) continue; - if (!is_qsa) fail_shape_inference("PackedSparseAttentionIndexer: prefix outputs require policy_mode='qsa'"); - propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, buffer); - propagateShapeFromInputToOutput(ctx, psai::kPastKvBuffer, buffer); - updateOutputElemType(ctx, lengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); - propagateShapeFromInputToOutput(ctx, psai::kPastStateLengths, lengths); - } if (!is_qsa) { // present_gate_buffer keeps its fixed positional slot (with an empty name) for policy_mode // 'qsa'; only propagate its type/shape when it is actually produced. @@ -3242,8 +3209,6 @@ Unlike SparseAttentionIndexer, this operator: (the optional state_status output reports which requests were rejected and why); * additionally emits selected_counts, the exact number of active (non -1) entries per query, so that no downstream consumer needs to scan selected_indices for its query's true count. - * for QSA, optionally emits the pending-state snapshot after the first request token and a - per-request validation status for overflow or inconsistent state metadata. Both policy_mode values keep the semantics of SparseAttentionIndexer, applied independently to each request's own packed token range and fixed-capacity state slice: @@ -3500,26 +3465,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "synchronizing inside the operator. Currently supported by the CUDA provider only.", "M", OpSchema::Optional) - .Output(9, "mid_kv_buffer", - "Only for CUDA policy_mode 'qsa' refresh: pending-buffer snapshot after the first token of each request. " - "Not supported in IndexShare reuse mode.", - "T", OpSchema::Optional) - .Output(10, "mid_state_lengths", - "Only for CUDA policy_mode 'qsa' refresh: state lengths after the first token of each request. " - "Not supported in IndexShare reuse mode.", - "M", OpSchema::Optional) - .Output(11, "validation_status", - "Only for CUDA policy_mode 'qsa' refresh: per-request overflow status; two-token rejection is encoded as 0x99. " - "Not supported in IndexShare reuse mode.", - "M", OpSchema::Optional) - .Output(12, "mid_kv_buffer2", "CUDA QSA pending-buffer snapshot after at most two request tokens.", - "T", OpSchema::Optional) - .Output(13, "mid_state_lengths2", "CUDA QSA state lengths after at most two request tokens.", - "M", OpSchema::Optional) - .Output(14, "mid_kv_buffer3", "CUDA QSA pending-buffer snapshot after at most three request tokens.", - "T", OpSchema::Optional) - .Output(15, "mid_state_lengths3", "CUDA QSA state lengths after at most three request tokens.", - "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain floating point tensors to float, float16 and bfloat16.") diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 921f846112201..8643770738dc5 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -558,7 +558,7 @@ TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsWrongOutputCount) { GraphOptions options; options.output_count = 4; ExpectResolveFailure([&options](ModelTestBuilder& builder) { AddNode(builder, options); }, - "output size 4 not in range [min=6, max=16]"); + "output size 4 not in range [min=6, max=9]"); } TEST(PackedSparseAttentionIndexerShapeInferenceTest, RejectsStateUpdateCapacityWithoutCaptureCount) { @@ -1096,26 +1096,12 @@ void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedP if (problem.state_update_capacity == 0) test.AddOptionalOutputEdge(); test.AddOutput("status", {total_tokens}, std::vector(static_cast(total_tokens), 0)); } - if (problem.report_state_status || problem.report_mid_state) { + if (problem.report_state_status) { if (padding == 0) { if (problem.state_update_capacity == 0) test.AddOptionalOutputEdge(); test.AddOptionalOutputEdge(); // indexshare_status } - if (problem.report_state_status) { - test.AddOutput("state_status", {batch_size}, expected.state_status); - } else { - test.AddOptionalOutputEdge(); - } - } - if (problem.report_mid_state) { - for (int batch = 0; batch < batch_size; ++batch) { - ASSERT_EQ(problem.cumulative_sequence_lengths[batch + 1] - problem.cumulative_sequence_lengths[batch], 1); - } - test.AddOutput("mid_kv_buffer", {batch_size, buffer_capacity, head_size}, - ToElementType(expected.present_kv_buffer), false, 0.0f, tolerance); - test.AddOutput("mid_state_lengths", {batch_size, 2}, expected.present_state_lengths); - test.AddOutput("validation_status", {batch_size}, - std::vector(static_cast(batch_size), 0)); + test.AddOutput("state_status", {batch_size}, expected.state_status); } RunOnProvider(test, std::move(provider), expected_error.empty() ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, @@ -1141,20 +1127,6 @@ TEST(PackedSparseAttentionIndexerTest, QsaFloat16) { RunQsaPackedTest TEST(PackedSparseAttentionIndexerTest, QsaBFloat16) { RunQsaPackedTest(2.0e-2f); } -TEST(PackedSparseAttentionIndexerTest, QsaMidStateOutputsPreserveStateStatusSlot) { - for (bool report_state_status : {false, true}) { - QsaPackedProblem problem; - problem.cumulative_sequence_lengths = {0, 1, 2}; - problem.past_sequence_lengths = {0, 0}; - problem.past_state_lengths = {0, 0, 0, 0}; - problem.report_state_status = report_state_status; - problem.report_mid_state = true; - problem = MakeQsaPackedProblem(std::move(problem)); - RunQsaPackedTest(1.0e-5f, problem, ProviderKind::Cuda); - RunQsaPackedTest(2.0e-2f, std::move(problem), ProviderKind::Cuda); - } -} - TEST(PackedSparseAttentionIndexerTest, QsaBFloat16FloatState) { RunQsaPackedTest(2.0e-2f); } @@ -1168,7 +1140,7 @@ TEST(PackedSparseAttentionIndexerTest, IndexShareRefreshPreservesBudgetAndPadsRo } TEST(PackedSparseAttentionIndexerTest, IndexShareReusePackedQkAndStatus) { - for (int invalid = 0; invalid < 9; ++invalid) { + for (int invalid = 0; invalid < 7; ++invalid) { SCOPED_TRACE(invalid); std::vector cumulative{0, 1, 2}; std::vector past_lengths{5, 5}; @@ -1217,17 +1189,7 @@ TEST(PackedSparseAttentionIndexerTest, IndexShareReusePackedQkAndStatus) { test.AddOptionalOutputEdge(); test.AddOutput("status", {2}, {invalid == 0 ? 0 : 1, 2}); test.AddOutput("state_status", {2}, {0, 0}); - if (invalid == 7) { - test.AddOutput("mid_buffer", {2, 3, 2}, std::vector(12)); - test.AddOutput("mid_lengths", {2, 2}, std::vector(4)); - } else if (invalid == 8) { - test.AddOptionalOutputEdge(); - test.AddOptionalOutputEdge(); - test.AddOutput("validation", {2}, {0, 0}); - } - RunOnProvider(test, std::move(provider), - invalid >= 7 ? OpTester::ExpectResult::kExpectFailure : OpTester::ExpectResult::kExpectSuccess, - invalid >= 7 ? "outputs 9-11 are not supported in IndexShare reuse mode" : ""); + RunOnProvider(test, std::move(provider)); } } From aa7877617243ef9aa1109ce65942b857bcdb38d3 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 05:34:05 +0000 Subject: [PATCH 15/23] perf(cuda): optimize INT8 KV DynamicSparseAttention prefill and PackedSparseAttentionIndexer - Extend DynamicSparseAttentionPrefillKernel to INT8 KV caches using async cp.async.ca.shared.global (CpAsync8) overlapped with Tensor Core MMA + online softmax, alias sQ with sK (-8.45 KB smem), and set __launch_bounds__(128, 3) (3 CTAs/SM, 5.3x lower long_scoreboard stall, 14.50 ms vs 34.66 ms at L=2048). - Parallelize PackedSparseAttentionIndexer QsaUpdateState across compressed blocks via QsaUpdateStateBlocksKernel and enable WarpHeads=true by default in QsaRaggedPrefillScoreKernel. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../bert/dynamic_sparse_attention_impl.cu | 29 ++-- .../bert/dynamic_sparse_attention_prefill.cuh | 96 ++++++++++-- .../packed_sparse_attention_indexer_impl.cu | 147 +++++++++++++++--- 3 files changed, 224 insertions(+), 48 deletions(-) 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 2fce55999d618..83d3536f0b443 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -2052,7 +2052,7 @@ Status LaunchDynamicSparseAttention( CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); constexpr int kMaxGridDimYZ = 65535; use_prefill_kernel = - enabled && !data.quantized_cache && cc_major >= 8 && parameters.sequence_length >= 8 && + enabled && cc_major >= 8 && parameters.sequence_length >= 8 && parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && parameters.selected_kv_source == DynamicSparseAttentionKvSource::kMain && !parameters.use_smooth_softmax && parameters.max_selected > 0 && @@ -2199,16 +2199,27 @@ Status LaunchDynamicSparseAttention( } } else if (use_prefill_kernel) { if constexpr (std::is_same_v) { - CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( - dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - static_cast(dsa_prefill::kSharedBytes))); const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), static_cast(parameters.kv_num_heads)); - dsa_prefill::DynamicSparseAttentionPrefillKernel<<>>( - data.prepared_query, data.present_key, data.present_value, data.selected_indices, data.selected_counts, - data.seqlens_k, data.output, parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, - parameters.cache_capacity, parameters.max_selected, parameters.scale); + if (data.quantized_cache) { + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kSharedBytes))); + dsa_prefill::DynamicSparseAttentionPrefillKernel<<>>( + data.prepared_query, data.present_key, data.present_value, data.selected_indices, data.selected_counts, + data.seqlens_k, data.output, parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, + parameters.cache_capacity, parameters.max_selected, parameters.scale, data.key_scale, data.value_scale); + } else { + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kSharedBytes))); + dsa_prefill::DynamicSparseAttentionPrefillKernel<<>>( + data.prepared_query, data.present_key, data.present_value, data.selected_indices, data.selected_counts, + data.seqlens_k, data.output, parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, + parameters.cache_capacity, parameters.max_selected, parameters.scale, nullptr, nullptr); + } } } else if (!data.quantized_cache && grouped_prefill_threads_supported && UseGroupedPrefillAttention(parameters, sizeof(T), max_shared_memory_per_block)) { diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh index 03cc1f084793e..a353787bf4b41 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh @@ -32,7 +32,8 @@ constexpr size_t kQBytes = kMaxGroup * kRowStride * sizeof(__nv_bfloat16); constexpr size_t kKBytes = kTile * kRowStride * sizeof(__nv_bfloat16); constexpr size_t kSBytes = kMaxGroup * kSStride * sizeof(float); constexpr size_t kPBytes = kMaxGroup * kPStride * sizeof(__nv_bfloat16); -constexpr size_t kSharedBytes = kQBytes + 2 * kKBytes + kSBytes + kPBytes + 2 * kMaxGroup * sizeof(float); +// sQ is only used to load qf[][] into registers before the tile loop and aliases sK thereafter. +constexpr size_t kSharedBytes = 2 * kKBytes + kSBytes + kPBytes + 2 * kMaxGroup * sizeof(float); #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 __device__ __forceinline__ uint32_t SmemAddr(const void* p) { @@ -42,6 +43,9 @@ __device__ __forceinline__ uint32_t SmemAddr(const void* p) { __device__ __forceinline__ void CpAsync16(void* smem_dst, const void* gmem_src) { asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(SmemAddr(smem_dst)), "l"(gmem_src)); } +__device__ __forceinline__ void CpAsync8(void* smem_dst, const void* gmem_src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" ::"r"(SmemAddr(smem_dst)), "l"(gmem_src)); +} __device__ __forceinline__ void CpAsyncCommit() { asm volatile("cp.async.commit_group;\n" ::); } template __device__ __forceinline__ void CpAsyncWait() { @@ -65,11 +69,31 @@ __device__ __forceinline__ void MmaBf16(float (&c)[4], const uint32_t (&a)[4], u : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); } + +__device__ __forceinline__ uint4 DequantizeInt8x8ToBf16(uint2 packed, float4 s0, float4 s1) { + const int32_t w0 = static_cast(packed.x); + const int32_t w1 = static_cast(packed.y); + const float b0 = static_cast(static_cast(w0 << 24) >> 24) * s0.x; + const float b1 = static_cast(static_cast(w0 << 16) >> 24) * s0.y; + const float b2 = static_cast(static_cast(w0 << 8) >> 24) * s0.z; + const float b3 = static_cast(static_cast(w0) >> 24) * s0.w; + const float b4 = static_cast(static_cast(w1 << 24) >> 24) * s1.x; + const float b5 = static_cast(static_cast(w1 << 16) >> 24) * s1.y; + const float b6 = static_cast(static_cast(w1 << 8) >> 24) * s1.z; + const float b7 = static_cast(static_cast(w1) >> 24) * s1.w; + uint4 out; + *reinterpret_cast<__nv_bfloat162*>(&out.x) = __floats2bfloat162_rn(b0, b1); + *reinterpret_cast<__nv_bfloat162*>(&out.y) = __floats2bfloat162_rn(b2, b3); + *reinterpret_cast<__nv_bfloat162*>(&out.z) = __floats2bfloat162_rn(b4, b5); + *reinterpret_cast<__nv_bfloat162*>(&out.w) = __floats2bfloat162_rn(b6, b7); + return out; +} #endif // grid = (batch * sequence_length, kv_num_heads), block = kThreads, dynamic smem = kSharedBytes. // query: [rows, num_heads, 256] (already normalized / rotated); key, value: [batch, kv_heads, capacity, 256]. -__global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKernel( +template +__global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillKernel( const __nv_bfloat16* __restrict__ query, const __nv_bfloat16* __restrict__ key, const __nv_bfloat16* __restrict__ value, @@ -82,11 +106,13 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern int kv_num_heads, int capacity, int max_selected, - float scale) { + float scale, + const float* __restrict__ key_scale = nullptr, + const float* __restrict__ value_scale = nullptr) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 extern __shared__ __align__(16) unsigned char smem_raw[]; - __nv_bfloat16* sQ = reinterpret_cast<__nv_bfloat16*>(smem_raw); - __nv_bfloat16* sK = sQ + kMaxGroup * kRowStride; + __nv_bfloat16* sK = reinterpret_cast<__nv_bfloat16*>(smem_raw); + __nv_bfloat16* sQ = sK; __nv_bfloat16* sV = sK + kTile * kRowStride; float* sS = reinterpret_cast(sV + kTile * kRowStride); __nv_bfloat16* sP = reinterpret_cast<__nv_bfloat16*>(sS + kMaxGroup * kSStride); @@ -108,8 +134,23 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern const int total = valid_length ? static_cast(total_length) : 0; const int count = valid_length ? max(0, min(selected_counts[row], max_selected)) : 0; const int32_t* row_indices = selected_indices + static_cast(row) * max_selected; - const __nv_bfloat16* key_base = key + (static_cast(batch) * kv_num_heads + kv_head) * capacity * kHeadDim; - const __nv_bfloat16* value_base = value + (static_cast(batch) * kv_num_heads + kv_head) * capacity * kHeadDim; + const int64_t kv_head_offset = (static_cast(batch) * kv_num_heads + kv_head) * capacity * kHeadDim; + const __nv_bfloat16* key_base = key + kv_head_offset; + const __nv_bfloat16* value_base = value + kv_head_offset; + const int8_t* key_base_i8 = reinterpret_cast(key) + kv_head_offset; + const int8_t* value_base_i8 = reinterpret_cast(value) + kv_head_offset; + float4 ks0 = make_float4(1.0f, 1.0f, 1.0f, 1.0f); + float4 ks1 = make_float4(1.0f, 1.0f, 1.0f, 1.0f); + float4 vs0 = make_float4(1.0f, 1.0f, 1.0f, 1.0f); + float4 vs1 = make_float4(1.0f, 1.0f, 1.0f, 1.0f); + if constexpr (kQuantizedCache) { + const float* ks_ptr = key_scale + kv_head * kHeadDim + lane * 8; + const float* vs_ptr = value_scale + kv_head * kHeadDim + lane * 8; + ks0 = *reinterpret_cast(ks_ptr); + ks1 = *reinterpret_cast(ks_ptr + 4); + vs0 = *reinterpret_cast(vs_ptr); + vs1 = *reinterpret_cast(vs_ptr + 4); + } // Stage the group's query heads; the padding rows are zero so their (ignored) scores stay finite. const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; @@ -129,6 +170,7 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern for (int k = 0; k < 16; ++k) { LdmatrixX4(qf[k], sQ + (lane & 15) * kRowStride + k * 16 + (lane >> 4) * 8); } + __syncthreads(); auto candidate_index = [&](int candidate) -> int { if (candidate >= count) return -1; @@ -136,26 +178,46 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern return (index >= 0 && index < total && index <= query_position && index < capacity) ? index : -1; }; - auto issue_rows = [&](__nv_bfloat16* dst, const __nv_bfloat16* base, int tile) { + auto issue_rows = [&](__nv_bfloat16* dst, const __nv_bfloat16* base, const int8_t* base_i8, int tile) { // 4 candidate rows per pass (one per warp), 32 lanes x 16 B = one 512 B row. -#pragma unroll 4 +#pragma unroll for (int it = 0; it < kTile / 4; ++it) { const int cand_in_tile = it * 4 + warp; const int index = candidate_index(tile * kTile + cand_in_tile); __nv_bfloat16* row_dst = dst + cand_in_tile * kRowStride + lane * 8; if (index >= 0) { - CpAsync16(row_dst, base + static_cast(index) * kHeadDim + lane * 8); + if constexpr (kQuantizedCache) { + CpAsync8(row_dst, base_i8 + static_cast(index) * kHeadDim + lane * 8); + } else { + CpAsync16(row_dst, base + static_cast(index) * kHeadDim + lane * 8); + } } else { - *reinterpret_cast(row_dst) = make_uint4(0, 0, 0, 0); + if constexpr (kQuantizedCache) { + *reinterpret_cast(row_dst) = make_uint2(0, 0); + } else { + *reinterpret_cast(row_dst) = make_uint4(0, 0, 0, 0); + } + } + } + }; + + auto dequantize_rows = [&](__nv_bfloat16* dst, float4 sc0, float4 sc1) { + if constexpr (kQuantizedCache) { +#pragma unroll + for (int it = 0; it < kTile / 4; ++it) { + const int cand_in_tile = it * 4 + warp; + __nv_bfloat16* row_dst = dst + cand_in_tile * kRowStride + lane * 8; + const uint2 packed = *reinterpret_cast(row_dst); + *reinterpret_cast(row_dst) = DequantizeInt8x8ToBf16(packed, sc0, sc1); } } }; const int tiles = (count + kTile - 1) / kTile; if (tiles > 0) { - issue_rows(sK, key_base, 0); + issue_rows(sK, key_base, key_base_i8, 0); CpAsyncCommit(); - issue_rows(sV, value_base, 0); + issue_rows(sV, value_base, value_base_i8, 0); CpAsyncCommit(); } @@ -172,6 +234,7 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern for (int tile = 0; tile < tiles; ++tile) { CpAsyncWait<1>(); // K(tile) has landed, V(tile) may still be in flight + dequantize_rows(sK, ks0, ks1); __syncthreads(); // S = Q K^T for this warp's 16 candidate columns. @@ -197,7 +260,7 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern __syncthreads(); // scores visible; sK is free if (tile + 1 < tiles) { - issue_rows(sK, key_base, tile + 1); + issue_rows(sK, key_base, key_base_i8, tile + 1); } CpAsyncCommit(); // K(tile+1) (possibly empty) @@ -247,6 +310,7 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern } CpAsyncWait<1>(); // V(tile) has landed (K(tile+1) may be in flight) + dequantize_rows(sV, vs0, vs1); __syncthreads(); // O += P V for this warp's 64 output dimensions. @@ -266,7 +330,7 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern __syncthreads(); // sV and sP are free if (tile + 1 < tiles) { - issue_rows(sV, value_base, tile + 1); + issue_rows(sV, value_base, value_base_i8, tile + 1); } CpAsyncCommit(); // V(tile+1) (possibly empty) } @@ -304,6 +368,8 @@ __global__ void __launch_bounds__(kThreads, 2) DynamicSparseAttentionPrefillKern (void)capacity; (void)max_selected; (void)scale; + (void)key_scale; + (void)value_scale; #endif } diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index c8e0800fba77c..9141b5e803f0c 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -138,11 +138,95 @@ __global__ void ElementwiseCopyKernel(const T* src, T* dst, int64_t count) { // policy_mode = "qsa" // --------------------------------------------------------------------------------------------- +template +__global__ void QsaUpdateStateBlocksKernel(const T* key, const T* key_norm_weight, const T* cos_cache, + const T* sin_cache, const int32_t* cumulative_sequence_lengths, + const int32_t* past_sequence_lengths, const int32_t* past_state_lengths, + const int32_t* state_update_capture_count, bool state_update_active, + TState* present_key_state, T* state_update, + PackedSparseAttentionIndexerParams params) { + extern __shared__ float shared[]; + float* pooled = shared; + float* reduction = shared + params.head_size; + + for (int b = static_cast(blockIdx.x); b < params.batch_size; b += static_cast(gridDim.x)) { + const int req_start = cumulative_sequence_lengths[b]; + const int req_end = cumulative_sequence_lengths[b + 1]; + const int raw_key_len = past_state_lengths[b * 2 + psai::kKeyStateLength]; + const int raw_buf_len = past_state_lengths[b * 2 + psai::kBufferLength]; + const int past_sequence_length = past_sequence_lengths[b]; + const bool invalid_metadata = + cumulative_sequence_lengths[0] != 0 || + cumulative_sequence_lengths[params.batch_size] != params.total_tokens || + req_start < 0 || req_end < req_start || req_end > params.total_tokens || + past_sequence_length < 0 || raw_key_len < 0 || raw_key_len > params.state_capacity || + raw_buf_len < 0 || raw_buf_len >= params.compress_ratio || + raw_key_len != past_sequence_length / params.compress_ratio || + raw_buf_len != past_sequence_length % params.compress_ratio; + const int req_len = invalid_metadata ? 0 : req_end - req_start; + const int old_key_len = min(max(raw_key_len, 0), params.state_capacity); + const int old_buf_len = min(max(raw_buf_len, 0), params.compress_ratio - 1); + const int pending = old_buf_len + req_len; + const int full_new_block_count = pending / params.compress_ratio; + const int capacity_left = params.state_capacity - old_key_len; + const bool rejected = invalid_metadata || full_new_block_count > capacity_left; + const int new_block_count = rejected ? 0 : full_new_block_count; + const bool capture_active = !rejected && state_update != nullptr && state_update_capture_count != nullptr && + state_update_active; + const int capture_count = capture_active + ? min(max(state_update_capture_count[b], 0), + min(req_len, params.state_update_capacity)) + : 0; + + for (int k = 1 + static_cast(blockIdx.y); k < new_block_count; k += static_cast(gridDim.y)) { + float sum_squares = 0.0f; + for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { + float sum = 0.0f; + for (int t = 0; t < params.compress_ratio; ++t) { + const int virtual_pos = k * params.compress_ratio + t; + sum += to_float(key[(static_cast(req_start) + (virtual_pos - old_buf_len)) * + params.key_row_stride + + d]); + } + const float p = sum / static_cast(params.compress_ratio); + pooled[d] = p; + sum_squares += p * p; + } + sum_squares = SaiBlockSum(sum_squares, reduction); + const float inverse_rms = rsqrtf(sum_squares / static_cast(params.head_size) + params.epsilon); + for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { + pooled[d] = pooled[d] * inverse_rms * to_float(key_norm_weight[d]); + } + __syncthreads(); + + const int entry = old_key_len + k; + const int rope_position = + SaiClampPosition(static_cast(entry) * params.compress_ratio, params.max_rotary_length); + const int64_t cache_offset = + (static_cast(params.cos_cache_batched ? b : 0) * params.max_rotary_length + rope_position) * + params.rotary_width; + const int64_t out_base = (static_cast(b) * params.state_capacity + entry) * params.head_size; + const int completion_token = (k + 1) * params.compress_ratio - old_buf_len - 1; + const int64_t capture_base = + (static_cast(b) * params.state_update_capacity + completion_token) * params.head_size; + for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { + const float rot = SaiLeadingRope(pooled, params.rotary_width, cos_cache + cache_offset, + sin_cache + cache_offset, d); + present_key_state[out_base + d] = from_float(rot); + if (completion_token < capture_count) { + state_update[capture_base + d] = from_float(rot); + } + } + __syncthreads(); + } + } +} + // One block per request: forms every newly-closed compress_ratio block (mean-pool -> RMSNorm -> // leading RoPE -> append into fixed-capacity key_state) and publishes the raw trailing buffer. // Reads only past_kv_buffer / key (never present_kv_buffer / present_key_state), so it is correct // whether or not the present/past tensors are the same aliased allocation. -template +template __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, const T* cos_cache, const T* sin_cache, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const T* past_kv_buffer, @@ -214,7 +298,9 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con } __syncthreads(); - for (int k = 0; k < new_block_count; ++k) { + const int loop_limit = kFirstBlockOnly ? min(new_block_count, 1) : new_block_count; + for (int k = 0; k < loop_limit; ++k) { + float sum_squares = 0.0f; for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { float sum = 0.0f; for (int t = 0; t < params.compress_ratio; ++t) { @@ -227,13 +313,9 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con params.key_row_stride + d]); } - pooled[d] = sum / static_cast(params.compress_ratio); - } - __syncthreads(); - - float sum_squares = 0.0f; - for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { - sum_squares += pooled[d] * pooled[d]; + const float p = sum / static_cast(params.compress_ratio); + pooled[d] = p; + sum_squares += p * p; } sum_squares = SaiBlockSum(sum_squares, reduction); const float inverse_rms = rsqrtf(sum_squares / static_cast(params.head_size) + params.epsilon); @@ -248,24 +330,21 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con const int64_t cache_offset = (static_cast(params.cos_cache_batched ? b : 0) * params.max_rotary_length + rope_position) * params.rotary_width; - for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { - rotated[d] = SaiLeadingRope(pooled, params.rotary_width, cos_cache + cache_offset, - sin_cache + cache_offset, d); - } - __syncthreads(); - const int64_t out_base = (static_cast(b) * params.state_capacity + entry) * params.head_size; const int completion_token = (k + 1) * params.compress_ratio - old_buf_len - 1; const int64_t capture_base = (static_cast(b) * params.state_update_capacity + completion_token) * params.head_size; for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { - present_key_state[out_base + d] = from_float(rotated[d]); + const float rot = SaiLeadingRope(pooled, params.rotary_width, cos_cache + cache_offset, + sin_cache + cache_offset, d); + present_key_state[out_base + d] = from_float(rot); if (completion_token < capture_count) { - state_update[capture_base + d] = from_float(rotated[d]); + state_update[capture_base + d] = from_float(rot); } } __syncthreads(); } + (void)rotated; // Publish the raw trailing buffer. Skipped entirely on overflow: present_kv_buffer already // holds past_kv_buffer's contents unchanged (from the baseline copy in the Launch function @@ -1647,13 +1726,33 @@ Status LaunchQsaPackedSparseAttentionIndexer( } const size_t value_bytes = static_cast(params.head_size) * sizeof(float); - ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( - QsaUpdateStateKernel, 2 * value_bytes + kThreads * sizeof(float), "QsaUpdateStateKernel")); const int state_blocks = static_cast(std::min(params.batch_size, kSaiMaxGridDimX)); - QsaUpdateStateKernel<<>>( - key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, past_kv_buffer, - past_state_lengths, state_update_capture_count, state_update_active, present_key_state, present_kv_buffer, - present_state_lengths, state_update, overflow_flags, params); + if (params.total_tokens >= 4 * params.compress_ratio && params.compress_ratio > 0) { + const int parallel_blocks = + std::min((params.total_tokens + params.compress_ratio - 1) / params.compress_ratio, 512); + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + QsaUpdateStateBlocksKernel, value_bytes + kThreads * sizeof(float), "QsaUpdateStateBlocksKernel")); + QsaUpdateStateBlocksKernel + <<>>( + key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, + past_state_lengths, state_update_capture_count, state_update_active, present_key_state, state_update, + params); + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + QsaUpdateStateKernel, 2 * value_bytes + kThreads * sizeof(float), "QsaUpdateStateKernel")); + QsaUpdateStateKernel + <<>>( + key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, + past_kv_buffer, past_state_lengths, state_update_capture_count, state_update_active, present_key_state, + present_kv_buffer, present_state_lengths, state_update, overflow_flags, params); + } else { + ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( + QsaUpdateStateKernel, 2 * value_bytes + kThreads * sizeof(float), "QsaUpdateStateKernel")); + QsaUpdateStateKernel + <<>>( + key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, + past_kv_buffer, past_state_lengths, state_update_capture_count, state_update_active, present_key_state, + present_kv_buffer, present_state_lengths, state_update, overflow_flags, params); + } if (params.total_tokens == 0) { return CUDA_CALL(cudaGetLastError()); @@ -1748,7 +1847,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( overflow_flags, key_stride, selected_indices, selected_counts, params); } else if (use_ragged_prefill) { static const bool warp_heads = - ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_WARP_HEADS", false); + ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_WARP_HEADS", true); if (warp_heads) { QsaRaggedPrefillScoreKernel<<>>( present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, From 18a13c7011414dd599f2bc17ff35a0343474aa12 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 06:20:16 +0000 Subject: [PATCH 16/23] fix(cuda): fix FP16 ratio-12 grouped decode MMA and expand prefill test coverage - Stage all kQueryHeadsPerBlock heads into mma_query and populate the second 8-row MMA fragment/score slots in GroupedDecodeDynamicSparseAttentionKernel so warps 8-11 read computed scores for GQA ratio 12 on SM80+. - Extend FloatingCacheGroupedReferenceParity_CUDA to sequence_length=9 to exercise GroupedPrefillDynamicSparseAttentionKernel, add GroupedDecodeRatio12ReferenceParity_CUDA, and add Int8CachePrefillBFloat16_CUDA. - Remove unused legacy QsaPairedScoreTileTopKKernel and QsaCompactTileTopKKernel from packed_sparse_attention_indexer_impl.cu and sync packed_sparse_attention_indexer.md with current defaults and CSA-only position_ids. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../packed_sparse_attention_indexer.md | 28 ++-- .../bert/dynamic_sparse_attention_impl.cu | 17 ++- .../packed_sparse_attention_indexer_impl.cu | 143 ------------------ .../dynamic_sparse_attention_op_test.cc | 62 +++++++- 4 files changed, 88 insertions(+), 162 deletions(-) diff --git a/docs/contrib_ops/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md index 300453da899eb..615d04c086217 100644 --- a/docs/contrib_ops/packed_sparse_attention_indexer.md +++ b/docs/contrib_ops/packed_sparse_attention_indexer.md @@ -3,8 +3,8 @@ ## Experimental CUDA decode specializations Set these process-level options before the first operator invocation. All three -default to `0`; only `0` and `1` are accepted. Existing dispatch remains available -when an option is disabled or its geometry is unsupported. +default to `1`; only `0` and `1` are accepted. Existing dispatch remains available +when an option is disabled (`0`) or its geometry is unsupported. - `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1`: compact QSA chunks of 2048 keys to 512 keys with device-side 512/1024/2048 sorting probes. Up to five chunks @@ -15,9 +15,9 @@ when an option is disabled or its geometry is unsupported. - `ORT_DSA_INT8_VECTOR_STAGING=1`: the companion `DynamicSparseAttention` grouped INT8 path uses eight-byte cache loads and independent per-channel key/value scales. Its existing guards still require head size 256, FP16/BF16 - queries, one or two tokens, GQA ratio at least four, more than 256 candidates, - and sufficient resources. Auxiliary floating-point caches retain their own - unscaled load path. + queries, one through eight tokens (`1 <= sequence_length <= 8`), GQA ratio at + least four, more than 256 candidates, and sufficient resources. Auxiliary + floating-point caches retain their own unscaled load path. - `ORT_DSA_SHARED_SCALE_MERGE=1`: grouped attention merges with at most 32 splits compute the exponential factor once per split and share it across output channels. A nonzero partial sum, not the maximum-score sentinel, @@ -99,7 +99,7 @@ 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 | optional qsa / required csa | +| 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 | @@ -240,9 +240,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"`. Optional QSA `position_ids` override query rotary positions and causal -visibility. Callers must keep these positions consistent with prepared block keys and the cache -indices consumed by attention. +`selected_kv_source="main"`. ## 6. Policy `csa` @@ -260,7 +258,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` optionally accepts an explicit position override). + `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`. @@ -325,13 +323,13 @@ 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. -- CUDA has two experimental, default-off dispatch controls. Set them before the first eligible - call; each value is cached per process. `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1` uses live-key - radix selection for eligible short QSA calls. `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1` uses +- CUDA has two additional dispatch controls. Set them before the first eligible + call; each value is cached per process. `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1` (default `0`, + mutually exclusive with `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1`) uses live-key + radix selection for eligible short QSA calls. `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1` (default `1`) uses warp-local scoring for eligible packed prefill calls with more than 64 tokens, four heads, head size 128 and compression ratio four. Both retain causal visibility, stable tie ordering, - selected counts and state-update semantics. Broad full-model quality and performance - qualification remains outstanding; neither option is a general tuning recommendation. + selected counts and state-update semantics. - 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_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index 83d3536f0b443..963900bfe6a3c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -1006,7 +1006,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( if constexpr (std::is_same_v && !kQuantizedCache) { const int group_head_base = head_tile * kQueryHeadsPerBlock; for (int index = static_cast(threadIdx.x); - index < kGroupedDecodeQueryHeadsPerBlock * head_size; + 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; @@ -1168,6 +1168,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( 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) { @@ -1176,10 +1177,16 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( const uint32_t query_fragment[4]{ *reinterpret_cast( mma_query + row_pair * head_size + first_k), - 0, + second_row < kQueryHeadsPerBlock + ? *reinterpret_cast( + mma_query + second_row * head_size + first_k) + : 0u, *reinterpret_cast( mma_query + row_pair * head_size + second_k), - 0}; + 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), @@ -1190,6 +1197,10 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( 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(); } diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 9141b5e803f0c..753a74b4b79d3 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -678,108 +678,6 @@ __global__ void QsaScoreTileTopKKernel(const TState* present_key_state, const fl } } -template -__global__ void QsaPairedScoreTileTopKKernel(const T* present_key_state, const float* query_rotated, - const int32_t* cumulative_sequence_lengths, - const int32_t* past_sequence_lengths, const int64_t* position_ids, - const int32_t* present_state_lengths, const int32_t* overflow_flags, - uint64_t* tile_keys, int tile_count, - PackedSparseAttentionIndexerParams params) { - __shared__ uint64_t keys[2][kHierarchicalTileBlocks]; - __shared__ int block_counts[2]; - const int first_token = static_cast(blockIdx.y) * 2; - const int warp = static_cast(threadIdx.x) / kWarpSize; - const int lane = static_cast(threadIdx.x) % kWarpSize; - if (threadIdx.x < 2) { - const int metadata_token = first_token + static_cast(threadIdx.x); - int count = 0; - if (metadata_token < params.total_tokens && overflow_flags[0] == 0) { - const int64_t position = params.has_position_ids ? position_ids[metadata_token] - : static_cast(past_sequence_lengths[0]) + metadata_token; - const int64_t causal = SaiCausalThreshold(position, kQwenCompressRatio); - const int64_t visible = min(causal, static_cast(present_state_lengths[psai::kKeyStateLength])); - count = static_cast(max(static_cast(0), visible)); - } - block_counts[threadIdx.x] = count; - } - __syncthreads(); - float4 query_values[2][kQwenNumHeads]; -#pragma unroll - for (int query_index = 0; query_index < 2; ++query_index) { -#pragma unroll - for (int head = 0; head < kQwenNumHeads; ++head) { - const int token = first_token + query_index; - query_values[query_index][head] = token < params.total_tokens - ? LoadQsaVector4(query_rotated + (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize + lane * 4) - : make_float4(0.0f, 0.0f, 0.0f, 0.0f); - } - } - const int key_stride = tile_count * kHierarchicalTileBlocks; - for (int query_index = 0; query_index < 2; ++query_index) { - if (first_token + query_index < params.total_tokens) { - const int active = (block_counts[query_index] + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; - const int64_t row_base = static_cast(first_token + query_index) * key_stride; - for (int index = active * kHierarchicalTileBlocks + blockIdx.x * blockDim.x + threadIdx.x; - index < key_stride; index += gridDim.x * blockDim.x) { - tile_keys[row_base + index] = topk::kPaddingSortKey; - } - } - } - const int visible = max(block_counts[0], block_counts[1]); - const int active_tiles = (visible + kHierarchicalTileBlocks - 1) / kHierarchicalTileBlocks; - for (int tile = static_cast(blockIdx.x); tile < active_tiles; tile += static_cast(gridDim.x)) { - const int block = tile * kHierarchicalTileBlocks + warp; - const float4 key_value = block < visible - ? LoadQsaVector4(present_key_state + static_cast(block) * kQwenHeadSize + lane * 4) - : make_float4(0.0f, 0.0f, 0.0f, 0.0f); -#pragma unroll - for (int query_index = 0; query_index < 2; ++query_index) { - float total = 0.0f; -#pragma unroll - for (int head = 0; head < kQwenNumHeads; ++head) { - const float4 query_value = query_values[query_index][head]; - float score = query_value.x * key_value.x + query_value.y * key_value.y + - query_value.z * key_value.z + query_value.w * key_value.w; - score = topk::WarpReduceSum(score); - total += fmaxf(score, 0.0f); - } - uint64_t key = topk::kPaddingSortKey; - if (block < block_counts[query_index]) { - const float scaled = total * params.scale; - key = topk::PackStableSortKey(scaled == 0.0f ? 0.0f : scaled, block); - } - if (lane == 0) { - keys[query_index][warp] = key; - } - } - __syncthreads(); - if (threadIdx.x < 2) { - const int query_index = static_cast(threadIdx.x); - for (int index = 1; index < kHierarchicalTileBlocks; ++index) { - const uint64_t key = keys[query_index][index]; - int insertion = index; - while (insertion > 0 && key > keys[query_index][insertion - 1]) { - keys[query_index][insertion] = keys[query_index][insertion - 1]; - --insertion; - } - keys[query_index][insertion] = key; - } - } - __syncthreads(); - if (threadIdx.x < 2 * kHierarchicalTileBlocks) { - const int query_index = static_cast(threadIdx.x) / kHierarchicalTileBlocks; - const int slot = static_cast(threadIdx.x) % kHierarchicalTileBlocks; - if (first_token + query_index < params.total_tokens && - tile * kHierarchicalTileBlocks < block_counts[query_index]) { - const int64_t output_base = static_cast(first_token + query_index) * key_stride + - tile * kHierarchicalTileBlocks; - tile_keys[output_base + slot] = keys[query_index][slot]; - } - } - __syncthreads(); - } -} - __device__ __forceinline__ uint64_t QsaMergeRank(const uint64_t* left, int left_count, const uint64_t* right, int right_count, int rank) { int low = max(0, rank - right_count); @@ -800,47 +698,6 @@ __device__ __forceinline__ uint64_t QsaMergeRank(const uint64_t* left, int left_ return topk::kPaddingSortKey; } -__global__ void QsaCompactTileTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, - int input_stride, int output_stride) { - constexpr int kChunkSize = 2048; - static_assert(kChunkSize % kHierarchicalTileBlocks == 0); - __shared__ uint64_t shared_keys[2][kChunkSize]; - const int chunk_start = static_cast(blockIdx.x) * kChunkSize; - const int64_t input_base = static_cast(blockIdx.y) * input_stride + chunk_start; - const int64_t output_base = static_cast(blockIdx.y) * output_stride + - static_cast(blockIdx.x) * kBoundedTopKMax; - if (input_keys[input_base] == topk::kPaddingSortKey) { - for (int rank = threadIdx.x; rank < kBoundedTopKMax; rank += blockDim.x) { - output_keys[output_base + rank] = topk::kPaddingSortKey; - } - return; - } - for (int index = threadIdx.x; index < kChunkSize; index += blockDim.x) { - shared_keys[0][index] = chunk_start + index < input_stride - ? input_keys[input_base + index] - : topk::kPaddingSortKey; - } - __syncthreads(); - int input_bank = 0; - for (int width = kHierarchicalTileBlocks; width < kChunkSize; width *= 2) { - const int count = min(width, kBoundedTopKMax); - const int output_count = min(count * 2, kBoundedTopKMax); - const int pairs = kChunkSize / (width * 2); - for (int index = threadIdx.x; index < pairs * output_count; index += blockDim.x) { - const int pair_start = (index / output_count) * width * 2; - const int rank = index % output_count; - shared_keys[1 - input_bank][pair_start + rank] = - QsaMergeRank(shared_keys[input_bank] + pair_start, count, - shared_keys[input_bank] + pair_start + width, count, rank); - } - __syncthreads(); - input_bank = 1 - input_bank; - } - for (int rank = threadIdx.x; rank < kBoundedTopKMax; rank += blockDim.x) { - output_keys[output_base + rank] = shared_keys[input_bank][rank]; - } -} - __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, int key_stride, int list_width, int rows) { const int pairs_per_row = (key_stride + 2 * list_width - 1) / (2 * list_width); diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 94f6081d8ad9e..2fcf479b92944 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -550,7 +550,7 @@ DynamicSparseAttentionCase MakeGroupedCacheCase(int sequence_length, bool auxili } TEST(DynamicSparseAttentionTest, FloatingCacheGroupedReferenceParity_CUDA) { - for (int sequence_length : {1, 3}) { + for (int sequence_length : {1, 3, 9}) { for (bool auxiliary : {false, true}) { for (bool bf16 : {false, true}) { auto cuda_ep = DefaultCudaExecutionProvider(); @@ -569,6 +569,66 @@ TEST(DynamicSparseAttentionTest, FloatingCacheGroupedReferenceParity_CUDA) { } } +TEST(DynamicSparseAttentionTest, GroupedDecodeRatio12ReferenceParity_CUDA) { + for (int sequence_length : {1, 3}) { + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + + ", auxiliary=" + std::to_string(auxiliary)); + const auto floating_case = MakeGroupedCacheCase(sequence_length, auxiliary, true, 2, 24, 2); + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP() << "CUDA EP not available."; + RunDynamicSparseAttentionCase(floating_case, std::move(cuda_ep)); + cuda_ep = DefaultCudaExecutionProvider(); + RunDynamicSparseAttentionCase(floating_case, std::move(cuda_ep)); + + const auto int8_case = MakeGroupedCacheCase(sequence_length, auxiliary, false, 2, 24, 2, true); + cuda_ep = DefaultCudaExecutionProvider(); + RunDynamicSparseAttentionCase(int8_case, std::move(cuda_ep)); + cuda_ep = DefaultCudaExecutionProvider(); + RunDynamicSparseAttentionCase(int8_case, std::move(cuda_ep)); + } + } +} + +TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + auto c = MakeGroupedCacheCase(9, false, false, 2, 20, 2, true); + c.head_sink.clear(); + for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { + const int b = row / c.sequence_length; + for (int h = 0; h < c.num_heads; ++h) { + const int64_t kv = h / (c.num_heads / c.kv_num_heads); + const int64_t out = (row * c.num_heads + h) * c.head_size; + std::vector numerator(c.head_size, 0.0f); + double denominator = 0.0; + for (int i = 0; i < c.selected_counts[row]; ++i) { + const int index = c.selected_indices[row * c.max_selected + i]; + const int64_t offset = ((b * c.kv_num_heads + kv) * c.cache_sequence_length + index) * c.head_size; + double dot = 0.0; + for (int d = 0; d < c.head_size; ++d) { + dot += c.query[out + d] * c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d]; + } + const double weight = std::exp(dot * c.scale); + denominator += weight; + for (int d = 0; d < c.head_size; ++d) { + numerator[d] += static_cast(weight * c.expected_present_value[offset + d] * + c.value_scale[kv * c.head_size + d]); + } + } + for (int d = 0; d < c.head_size; ++d) { + c.expected_output[out + d] = denominator == 0.0 ? 0.0f : static_cast(numerator[d] / denominator); + } + } + } + for (const char* prefill_env : {"1", "0"}) { + SCOPED_TRACE(prefill_env); + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", prefill_env}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider()); + } +} + TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) { for (bool bf16 : {false, true}) { auto cuda_ep = DefaultCudaExecutionProvider(); From fea46be4d8f8b4d5d4b61d139095c18e9ca26d2d Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 16:57:27 +0000 Subject: [PATCH 17/23] fix(cuda): avoid narrowing in sparse attention test Keep the BF16 prefill reference batch index in int64_t to match the sequence length and fix MSVC C4244 warnings treated as errors in Windows CUDA, CUDA plugin, and TensorRT CI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test/contrib_ops/dynamic_sparse_attention_op_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 2fcf479b92944..46ee626bb9c0a 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -596,7 +596,7 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { auto c = MakeGroupedCacheCase(9, false, false, 2, 20, 2, true); c.head_sink.clear(); for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { - const int b = row / c.sequence_length; + const int64_t b = row / c.sequence_length; for (int h = 0; h < c.num_heads; ++h) { const int64_t kv = h / (c.num_heads / c.kv_num_heads); const int64_t out = (row * c.num_heads + h) * c.head_size; From 9116cb9dba45ef1cccf8b452895555b091dd57b1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 13:40:07 -0700 Subject: [PATCH 18/23] [CUDA] Support per-batch committed_length in DynamicSparseAttention (#33627) ### Summary Extends com.microsoft::DynamicSparseAttention (committed_length input 17, introduced in #33472) to accept either [1] or [batch_size] int32 elements and validates committed_length[b * committed_length_stride] per batch item in ValidateInputsKernel. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/dynamic_sparse_attention.cc | 13 ++++++++++--- .../cuda/bert/dynamic_sparse_attention_impl.cu | 14 ++++++++------ .../cuda/bert/dynamic_sparse_attention_impl.h | 3 ++- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 4 ++-- 4 files changed, 22 insertions(+), 12 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index 9555b9f919cd6..d3e020eb16e2e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -165,8 +165,9 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons 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 && - (parameters.batch_size != 1 || committed_length->Shape().Size() != 1), - "DynamicSparseAttention: committed_length requires batch_size=1 and one int32 element."); + 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}); @@ -253,11 +254,17 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons 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, validation_workspace.get(), GetDeviceProp().sharedMemPerBlock, strict_validation_, data.key_scale, data.value_scale, - committed_length == nullptr ? nullptr : committed_length->Data())); + committed_length == nullptr ? nullptr : committed_length->Data(), + committed_length_stride)); } const size_t attention_workspace_elements = 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 963900bfe6a3c..435660592f9f4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -105,7 +105,8 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, const float* key_scale, const float* value_scale, int cache_channels, - const int32_t* committed_length) { + 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; @@ -232,10 +233,10 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, return; } - if (row == 0 && threadIdx.x == 0 && committed_length != nullptr) { - const int64_t committed = *committed_length; + 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[0] != total - 1) { + if (committed < 0 || total > main_capacity || seqlens_k[b] != total - 1) { atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, 9); } } @@ -1876,7 +1877,8 @@ Status ValidateDynamicSparseAttentionOnDevice( bool copy_result_to_host, const float* key_scale, const float* value_scale, - const int32_t* committed_length) { + 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(parameters.batch_size) * parameters.sequence_length; @@ -1906,7 +1908,7 @@ Status ValidateDynamicSparseAttentionOnDevice( parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, validation_bitmap, validation_hash_entries, use_shared_hash_table, error_flag, key_scale, value_scale, parameters.kv_num_heads * parameters.head_size, - committed_length); + committed_length, committed_length_stride); CUDA_RETURN_IF_ERROR(cudaGetLastError()); if (!copy_result_to_host) { diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index 6050386bd9c71..ed71674276a95 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -76,7 +76,8 @@ Status ValidateDynamicSparseAttentionOnDevice( bool copy_result_to_host, const float* key_scale = nullptr, const float* value_scale = nullptr, - const int32_t* committed_length = nullptr); + const int32_t* committed_length = nullptr, + int committed_length_stride = 0); template Status LaunchDynamicSparseAttention( diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 768f88cde070f..d31d2fdfc01d9 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -659,8 +659,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.", "T", OpSchema::Optional) .Input(17, "committed_length", - "Optional single-request committed token count for cursor validation. When supplied, validation_status " - "is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent. CUDA only.", + "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. CUDA only.", "M", OpSchema::Optional) .Input(18, "key_scale", "Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].", "S", OpSchema::Optional) From a4151cc48260715191069cf0c606e218e42d0000 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:51:08 +0000 Subject: [PATCH 19/23] [CUDA] DSA: fail strict validation on inconsistent committed cursor Name validation code 9 (kDynamicSparseAttentionInvalidCommittedLength) and return a cursor-specific error from the strict-validation host check instead of treating it as success when no validation_status output is connected. Add a regression test (strict env var, no status output) and clear smooth_softmax in the INT8 prefill test so the reference matches and the INT8 prefill kernel is reached. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/dynamic_sparse_attention_impl.cu | 12 ++++++++---- .../cuda/bert/dynamic_sparse_attention_impl.h | 1 + .../dynamic_sparse_attention_op_test.cc | 14 ++++++++++++++ 3 files changed, 23 insertions(+), 4 deletions(-) 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 435660592f9f4..17cd24817eae8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -80,8 +80,8 @@ __device__ __forceinline__ void SetValidationError(int32_t* error_flag, DynamicSparseAttentionValidationError error) { const int32_t previous = atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, static_cast(error)); - if (previous == 9) { - atomicCAS(error_flag, 9, static_cast(error)); + if (previous == kDynamicSparseAttentionInvalidCommittedLength) { + atomicCAS(error_flag, kDynamicSparseAttentionInvalidCommittedLength, static_cast(error)); } } @@ -237,7 +237,8 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, 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, 9); + atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, + kDynamicSparseAttentionInvalidCommittedLength); } } } @@ -1919,7 +1920,7 @@ Status ValidateDynamicSparseAttentionOnDevice( CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( &host_error, error_flag, sizeof(host_error), cudaMemcpyDeviceToHost, stream)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); - if (host_error == kDynamicSparseAttentionValidationOk || host_error == 9) { + if (host_error == kDynamicSparseAttentionValidationOk) { return Status::OK(); } @@ -1949,6 +1950,9 @@ Status ValidateDynamicSparseAttentionOnDevice( case kDynamicSparseAttentionInvalidPosition: message = "a rotary position is outside the rotary cache"; break; + case kDynamicSparseAttentionInvalidCommittedLength: + message = "committed_length is inconsistent with seqlens_k, sequence_length, or cache_capacity"; + break; case kDynamicSparseAttentionInvalidCacheScale: message = "INT8 cache scales must be positive and finite"; break; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index ed71674276a95..baa7e566eec17 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -22,6 +22,7 @@ enum DynamicSparseAttentionValidationError : int32_t { kDynamicSparseAttentionNonCausalIndex = 6, kDynamicSparseAttentionInvalidSequenceLength = 7, kDynamicSparseAttentionInvalidPosition = 8, + kDynamicSparseAttentionInvalidCommittedLength = 9, kDynamicSparseAttentionInvalidCacheScale = 10, }; diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 46ee626bb9c0a..63e24186e2e08 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -595,6 +595,7 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; auto c = MakeGroupedCacheCase(9, false, false, 2, 20, 2, true); c.head_sink.clear(); + c.smooth_softmax = 0; for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { const int64_t b = row / c.sequence_length; for (int h = 0; h < c.num_heads; ++h) { @@ -2120,6 +2121,19 @@ TEST(DynamicSparseAttentionTest, RejectsInvalidSelectionMetadata_CUDA) { RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectFailure, "incompatible with the current sequence"); } + { + SCOPED_TRACE("committed cursor disagrees with seqlens_k and no validation_status output"); + auto c = MakeSingleTokenSelectedOnlyCase(); + c.cache_sequence_length = 2; + c.total_sequence_length = 2; + c.past_key.assign(16, 0.0f); + c.past_value.assign(16, 0.0f); + c.expected_present_key.assign(16, 0.0f); + c.expected_present_value.assign(16, 0.0f); + c.committed_length = 1; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectFailure, + "committed_length is inconsistent"); + } } #ifndef ORT_NO_EXCEPTIONS From 977e9813a42b424a639b3a3bf64c8249aa995ed1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:52:25 +0000 Subject: [PATCH 20/23] test(cuda): run alternate DSA/indexer switches in fresh processes The INT8 vector-staging, shared-scale merge, compact/radix top-k and warp-heads switches are cached in function-local statics, so add a small helper that re-runs the relevant existing cases in a new process of the test binary with each non-default combination set before the first eligible call. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../dynamic_sparse_attention_op_test.cc | 29 ++++++++++++ .../contrib_ops/fresh_process_test_util.h | 44 +++++++++++++++++++ ...packed_sparse_attention_indexer_op_test.cc | 44 +++++++++++++++++++ 3 files changed, 117 insertions(+) create mode 100644 onnxruntime/test/contrib_ops/fresh_process_test_util.h diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 63e24186e2e08..14bdc9ed6741c 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -18,6 +18,7 @@ #if defined(USE_CUDA) || defined(USE_WEBGPU) #include "test/common/tensor_op_test_utils.h" +#include "test/contrib_ops/fresh_process_test_util.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/scoped_env_vars.h" #include "test/util/include/default_providers.h" @@ -691,6 +692,34 @@ TEST(DynamicSparseAttentionTest, GroupedMergeActiveSplitBoundaries_CUDA) { } } +// ORT_DSA_INT8_VECTOR_STAGING and ORT_DSA_SHARED_SCALE_MERGE are cached in function-local statics, so +// the non-default INT8 grouped-decode staging and the per-channel merge are checked in fresh processes. +TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeAlternateSwitchesInFreshProcess_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + + const std::string filter = + "DynamicSparseAttentionTest.Int8CacheGroupedDecode_CUDA:" + "DynamicSparseAttentionTest.Int8CacheGroupedDecodeFloat16_CUDA:" + "DynamicSparseAttentionTest.Int8CacheGroupedDecodeVectorStagingChannels_CUDA:" + "DynamicSparseAttentionTest.GroupedMergeActiveSplitBoundaries_CUDA:" + "DynamicSparseAttentionTest.GroupedDecodeRatio12ReferenceParity_CUDA"; + + struct Mode { + const char* vector_staging; + const char* shared_scale_merge; + }; + const Mode modes[] = {{"0", "1"}, {"1", "0"}, {"0", "0"}}; + for (const Mode& mode : modes) { + SCOPED_TRACE(std::string("INT8_VECTOR_STAGING=") + mode.vector_staging + + " SHARED_SCALE_MERGE=" + mode.shared_scale_merge); + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{ + {"ORT_DSA_INT8_VECTOR_STAGING", mode.vector_staging}, + {"ORT_DSA_SHARED_SCALE_MERGE", mode.shared_scale_merge}, + }}; + ASSERT_EQ(RunGTestFilterInFreshProcess(filter), 0); + } +} + TEST(DynamicSparseAttentionTest, Int8CacheRejectsNonpositiveOrNonfiniteScales_CUDA) { ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{{"ORT_DYNAMIC_SPARSE_ATTENTION_STRICT_VALIDATION", "1"}}}; diff --git a/onnxruntime/test/contrib_ops/fresh_process_test_util.h b/onnxruntime/test/contrib_ops/fresh_process_test_util.h new file mode 100644 index 0000000000000..cb51ba9845721 --- /dev/null +++ b/onnxruntime/test/contrib_ops/fresh_process_test_util.h @@ -0,0 +1,44 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/common/common.h" + +#ifdef _WIN32 +#include +#else +#include +#include +#endif + +namespace onnxruntime { +namespace test { + +// Kernels cache some environment switches in function-local statics, so a switch can only be +// changed before the first eligible call in a process. This helper re-runs the selected gtest cases +// in a new process of the current test binary. The caller sets the environment variables (for +// example with ScopedEnvironmentVariables) before calling; the child inherits them. Returns the +// child exit status (0 on success). +inline int RunGTestFilterInFreshProcess(const std::string& gtest_filter) { +#ifdef _WIN32 + std::string path(MAX_PATH, '\0'); + const DWORD length = GetModuleFileNameA(nullptr, path.data(), static_cast(path.size())); + ORT_ENFORCE(length != 0 && length < path.size(), "GetModuleFileNameA failed."); + path.resize(length); +#else + std::string path(PATH_MAX, '\0'); + const ssize_t length = readlink("/proc/self/exe", path.data(), path.size()); + ORT_ENFORCE(length > 0 && static_cast(length) < path.size(), "readlink(/proc/self/exe) failed."); + path.resize(static_cast(length)); +#endif + const std::string command = "\"" + path + "\" --gtest_filter=" + gtest_filter + + " --gtest_fail_if_no_test_selected --gtest_color=no"; + return std::system(command.c_str()); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 8643770738dc5..94fa71f817763 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -32,11 +32,13 @@ #include "core/graph/constants.h" #include "core/graph/model.h" #include "test/common/tensor_op_test_utils.h" +#include "test/contrib_ops/fresh_process_test_util.h" #include "test/providers/provider_test_utils.h" #include "test/test_environment.h" #include "test/unittest_util/graph_transform_test_builder.h" #include "test/util/include/asserts.h" #include "test/util/include/default_providers.h" +#include "test/util/include/scoped_env_vars.h" #ifdef USE_WEBGPU #include "core/providers/webgpu/webgpu_provider_options.h" @@ -1539,6 +1541,48 @@ TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalScoringReusesKeysAcrossHea } } +// The CUDA indexer caches ORT_PACKED_SPARSE_INDEXER_{COMPACT_TOPK,RADIX_TOPK,WARP_HEADS} in +// function-local statics, so every alternate selection mode is exercised by re-running the +// hierarchical top-k and ragged-prefill cases (nonuniform scores, stable ties, partial merge groups) +// in a fresh process that has the switches set before the first eligible call. +TEST(PackedSparseAttentionIndexerTest, QsaAlternateTopKModesInFreshProcess) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + + const std::string filter = + "PackedSparseAttentionIndexerTest.QsaHierarchicalTopK:" + "PackedSparseAttentionIndexerTest.QsaHierarchicalPairedVerification:" + "PackedSparseAttentionIndexerTest.QsaCompactChunkAndFallbackBoundaries:" + "PackedSparseAttentionIndexerTest.QsaHierarchicalTileBoundariesPreserveStableTies:" + "PackedSparseAttentionIndexerTest.QsaHierarchicalPackedRowsWithLargeCapacity:" + "PackedSparseAttentionIndexerTest.QsaHierarchicalMergeLevelsWithPartialGroups:" + "PackedSparseAttentionIndexerTest.QsaHierarchicalScoringReusesKeysAcrossHeads:" + "PackedSparseAttentionIndexerTest.QsaRaggedPrefill"; + + struct Mode { + const char* compact; + const char* radix; + const char* warp_heads; + }; + const Mode modes[] = { + {"1", "0", "1"}, // explicit defaults + {"0", "0", "1"}, // tile-level two-level merge + {"0", "1", "1"}, // live-key radix selection + {"1", "0", "0"}, // non-warp ragged prefill scorer + }; + for (const Mode& mode : modes) { + SCOPED_TRACE(std::string("COMPACT_TOPK=") + mode.compact + " RADIX_TOPK=" + mode.radix + + " WARP_HEADS=" + mode.warp_heads); + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{ + {"ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK", mode.compact}, + {"ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", mode.radix}, + {"ORT_PACKED_SPARSE_INDEXER_WARP_HEADS", mode.warp_heads}, + }}; + ASSERT_EQ(RunGTestFilterInFreshProcess(filter), 0); + } +} + // state_status is optional: requesting it must not change any other output, and every healthy // request reports 0. TEST(PackedSparseAttentionIndexerTest, QsaStateStatusOkForValidRequests) { From 0639efa72c4a8c37fa64bcd054c4cd19e8075cab Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:53:49 +0000 Subject: [PATCH 21/23] docs: document DSA/indexer dispatch switches, INT8 encoding and status codes Give validation code 9 its documented name, state that the INT8 KV cache is symmetric signed int8 like GroupQueryAttention/PagedAttention, list the exact conditions and values for every CUDA dispatch switch, and sync ContribOperators.md with the schema (committed_length is [1] or [batch]; validation_status code list). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 8 ++- .../cuda/dynamic_sparse_attention.md | 61 ++++++++++++++++++ .../packed_sparse_attention_indexer.md | 63 +++++++++---------- .../core/graph/contrib_ops/bert_defs.cc | 12 +++- 4 files changed, 104 insertions(+), 40 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index acb492a5c8174..459e4cf71a98b 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1791,7 +1791,9 @@ This version of the operator has been available since version 1 of the 'com.micr 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. Newly appended keys (after optional normalization and RoPE) and values are divided by their scales, + 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. @@ -1869,7 +1871,7 @@ This version of the operator has been available since version 1 of the 'com.micr
head_sink (optional) : T
Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.
committed_length (optional) : M
-
Optional single-request committed token count for cursor validation. When supplied, validation_status is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent. CUDA only.
+
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
@@ -1886,7 +1888,7 @@ This version of the operator has been available since version 1 of the 'com.micr
present_value (optional) : TCache
Updated main value cache in BNSH layout.
validation_status (optional) : M
-
Optional CUDA-only int32 status: 0 when valid, a validation code when invalid.
+
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 diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index 0cbccd1bc25c7..ae3e1050f9deb 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -7,3 +7,64 @@ 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. diff --git a/docs/contrib_ops/packed_sparse_attention_indexer.md b/docs/contrib_ops/packed_sparse_attention_indexer.md index 615d04c086217..f0289b6aa6268 100644 --- a/docs/contrib_ops/packed_sparse_attention_indexer.md +++ b/docs/contrib_ops/packed_sparse_attention_indexer.md @@ -1,32 +1,31 @@ # PackedSparseAttentionIndexer — Operator Documentation -## Experimental CUDA decode specializations - -Set these process-level options before the first operator invocation. All three -default to `1`; only `0` and `1` are accepted. Existing dispatch remains available -when an option is disabled (`0`) or its geometry is unsupported. - -- `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1`: compact QSA chunks of 2048 keys - to 512 keys with device-side 512/1024/2048 sorting probes. Up to five chunks - use a fused final merge/emit; larger selections use the existing pairwise - merge. Stable tie ordering, causal tails, capacity stride/padding, and - overflow handling are preserved. This option is mutually exclusive with - `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1`. -- `ORT_DSA_INT8_VECTOR_STAGING=1`: the companion `DynamicSparseAttention` - grouped INT8 path uses eight-byte cache loads and independent per-channel - key/value scales. Its existing guards still require head size 256, FP16/BF16 - queries, one through eight tokens (`1 <= sequence_length <= 8`), GQA ratio at - least four, more than 256 candidates, and sufficient resources. Auxiliary - floating-point caches retain their own unscaled load path. -- `ORT_DSA_SHARED_SCALE_MERGE=1`: grouped attention merges with at most 32 - splits compute the exponential factor once per split and share it across - output channels. A nonzero partial sum, not the maximum-score sentinel, - determines whether a split is active. Active factors that underflow to zero - remain in the accumulation. - -These options change kernels, not graph schemas or cache layouts. Qualify -accuracy and performance on the target runtime, hardware, and workload before -enabling them in a deployment. +## CUDA QSA dispatch switches + +Three 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_RADIX_TOPK` | `0` (boolean) | hierarchical QSA TopK | `1` | Live-key radix selection (`QsaPartialTopKKernel`) over the tile-sorted keys, then `QsaEmitHierarchicalTopKKernel`. Requires `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=0`; enabling both fails with an error. | +| `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). | + +Effective hierarchical modes are therefore `COMPACT=1,RADIX=0` (default), `COMPACT=0,RADIX=0` and +`COMPACT=0,RADIX=1`. 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 +(`QsaAlternateTopKModesInFreshProcess`). 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 @@ -323,13 +322,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. -- CUDA has two additional dispatch controls. Set them before the first eligible - call; each value is cached per process. `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1` (default `0`, - mutually exclusive with `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=1`) uses live-key - radix selection for eligible short QSA calls. `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS=1` (default `1`) uses - warp-local scoring for eligible packed prefill calls with more than 64 tokens, four heads, - head size 128 and compression ratio four. Both retain causal visibility, stable tie ordering, - selected counts and state-update semantics. +- The CUDA dispatch switches `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK`, + `ORT_PACKED_SPARSE_INDEXER_RADIX_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/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index d31d2fdfc01d9..de18f6c841462 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -591,7 +591,9 @@ auxiliary_key is used as both K and V. DeepSeek-style post-attention inverse/con 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. Newly appended keys (after optional normalization and RoPE) and values are divided by their scales, +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. @@ -660,7 +662,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "T", OpSchema::Optional) .Input(17, "committed_length", "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. CUDA only.", + "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.", "M", OpSchema::Optional) .Input(18, "key_scale", "Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].", "S", OpSchema::Optional) @@ -669,7 +672,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Output(0, "output", "Attention output [batch, sequence, num_heads * head_size].", "T") .Output(1, "present_key", "Updated main key cache in BNSH layout.", "TCache", OpSchema::Optional) .Output(2, "present_value", "Updated main value cache in BNSH layout.", "TCache", OpSchema::Optional) - .Output(3, "validation_status", "Optional CUDA-only int32 status: 0 when valid, a validation code when invalid.", + .Output(3, "validation_status", + "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).", "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain all floating-point inputs and outputs to one element type.") From cd4e2e4291ab961b8ddff9c50da6068b8801ea37 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 23:37:52 +0000 Subject: [PATCH 22/23] [CUDA] Remove packed sparse indexer radix top-K path Benchmarks on the rc3 model (H200, DFlash2 S=8 decode) show ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK=1 is 0.15-0.34 ms/cycle slower than the default. Remove the env var, the packed-key QsaPartialTopKKernel specialization and the COMPACT/RADIX mutual-exclusion check; restore QsaPartialTopKKernel to its non-templated form. Alternate-mode fresh-process tests now cover only COMPACT_TOPK=0 and WARP_HEADS=0. Document the remaining switches with benchmark deltas; schema and ContribOperators.md are unaffected. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/dynamic_sparse_attention.md | 14 +++++ .../packed_sparse_attention_indexer.md | 28 +++++++--- .../packed_sparse_attention_indexer_impl.cu | 52 +++---------------- ...packed_sparse_attention_indexer_op_test.cc | 19 +++---- 4 files changed, 48 insertions(+), 65 deletions(-) diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index ae3e1050f9deb..8104053650175 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -68,3 +68,17 @@ 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. + +### 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 f0289b6aa6268..6c47de3386eb9 100644 --- a/docs/contrib_ops/packed_sparse_attention_indexer.md +++ b/docs/contrib_ops/packed_sparse_attention_indexer.md @@ -2,7 +2,7 @@ ## CUDA QSA dispatch switches -Three process-level environment variables select alternate CUDA QSA kernels. Each value is read once +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 @@ -17,15 +17,27 @@ prefill is eligible when `total_tokens > 64`, `head_size == 128`, `compress_rati |---|---|---|---|---| | `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_RADIX_TOPK` | `0` (boolean) | hierarchical QSA TopK | `1` | Live-key radix selection (`QsaPartialTopKKernel`) over the tile-sorted keys, then `QsaEmitHierarchicalTopKKernel`. Requires `ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK=0`; enabling both fails with an error. | | `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). | -Effective hierarchical modes are therefore `COMPACT=1,RADIX=0` (default), `COMPACT=0,RADIX=0` and -`COMPACT=0,RADIX=1`. The packed-QSA tests re-run the hierarchical, stable-tie, partial-merge-group +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 -(`QsaAlternateTopKModesInFreshProcess`). Qualify accuracy and performance on the target runtime, -hardware and workload before changing a default. +(`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 @@ -322,8 +334,8 @@ 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`, - `ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK` and `ORT_PACKED_SPARSE_INDEXER_WARP_HEADS` are described in +- 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 diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 753a74b4b79d3..1d57295963567 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -981,18 +981,8 @@ __global__ void QsaSelectKernel(const float* block_scores, const int32_t* topk_i } } -template -__device__ __forceinline__ uint64_t QsaSelectionKey(const float* scores, int index) { - if constexpr (PackedKeys) { - return reinterpret_cast(scores)[index]; - } else { - return topk::PackStableSortKey(scores[index], index); - } -} - // One block per packed query token. Radix-selects a bounded candidate set, then sorts only that // set. The output aliases the corresponding score row after all scores for that row have been read. -template __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, const int32_t* present_state_lengths, const int32_t* overflow_flags, @@ -1026,10 +1016,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat if (selected == 0) { continue; } - const int64_t score_stride = PackedKeys - ? 2 * ((static_cast(params.state_capacity) + kHierarchicalTopKTileBlocks - 1) / kHierarchicalTopKTileBlocks) * kHierarchicalTopKTileBlocks - : params.state_capacity; - float* scores = block_scores + static_cast(token) * score_stride; + float* scores = block_scores + static_cast(token) * params.state_capacity; if (threadIdx.x == 0) { prefix = 0; @@ -1044,7 +1031,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const uint32_t current_prefix = prefix; for (int index = threadIdx.x; index < block_count; index += blockDim.x) { - const uint32_t score_key = static_cast(QsaSelectionKey(scores, index) >> 32); + const uint32_t score_key = static_cast(topk::PackStableSortKey(scores[index], 0) >> 32); if (shift == 24 || (score_key >> (shift + 8)) == (current_prefix >> (shift + 8))) { atomicAdd(&temp.histogram[(score_key >> shift) & 0xffu], 1u); } @@ -1059,23 +1046,11 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat } else { prefix |= static_cast(bucket) << shift; remaining = rank; - if constexpr (PackedKeys) { - // A fully selected radix bucket needs no finer threshold. - if (rank == count && prefix != 0) { - --prefix; - remaining = 0; - } - } break; } } } __syncthreads(); - if constexpr (PackedKeys) { - if (remaining == 0) { - break; - } - } } if (threadIdx.x == 0) { @@ -1084,7 +1059,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const uint32_t threshold = prefix; for (int index = threadIdx.x; index < block_count; index += blockDim.x) { - const uint64_t key = QsaSelectionKey(scores, index); + const uint64_t key = topk::PackStableSortKey(scores[index], index); if (static_cast(key >> 32) > threshold) { const int slot = atomicAdd(&gathered, 1); selected_keys[slot] = key; @@ -1096,11 +1071,10 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat equal_seen = 0; } __syncthreads(); - // Sorted 32-key tiles still encounter equal scores in ascending original-index order. for (int base = 0; base < block_count && equal_seen < remaining; base += blockDim.x) { const int index = base + static_cast(threadIdx.x); const bool equal = index < block_count && - static_cast(QsaSelectionKey(scores, index) >> 32) == threshold; + static_cast(topk::PackStableSortKey(scores[index], 0) >> 32) == threshold; const unsigned int warp_mask = __ballot_sync(0xffffffffu, equal); const int lane = static_cast(threadIdx.x) % kWarpSize; const int warp = static_cast(threadIdx.x) / kWarpSize; @@ -1121,7 +1095,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat __syncthreads(); const int equal_rank = equal_seen + equal_warp_offsets[warp] + lane_rank; if (equal && equal_rank < remaining) { - selected_keys[gathered + equal_rank] = QsaSelectionKey(scores, index); + selected_keys[gathered + equal_rank] = topk::PackStableSortKey(scores[index], index); } __syncthreads(); if (threadIdx.x == 0) { @@ -1142,11 +1116,7 @@ __global__ void QsaPartialTopKKernel(float* block_scores, const int32_t* cumulat for (int item = 0; item < kBoundedTopKItemsPerThread; ++item) { const int rank = threadIdx.x * kBoundedTopKItemsPerThread + item; if (rank < selected) { - if constexpr (PackedKeys) { - reinterpret_cast(scores)[rank] = keys[item]; - } else { - topk_row[rank] = topk::UnpackStableSortIndex(keys[item]); - } + topk_row[rank] = topk::UnpackStableSortIndex(keys[item]); } } __syncthreads(); @@ -1643,14 +1613,10 @@ Status LaunchQsaPackedSparseAttentionIndexer( kHierarchicalTileScoreThreads, 0, stream>>>( present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, overflow_flags, merge_input, tile_count, params); - static const bool use_radix_topk = - ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", false); static const int compact_topk = ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK", 1); ORT_ENFORCE(compact_topk == 0 || compact_topk == 1, "ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK must be 0 or 1."); - ORT_ENFORCE(!compact_topk || !use_radix_topk, - "Compact and radix packed sparse indexer top-K options cannot both be enabled."); if (compact_topk) { const int chunks = (key_stride + qsa_topk::kChunkSize - 1) / qsa_topk::kChunkSize; const int compact_stride = chunks * kBoundedTopKMax; @@ -1675,10 +1641,6 @@ Status LaunchQsaPackedSparseAttentionIndexer( overflow_flags, compact_stride, selected_indices, selected_counts, params); } return CUDA_CALL(cudaGetLastError()); - } else if (use_radix_topk) { - QsaPartialTopKKernel<<>>( - selection_workspace, cumulative_sequence_lengths, past_sequence_lengths, position_ids, - present_state_lengths, overflow_flags, params); } else { for (int list_width = kHierarchicalTopKTileBlocks; list_width < key_stride;) { if (2 * list_width < key_stride) { @@ -1726,7 +1688,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( !use_hierarchical_topk && params.block_topk > kSaiFastTopKMax && params.block_topk <= kBoundedTopKMax; int32_t* topk_indices = use_bounded_topk ? reinterpret_cast(block_scores) : nullptr; if (use_bounded_topk) { - QsaPartialTopKKernel<<>>( + QsaPartialTopKKernel<<>>( block_scores, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, overflow_flags, params); } diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 94fa71f817763..d653b69b45135 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -1541,10 +1541,10 @@ TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalScoringReusesKeysAcrossHea } } -// The CUDA indexer caches ORT_PACKED_SPARSE_INDEXER_{COMPACT_TOPK,RADIX_TOPK,WARP_HEADS} in -// function-local statics, so every alternate selection mode is exercised by re-running the -// hierarchical top-k and ragged-prefill cases (nonuniform scores, stable ties, partial merge groups) -// in a fresh process that has the switches set before the first eligible call. +// The CUDA indexer caches ORT_PACKED_SPARSE_INDEXER_{COMPACT_TOPK,WARP_HEADS} in function-local +// statics, so each non-default mode (tile-level two-level merge, non-warp ragged prefill scorer) is +// exercised by re-running the hierarchical top-k and ragged-prefill cases (nonuniform scores, stable +// ties, partial merge groups) in a fresh process that has the switch set before the first eligible call. TEST(PackedSparseAttentionIndexerTest, QsaAlternateTopKModesInFreshProcess) { if (DefaultCudaExecutionProvider() == nullptr) { GTEST_SKIP() << "CUDA EP not available."; @@ -1562,21 +1562,16 @@ TEST(PackedSparseAttentionIndexerTest, QsaAlternateTopKModesInFreshProcess) { struct Mode { const char* compact; - const char* radix; const char* warp_heads; }; const Mode modes[] = { - {"1", "0", "1"}, // explicit defaults - {"0", "0", "1"}, // tile-level two-level merge - {"0", "1", "1"}, // live-key radix selection - {"1", "0", "0"}, // non-warp ragged prefill scorer + {"0", "1"}, // tile-level two-level merge + {"1", "0"}, // non-warp ragged prefill scorer }; for (const Mode& mode : modes) { - SCOPED_TRACE(std::string("COMPACT_TOPK=") + mode.compact + " RADIX_TOPK=" + mode.radix + - " WARP_HEADS=" + mode.warp_heads); + SCOPED_TRACE(std::string("COMPACT_TOPK=") + mode.compact + " WARP_HEADS=" + mode.warp_heads); ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{ {"ORT_PACKED_SPARSE_INDEXER_COMPACT_TOPK", mode.compact}, - {"ORT_PACKED_SPARSE_INDEXER_RADIX_TOPK", mode.radix}, {"ORT_PACKED_SPARSE_INDEXER_WARP_HEADS", mode.warp_heads}, }}; ASSERT_EQ(RunGTestFilterInFreshProcess(filter), 0); From c8c7b77c6223a2306347c8f7613f9b9a552e6c2e Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 27 Oct 2026 21:07:10 +0000 Subject: [PATCH 23/23] fix(cuda): guard DSA wide loads on buffer alignment and fix iOS test build DynamicSparseAttention: the default-on fast paths issued 16-byte (uint4/float4), 8-byte (uint2, cp.async) and 4-byte (bfloat162) accesses on caller-provided buffers. Tensors bound through IOBinding can start at any element offset, so a misaligned scale, cache, auxiliary or output buffer faulted with cudaErrorMisalignedAddress (716) and poisoned the CUDA context. - INT8 grouped decode: use vector staging only when the scales, auxiliary K/V and INT8 caches are aligned; otherwise use the scalar staging kernel. - Floating grouped decode and grouped prefill: StageHeadRow copies a K/V row with one 16-byte chunk per lane only when the row is 16-byte aligned and element-wise otherwise, so the workspace layout is unchanged. - Tensor-core prefill kernel: also require aligned caches, INT8 scales and output; otherwise the generic kernels run. - Add IOBinding tests that offset each cache, scale, auxiliary and output buffer for INT8 and floating caches, for decode and prefill, and document the alignment requirements and fallbacks. iOS CI: fresh_process_test_util.h uses std::system, which is unavailable on iOS. Include it and define QsaAlternateTopKModesInFreshProcess only under USE_CUDA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/dynamic_sparse_attention.md | 16 + .../bert/dynamic_sparse_attention_impl.cu | 53 ++- .../dynamic_sparse_attention_op_test.cc | 365 +++++++++++++++++- .../contrib_ops/fresh_process_test_util.h | 3 + ...packed_sparse_attention_indexer_op_test.cc | 8 +- 5 files changed, 425 insertions(+), 20 deletions(-) diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index 8104053650175..86ca09455b0a3 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -69,6 +69,22 @@ When `split_count > 32` (floating-point caches only) the generic `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 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 17cd24817eae8..a612c7071aab9 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -76,6 +77,20 @@ __device__ __forceinline__ void MmaM16N8K16(float (&output)[4], #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) { const int32_t previous = @@ -821,10 +836,8 @@ __global__ void GroupedPrefillDynamicSparseAttentionKernel( const T* value_head = candidate_is_auxiliary ? auxiliary_value + offset : main_value + offset; - reinterpret_cast(shared_key + shared_offset)[load_lane] = - reinterpret_cast(key_head)[load_lane]; - reinterpret_cast(shared_value + shared_offset)[load_lane] = - reinterpret_cast(value_head)[load_lane]; + StageHeadRow(shared_key + shared_offset, key_head, load_lane, head_size); + StageHeadRow(shared_value + shared_offset, value_head, load_lane, head_size); } } __syncthreads(); @@ -1156,10 +1169,8 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( } 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; - reinterpret_cast(shared_key + shared_offset)[load_lane] = - reinterpret_cast(key_head)[load_lane]; - reinterpret_cast(shared_value + shared_offset)[load_lane] = - reinterpret_cast(value_head)[load_lane]; + StageHeadRow(shared_key + shared_offset, key_head, load_lane, head_size); + StageHeadRow(shared_value + shared_offset, value_head, load_lane, head_size); } } } @@ -1711,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) { @@ -2068,8 +2084,16 @@ Status LaunchDynamicSparseAttention( CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device)); CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); constexpr int kMaxGridDimYZ = 65535; + // The kernel issues 16-byte (INT8: 8-byte) cp.async copies from the caches, float4 loads of the INT8 + // scales and bfloat162 stores to the output. Tensors bound through IOBinding can start at any element + // offset, so misaligned buffers use the generic kernels instead of faulting. + const size_t cache_alignment = data.quantized_cache ? 8 : 16; + const bool buffers_aligned = + IsAligned(data.present_key, cache_alignment) && IsAligned(data.present_value, cache_alignment) && + IsAligned(data.output, 4) && + (!data.quantized_cache || (IsAligned(data.key_scale, 16) && IsAligned(data.value_scale, 16))); use_prefill_kernel = - enabled && cc_major >= 8 && parameters.sequence_length >= 8 && + enabled && buffers_aligned && cc_major >= 8 && parameters.sequence_length >= 8 && parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && parameters.selected_kv_source == DynamicSparseAttentionKvSource::kMain && !parameters.use_smooth_softmax && parameters.max_selected > 0 && @@ -2099,6 +2123,13 @@ Status LaunchDynamicSparseAttention( ParseEnvironmentVariableWithDefault("ORT_DSA_INT8_VECTOR_STAGING", 1); ORT_ENFORCE(vectorized_staging == 0 || vectorized_staging == 1, "ORT_DSA_INT8_VECTOR_STAGING must be 0 or 1."); + // Vector staging loads 16-byte auxiliary K/V and scale chunks and 8-byte INT8 cache chunks. Tensors bound + // through IOBinding can start at any element offset, so misaligned buffers use scalar staging. + const bool use_vector_staging = + vectorized_staging != 0 && + IsAligned(data.key_scale, 16) && IsAligned(data.value_scale, 16) && + IsAligned(data.auxiliary_key, 16) && IsAligned(data.auxiliary_value, 16) && + IsAligned(data.present_key, 8) && IsAligned(data.present_value, 8); const auto launch = [&](auto heads_tag, auto vectorized) { constexpr int kHeads = decltype(heads_tag)::value; GroupedDecodeDynamicSparseAttentionKernel{}, std::true_type{}); } else { launch(std::integral_constant{}, std::false_type{}); } } else { - if (vectorized_staging) { + if (use_vector_staging) { launch(std::integral_constant{}, std::true_type{}); } else { launch(std::integral_constant{}, std::false_type{}); diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 14bdc9ed6741c..003012a5a6a1b 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -18,7 +18,6 @@ #if defined(USE_CUDA) || defined(USE_WEBGPU) #include "test/common/tensor_op_test_utils.h" -#include "test/contrib_ops/fresh_process_test_util.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/scoped_env_vars.h" #include "test/util/include/default_providers.h" @@ -30,8 +29,11 @@ #endif #ifdef USE_CUDA +#include "core/framework/to_tensor_proto_element_type.h" #include "core/providers/cuda/cuda_provider_options.h" #include "test/common/cuda_op_test_utils.h" +// Uses std::system, which is unavailable on iOS, so only CUDA builds (Windows/Linux) include it. +#include "test/contrib_ops/fresh_process_test_util.h" #endif #ifdef USE_WEBGPU @@ -591,10 +593,10 @@ TEST(DynamicSparseAttentionTest, GroupedDecodeRatio12ReferenceParity_CUDA) { } } -TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { - if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; - if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; - auto c = MakeGroupedCacheCase(9, false, false, 2, 20, 2, true); +// Selected-only prefill without a softmax sink, the configuration the tensor-core BF16 prefill kernel serves. +DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache) { + auto c = int8_cache ? MakeGroupedCacheCase(9, false, false, 2, 20, 2, true) + : MakeGroupedCacheCase(9, false, true, 2, 20, 2); c.head_sink.clear(); c.smooth_softmax = 0; for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { @@ -609,13 +611,14 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { const int64_t offset = ((b * c.kv_num_heads + kv) * c.cache_sequence_length + index) * c.head_size; double dot = 0.0; for (int d = 0; d < c.head_size; ++d) { - dot += c.query[out + d] * c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d]; + const float key_scale = int8_cache ? c.key_scale[kv * c.head_size + d] : 1.0f; + dot += c.query[out + d] * c.expected_present_key[offset + d] * key_scale; } const double weight = std::exp(dot * c.scale); denominator += weight; for (int d = 0; d < c.head_size; ++d) { - numerator[d] += static_cast(weight * c.expected_present_value[offset + d] * - c.value_scale[kv * c.head_size + d]); + const float value_scale = int8_cache ? c.value_scale[kv * c.head_size + d] : 1.0f; + numerator[d] += static_cast(weight * c.expected_present_value[offset + d] * value_scale); } } for (int d = 0; d < c.head_size; ++d) { @@ -623,6 +626,13 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { } } } + return c; +} + +TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + const auto c = MakeSelectedOnlyPrefillCase(true); for (const char* prefill_env : {"1", "0"}) { SCOPED_TRACE(prefill_env); ScopedEnvironmentVariables scoped_env_vars{ @@ -2188,6 +2198,345 @@ TEST(DynamicSparseAttentionTest, RejectsUnknownModeAndSource_CUDA) { } #endif +// Tensors bound through IOBinding can start at any element offset inside a larger allocation, so the kernels +// must not assume more than element alignment from caller-provided buffers. The tests below place each such +// buffer at a misaligned address and expect the aligned-buffer results instead of cudaErrorMisalignedAddress. +namespace { + +// Byte offsets of the caller-provided device buffers from a 256-byte aligned allocation. +struct BufferOffsets { + size_t key_cache = 0; + size_t value_cache = 0; + size_t auxiliary_key = 0; + size_t auxiliary_value = 0; + size_t key_scale = 0; + size_t value_scale = 0; + size_t output = 0; +}; + +struct OffsetBufferCase { + const char* name; + BufferOffsets offsets; +}; + +OffsetBufferCase MisalignedBuffer(const char* name, size_t BufferOffsets::* buffer, size_t bytes) { + OffsetBufferCase offset_case{name, {}}; + offset_case.offsets.*buffer = bytes; + return offset_case; +} + +// Offsets are the smallest that keep each buffer naturally aligned for its element type. +std::vector MakeOffsetBufferCases(bool int8_cache, bool auxiliary) { + const size_t cache_bytes = int8_cache ? 1 : 2; + const size_t scale_bytes = 4; + const size_t half_bytes = 2; + std::vector cases{ + {"aligned", {}}, + MisalignedBuffer("key_cache", &BufferOffsets::key_cache, cache_bytes), + MisalignedBuffer("value_cache", &BufferOffsets::value_cache, cache_bytes), + MisalignedBuffer("output", &BufferOffsets::output, half_bytes)}; + BufferOffsets all; + all.key_cache = all.value_cache = cache_bytes; + all.output = half_bytes; + if (int8_cache) { + cases.push_back(MisalignedBuffer("key_scale", &BufferOffsets::key_scale, scale_bytes)); + cases.push_back(MisalignedBuffer("value_scale", &BufferOffsets::value_scale, scale_bytes)); + all.key_scale = all.value_scale = scale_bytes; + } + if (auxiliary) { + cases.push_back(MisalignedBuffer("auxiliary_key", &BufferOffsets::auxiliary_key, half_bytes)); + cases.push_back(MisalignedBuffer("auxiliary_value", &BufferOffsets::auxiliary_value, half_bytes)); + all.auxiliary_key = all.auxiliary_value = half_bytes; + } + cases.push_back({"all", all}); + return cases; +} + +// Allocates device tensors whose data pointers sit a chosen number of bytes into their allocations. +class OffsetDeviceBuffers { + public: + OffsetDeviceBuffers(IExecutionProvider& ep, AllocatorPtr device_allocator, AllocatorPtr host_allocator) + : ep_(ep), device_allocator_(std::move(device_allocator)), host_allocator_(std::move(host_allocator)) {} + + template + OrtValue Empty(const TensorShape& shape, size_t byte_offset) { + const size_t bytes = static_cast(shape.Size()) * sizeof(Element); + allocations_.push_back(IAllocator::MakeUniquePtr(device_allocator_, bytes + byte_offset)); + uint8_t* base = allocations_.back().get(); + ORT_ENFORCE(reinterpret_cast(base) % 256 == 0, "Device allocation is not 256-byte aligned."); + OrtValue value; + Tensor::InitOrtValue(DataTypeImpl::GetType(), shape, base + byte_offset, + device_allocator_->Info(), value); + return value; + } + + template + OrtValue FromHost(const std::vector& values, const TensorShape& shape, size_t byte_offset = 0) { + OrtValue value = Empty(shape, byte_offset); + Tensor host_tensor(DataTypeImpl::GetType(), shape, const_cast(values.data()), + host_allocator_->Info()); + ORT_THROW_IF_ERROR(ep_.GetDataTransfer()->CopyTensor(host_tensor, *value.GetMutable())); + return value; + } + + template + std::vector ToHost(const OrtValue& device_value) { + const Tensor& device_tensor = device_value.Get(); + Tensor host_tensor(DataTypeImpl::GetType(), device_tensor.Shape(), host_allocator_); + ORT_THROW_IF_ERROR(ep_.GetDataTransfer()->CopyTensor(device_tensor, host_tensor)); + const auto values = host_tensor.DataAsSpan(); + return std::vector(values.begin(), values.end()); + } + + private: + IExecutionProvider& ep_; + AllocatorPtr device_allocator_; + AllocatorPtr host_allocator_; + std::vector> allocations_; +}; + +template +float ElementToFloat(Element value) { + if constexpr (std::is_arithmetic_v) { + return static_cast(value); + } else { + return value.ToFloat(); + } +} + +template +void ExpectMatchesExpected(const std::vector& actual, const std::vector& expected, + const char* name) { + ASSERT_EQ(actual.size(), expected.size()) << name; + // Like OpTester, compare against the expected values rounded to the tensor type. + const std::vector rounded = ToTensorData(expected); + for (size_t i = 0; i < actual.size(); ++i) { + const float expected_value = ElementToFloat(rounded[i]); + ASSERT_NEAR(ElementToFloat(actual[i]), expected_value, 0.005f + 0.01f * std::abs(expected_value)) + << name << " at " << i; + } +} + +// Runs the case through IOBinding with the cache buffers bound in place (present aliases past). +template +void RunCaseWithOffsetBuffers(const DynamicSparseAttentionCase& c, const BufferOffsets& offsets) { + auto ep = DefaultCudaExecutionProvider(); + ASSERT_NE(ep, nullptr); + + const bool has_auxiliary = c.auxiliary_sequence_length > 0; + const bool has_scales = !c.key_scale.empty(); + const int64_t query_count = c.batch_size * c.sequence_length; + const TensorShape query_shape{c.batch_size, c.sequence_length, c.num_heads * c.head_size}; + const TensorShape kv_shape{c.batch_size, c.sequence_length, c.kv_num_heads * c.head_size}; + const TensorShape cache_shape{c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}; + const TensorShape auxiliary_shape{c.batch_size, c.kv_num_heads, c.auxiliary_sequence_length, c.head_size}; + + Model model("dynamic_sparse_attention_offset_buffers", true, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 17}, {kMSDomain, 1}}, + {}, DefaultLoggingManager().DefaultLogger(), ModelOptions(true, true)); + auto& graph = model.MainGraph(); + const auto make_type = [](ONNX_NAMESPACE::TensorProto_DataType element_type) { + ONNX_NAMESPACE::TypeProto type; + type.mutable_tensor_type()->set_elem_type(element_type); + return type; + }; + const auto t_type = make_type(utils::ToTensorProtoElementType()); + const auto cache_type = make_type(utils::ToTensorProtoElementType()); + const auto int32_type = make_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); + const auto float_type = make_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + auto& empty = graph.GetOrCreateNodeArg("", nullptr); + const auto arg = [&graph](const char* name, const ONNX_NAMESPACE::TypeProto& type) { + return &graph.GetOrCreateNodeArg(name, &type); + }; + + std::vector inputs{ + arg("query", t_type), + arg("key", t_type), + arg("value", t_type), + arg("past_key", cache_type), + arg("past_value", cache_type), + has_auxiliary ? arg("auxiliary_key", t_type) : &empty, + has_auxiliary ? arg("auxiliary_value", t_type) : &empty, + arg("selected_indices", int32_type), + arg("selected_counts", int32_type), + arg("seqlens_k", int32_type), + arg("total_sequence_length", int32_type), + &empty, // cos_cache + &empty, // sin_cache + &empty, // position_ids + &empty, // q_norm_weight + &empty, // k_norm_weight + c.head_sink.empty() ? &empty : arg("head_sink", t_type)}; + if (has_scales) { + inputs.push_back(&empty); // committed_length + inputs.push_back(arg("key_scale", float_type)); + inputs.push_back(arg("value_scale", float_type)); + } + std::vector outputs{ + arg("output", t_type), arg("present_key", cache_type), arg("present_value", cache_type)}; + auto& node = graph.AddNode("dynamic_sparse_attention", "DynamicSparseAttention", "", + inputs, outputs, nullptr, kMSDomain); + node.AddAttribute("num_heads", c.num_heads); + node.AddAttribute("kv_num_heads", c.kv_num_heads); + node.AddAttribute("scale", c.scale); + node.AddAttribute("local_window_size", c.local_window_size); + node.AddAttribute("attention_mode", c.attention_mode); + node.AddAttribute("selected_kv_source", c.selected_kv_source); + node.AddAttribute("smooth_softmax", c.smooth_softmax); + ASSERT_STATUS_OK(graph.Resolve()); + + std::string model_data; + ASSERT_TRUE(model.ToProto().SerializeToString(&model_data)); + SessionOptions session_options; + InferenceSession session(session_options, GetEnvironment()); + IExecutionProvider* ep_ptr = ep.get(); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(ep))); + std::istringstream model_stream(model_data); + ASSERT_STATUS_OK(session.Load(model_stream)); + ASSERT_STATUS_OK(session.Initialize()); + + auto device_allocators = ep_ptr->CreatePreferredAllocators(); + auto device_allocator = std::find_if(device_allocators.begin(), device_allocators.end(), [](const auto& allocator) { + return allocator->Info().device.Type() == OrtDevice::GPU && + allocator->Info().alloc_type != OrtAllocatorType::OrtReadOnlyAllocator && + allocator->Info().mem_type == OrtMemTypeDefault; + }); + ASSERT_NE(device_allocator, device_allocators.end()); + auto allocator = session.GetAllocator((*device_allocator)->Info()); + ASSERT_NE(allocator, nullptr); + auto cpu_allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + + // Declared after the session so the device allocations are released before it is destroyed. + OffsetDeviceBuffers buffers(*ep_ptr, allocator, cpu_allocator); + std::unique_ptr binding; + ASSERT_STATUS_OK(session.NewIOBinding(&binding)); + const auto bind_input = [&](const char* name, const OrtValue& value) { + ORT_THROW_IF_ERROR(binding->BindInput(name, value)); + }; + bind_input("query", buffers.FromHost(ToTensorData(c.query), query_shape)); + bind_input("key", buffers.FromHost(ToTensorData(c.key), kv_shape)); + bind_input("value", buffers.FromHost(ToTensorData(c.value), kv_shape)); + OrtValue key_cache = buffers.FromHost(ToTensorData(c.past_key), cache_shape, offsets.key_cache); + OrtValue value_cache = buffers.FromHost(ToTensorData(c.past_value), cache_shape, offsets.value_cache); + bind_input("past_key", key_cache); + bind_input("past_value", value_cache); + if (has_auxiliary) { + ASSERT_EQ(c.auxiliary_value.size(), c.auxiliary_key.size()); + bind_input("auxiliary_key", + buffers.FromHost(ToTensorData(c.auxiliary_key), auxiliary_shape, offsets.auxiliary_key)); + bind_input("auxiliary_value", + buffers.FromHost(ToTensorData(c.auxiliary_value), auxiliary_shape, offsets.auxiliary_value)); + } + bind_input("selected_indices", + buffers.FromHost(c.selected_indices, TensorShape{query_count, c.max_selected})); + bind_input("selected_counts", buffers.FromHost(c.selected_counts, TensorShape{query_count})); + bind_input("seqlens_k", buffers.FromHost(c.seqlens_k, TensorShape{c.batch_size})); + std::vector total_length_data{c.total_sequence_length}; + OrtValue total_length_value; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape{}, total_length_data.data(), + cpu_allocator->Info(), total_length_value); + bind_input("total_sequence_length", total_length_value); + if (!c.head_sink.empty()) { + bind_input("head_sink", buffers.FromHost(ToTensorData(c.head_sink), TensorShape{c.num_heads})); + } + if (has_scales) { + const TensorShape scale_shape{static_cast(c.key_scale.size())}; + bind_input("key_scale", buffers.FromHost(c.key_scale, scale_shape, offsets.key_scale)); + bind_input("value_scale", buffers.FromHost(c.value_scale, scale_shape, offsets.value_scale)); + } + ASSERT_STATUS_OK(binding->BindOutput("output", buffers.Empty(query_shape, offsets.output))); + ASSERT_STATUS_OK(binding->BindOutput("present_key", key_cache)); + ASSERT_STATUS_OK(binding->BindOutput("present_value", value_cache)); + + ASSERT_STATUS_OK(binding->SynchronizeInputs()); + RunOptions run_options; + ASSERT_STATUS_OK(session.Run(run_options, *binding)); + ASSERT_STATUS_OK(binding->SynchronizeOutputs()); + + const auto& results = binding->GetOutputs(); + ASSERT_EQ(results.size(), 3u); + ExpectMatchesExpected(buffers.ToHost(results[0]), c.expected_output, "output"); + ExpectMatchesExpected(buffers.ToHost(results[1]), c.expected_present_key, "present_key"); + ExpectMatchesExpected(buffers.ToHost(results[2]), c.expected_present_value, "present_value"); +} + +template +void ExpectMisalignedBuffersMatch(const DynamicSparseAttentionCase& c, bool int8_cache) { + for (const auto& offset_case : MakeOffsetBufferCases(int8_cache, c.auxiliary_sequence_length > 0)) { + SCOPED_TRACE(offset_case.name); + RunCaseWithOffsetBuffers(c, offset_case.offsets); + // A misaligned access poisons the CUDA context, so later cases would only add noise. + if (::testing::Test::HasFailure()) return; + } +} + +} // namespace + +TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeMisalignedBuffers_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + for (bool bf16 : {true, false}) { + if (bf16 && !CudaHasBF16Support()) continue; + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("bf16=" + std::to_string(bf16) + ", auxiliary=" + std::to_string(auxiliary)); + // Twelve query heads share one KV head: the grouped decode kernel with INT8 vector staging. + const auto c = MakeGroupedCacheCase(2, auxiliary, false, 1, 12, 1, true); + if (bf16) { + ExpectMisalignedBuffersMatch(c, true); + } else { + ExpectMisalignedBuffersMatch(c, true); + } + if (::testing::Test::HasFailure()) return; + } + } +} + +TEST(DynamicSparseAttentionTest, Int8CachePrefillMisalignedBuffers_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + // Aligned buffers use the tensor-core prefill kernel; misaligned ones must fall back to the generic kernels. + ExpectMisalignedBuffersMatch(MakeSelectedOnlyPrefillCase(true), true); +} + +TEST(DynamicSparseAttentionTest, FloatingCacheGroupedDecodeMisalignedBuffers_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + for (bool bf16 : {true, false}) { + if (bf16 && !CudaHasBF16Support()) continue; + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("bf16=" + std::to_string(bf16) + ", auxiliary=" + std::to_string(auxiliary)); + const auto c = MakeGroupedCacheCase(2, auxiliary, true, 1, 12, 1); + if (bf16) { + ExpectMisalignedBuffersMatch(c, false); + } else { + ExpectMisalignedBuffersMatch(c, false); + } + if (::testing::Test::HasFailure()) return; + } + } +} + +TEST(DynamicSparseAttentionTest, FloatingCachePrefillMisalignedBuffers_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (CudaHasBF16Support()) { + SCOPED_TRACE("bf16 selected-only without sink"); + // Aligned buffers use the tensor-core prefill kernel; misaligned ones fall back to grouped prefill. + ExpectMisalignedBuffersMatch(MakeSelectedOnlyPrefillCase(false), false); + if (::testing::Test::HasFailure()) return; + } + for (bool bf16 : {true, false}) { + if (bf16 && !CudaHasBF16Support()) continue; + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("grouped prefill: bf16=" + std::to_string(bf16) + ", auxiliary=" + std::to_string(auxiliary)); + const auto c = MakeGroupedCacheCase(9, auxiliary, true, 2, 20, 2); + if (bf16) { + ExpectMisalignedBuffersMatch(c, false); + } else { + ExpectMisalignedBuffersMatch(c, false); + } + if (::testing::Test::HasFailure()) return; + } + } +} + #endif // USE_CUDA #ifdef USE_WEBGPU diff --git a/onnxruntime/test/contrib_ops/fresh_process_test_util.h b/onnxruntime/test/contrib_ops/fresh_process_test_util.h index cb51ba9845721..65a800ed39a22 100644 --- a/onnxruntime/test/contrib_ops/fresh_process_test_util.h +++ b/onnxruntime/test/contrib_ops/fresh_process_test_util.h @@ -23,6 +23,9 @@ namespace test { // in a new process of the current test binary. The caller sets the environment variables (for // example with ScopedEnvironmentVariables) before calling; the child inherits them. Returns the // child exit status (0 on success). +// +// Only Windows and Linux are supported (std::system is unavailable on iOS and /proc/self/exe is +// Linux-specific), so include this header only from tests compiled under USE_CUDA. inline int RunGTestFilterInFreshProcess(const std::string& gtest_filter) { #ifdef _WIN32 std::string path(MAX_PATH, '\0'); diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index d653b69b45135..31fbcdac79262 100644 --- a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc @@ -32,7 +32,6 @@ #include "core/graph/constants.h" #include "core/graph/model.h" #include "test/common/tensor_op_test_utils.h" -#include "test/contrib_ops/fresh_process_test_util.h" #include "test/providers/provider_test_utils.h" #include "test/test_environment.h" #include "test/unittest_util/graph_transform_test_builder.h" @@ -40,6 +39,11 @@ #include "test/util/include/default_providers.h" #include "test/util/include/scoped_env_vars.h" +#ifdef USE_CUDA +// Uses std::system, which is unavailable on iOS, so only CUDA builds (Windows/Linux) include it. +#include "test/contrib_ops/fresh_process_test_util.h" +#endif + #ifdef USE_WEBGPU #include "core/providers/webgpu/webgpu_provider_options.h" #endif @@ -1541,6 +1545,7 @@ TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalScoringReusesKeysAcrossHea } } +#ifdef USE_CUDA // The CUDA indexer caches ORT_PACKED_SPARSE_INDEXER_{COMPACT_TOPK,WARP_HEADS} in function-local // statics, so each non-default mode (tile-level two-level merge, non-warp ragged prefill scorer) is // exercised by re-running the hierarchical top-k and ragged-prefill cases (nonuniform scores, stable @@ -1577,6 +1582,7 @@ TEST(PackedSparseAttentionIndexerTest, QsaAlternateTopKModesInFreshProcess) { ASSERT_EQ(RunGTestFilterInFreshProcess(filter), 0); } } +#endif // USE_CUDA // state_status is optional: requesting it must not change any other output, and every healthy // request reports 0.