From 9b8aaa0c68360817a2fb160b056c717f1daa6329 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 17:56:51 +0000 Subject: [PATCH 1/2] Extend sparse attention caches and optimize packed indexing Part of the grouped generic CUDA operator stack. Preserve the qualified source snapshot. --- .../packed_sparse_attention_indexer_common.h | 4 + .../cuda/bert/dynamic_sparse_attention.cc | 68 +++- .../bert/dynamic_sparse_attention_impl.cu | 338 ++++++++++++----- .../cuda/bert/dynamic_sparse_attention_impl.h | 19 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 14 +- .../sparse/packed_sparse_attention_indexer.cc | 79 ++-- .../sparse/packed_sparse_attention_indexer.h | 2 +- .../packed_sparse_attention_indexer_impl.cu | 240 ++++++++++--- .../packed_sparse_attention_indexer_impl.h | 15 +- .../webgpu/bert/dynamic_sparse_attention.cc | 7 + .../bert/packed_sparse_attention_indexer.cc | 6 +- .../webgpu/bert/sparse_attention_indexer.h | 8 +- .../core/graph/contrib_ops/bert_defs.cc | 103 +++++- .../dynamic_sparse_attention_op_test.cc | 339 +++++++++++++++++- ...packed_sparse_attention_indexer_op_test.cc | 85 ++++- .../sparse_attention_indexer_op_test.cc | 5 +- 16 files changed, 1107 insertions(+), 225 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 c1a467fe3773d..fdc636bc9626d 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..b11283cff7599 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -28,6 +28,8 @@ using namespace onnxruntime::cuda; kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TCache", BuildKernelDefConstraints()) \ + .TypeConstraint("S", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ .MayInplace(3, 1) \ .MayInplace(4, 2) \ @@ -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 1b5a2b0404f53..0e3e8746de44d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -30,6 +31,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) { @@ -41,6 +44,10 @@ __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); @@ -85,13 +92,24 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, uint32_t* validation_hash_table, size_t validation_hash_entries, bool use_shared_hash_table, - int32_t* error_flag) { + 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_hash_table[]; if (use_shared_hash_table) { @@ -339,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; @@ -405,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]; + } } } @@ -434,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; @@ -485,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) { @@ -503,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; } } } @@ -511,14 +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 : -FLT_MAX; + logits[candidate] = valid ? dot * scale : -CUDART_INF_F; } } __syncthreads(); @@ -547,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 = logit == -FLT_MAX ? 0.0f : expf(logit - max_logit); + const float weight = logit == -CUDART_INF_F ? 0.0f : expf(logit - max_logit); logits[candidate] = weight; thread_sum += weight; } @@ -566,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 (weight == 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; @@ -586,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 = @@ -844,7 +892,7 @@ __global__ void GroupedPrefillDynamicSparseAttentionKernel( } } -template +template __global__ void GroupedDecodeDynamicSparseAttentionKernel( const T* query, const T* main_key, @@ -860,6 +908,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( float* partial_output, int grouped_block_count, int split_count, + int sequence_length, int num_heads, int kv_num_heads, int head_size, @@ -870,7 +919,9 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( float scale, bool local_plus_selected, bool selected_from_auxiliary, - bool use_smooth_softmax) { + 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; @@ -880,16 +931,20 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( 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 b = grouped_block / (kv_num_heads * head_tiles); - const int kv_tile = grouped_block - b * kv_num_heads * head_tiles; + 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 int total_length = seqlens_k[b] + 1; - const int query_position = total_length - 1; + 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; @@ -898,15 +953,18 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( 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[b], max_selected)); + const int selected_count = max(0, min(selected_counts[row], max_selected)); const int candidate_count = local_count + selected_count; - const int candidate_begin = split * kGroupedDecodeCandidatesPerSplit; + 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(b) * max_selected; + 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; + 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]; @@ -917,7 +975,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( float query_values[kValuesPerLane]{}; float accumulator[kValuesPerLane]{}; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { + 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; @@ -926,7 +984,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( 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(b) * num_heads + + ? query[(static_cast(row) * num_heads + kv_head * query_heads_per_kv + staged_group_head) * head_size + head_offset] @@ -936,7 +994,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( } else #endif if (active_head) { - const T* query_head = query + (static_cast(b) * num_heads + head) * head_size; + 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]); @@ -1026,22 +1084,34 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( ((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]; + 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) { + 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; @@ -1077,7 +1147,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( const int shared_offset = tile_candidate * head_size; float score; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v && !kQuantizedCache) { score = mma_scores[static_cast(threadIdx.x) / 32 * kGroupedPrefillCandidatesPerTile + tile_candidate] * @@ -1110,13 +1180,13 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( __syncthreads(); candidate_start += kGroupedPrefillCandidatesPerTile; - if (candidate_start % kGroupedDecodeCandidatesPerSplit == 0) { - candidate_start += (split_count - 1) * kGroupedDecodeCandidatesPerSplit; + if (candidate_start % kCandidatesPerSplit == 0) { + candidate_start += (split_count - 1) * kCandidatesPerSplit; } } if (active_head) { - const int partial_block = (b * num_heads + head) * split_count + split; + 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; @@ -1155,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; @@ -1221,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) { @@ -1239,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; } } } @@ -1247,14 +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 : -FLT_MAX; + logits[split_candidate] = valid ? dot * scale : -CUDART_INF_F; } } __syncthreads(); @@ -1281,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 = logit == -FLT_MAX ? 0.0f : expf(logit - next_max); + const float weight = logit == -CUDART_INF_F ? 0.0f : expf(logit - next_max); logits[candidate] = weight; thread_sum += weight; } @@ -1300,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 (weight == 0.0f) { - continue; - } - const int candidate = candidate_start + split_candidate; int index; const T* value_head; @@ -1311,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; @@ -1321,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; @@ -1387,7 +1476,7 @@ __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]; - if (split_max != -FLT_MAX) { + if (partial_sum[partial_start + split] != 0.0f) { thread_sum += partial_sum[partial_start + split] * expf(split_max - max_logit); } } @@ -1406,7 +1495,7 @@ __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]; - if (split_max != -FLT_MAX) { + 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; @@ -1450,7 +1539,7 @@ __global__ void MergeGroupedDecodeDynamicSparseAttentionKernel(const float* part if (h < 32) { float value = 0.0f; - if (lane < split_count && partial_max[partial_start + lane] != -FLT_MAX) { + 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); } @@ -1467,7 +1556,7 @@ __global__ void MergeGroupedDecodeDynamicSparseAttentionKernel(const float* part float accumulator = 0.0f; for (int split = 0; split < split_count; ++split) { const float split_max = partial_max[partial_start + split]; - if (split_max != -FLT_MAX) { + 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); @@ -1558,22 +1647,28 @@ bool UseGroupedPrefillAttention(const DynamicSparseAttentionParameters& paramete bool UseGroupedDecodeAttention(const DynamicSparseAttentionParameters& parameters, size_t element_size, - size_t max_shared_memory_per_block) { - return parameters.sequence_length == 1 && + 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); + 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) { +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) + kGroupedDecodeCandidatesPerSplit - 1) / - kGroupedDecodeCandidatesPerSplit; + (GetCandidateCapacity(parameters) + candidates_per_split - 1) / candidates_per_split; const int64_t grouped_head_blocks = - static_cast(parameters.batch_size) * parameters.kv_num_heads * + static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads * GetGroupedHeadTiles(parameters, kGroupedDecodeQueryHeadsPerBlock); const int64_t occupancy_splits = - std::max(1, (kTargetGroupedDecodeBlocks + grouped_head_blocks - 1) / + std::max(1, (target_blocks + grouped_head_blocks - 1) / grouped_head_blocks); return static_cast(std::min(candidate_splits, occupancy_splits)); } @@ -1599,6 +1694,23 @@ 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( @@ -1616,7 +1728,19 @@ 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) { + 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; } @@ -1647,7 +1771,9 @@ 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; @@ -1675,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_hash_entries, use_shared_hash_table, 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) { @@ -1716,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; } @@ -1723,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, @@ -1740,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, @@ -1792,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")); @@ -1813,11 +1955,11 @@ Status LaunchDynamicSparseAttention( 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)) { - const int split_count = GetGroupedDecodeSplitCount(parameters); + 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.kv_num_heads * + 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")); @@ -1827,21 +1969,35 @@ Status LaunchDynamicSparseAttention( float* partial_sum = partial_max + partial_blocks; float* partial_output = partial_sum + partial_blocks; const size_t grouped_shared_bytes = - GetGroupedAttentionSharedBytes(parameters, sizeof(T)); - if (parameters.cache_capacity <= 131072) { + 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, + 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); + parameters.use_smooth_softmax, nullptr, nullptr); } else { GroupedDecodeDynamicSparseAttentionKernel <<(grouped_blocks), kGroupedDecodeThreads, grouped_shared_bytes, stream>>>( @@ -1849,13 +2005,13 @@ Status LaunchDynamicSparseAttention( 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, + 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); + parameters.use_smooth_softmax, nullptr, nullptr); } CUDA_RETURN_IF_ERROR(cudaGetLastError()); @@ -1869,7 +2025,7 @@ Status LaunchDynamicSparseAttention( partial_max, partial_sum, partial_output, data.output, static_cast(query_blocks), split_count, parameters.head_size); } - } else if (grouped_prefill_threads_supported && + } 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 * @@ -1888,7 +2044,11 @@ Status LaunchDynamicSparseAttention( parameters.local_window_size, parameters.scale, local_plus_selected, parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, parameters.use_smooth_softmax); - } else if (UseFusedAttention(parameters, sizeof(T), max_shared_memory_per_block)) { + } 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, @@ -1899,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; @@ -1927,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 9a68c977425b7..11ef53dd60d45 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -293,9 +293,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); @@ -679,9 +680,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 daf1305969249..5acad33e26b1a 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) \ @@ -37,11 +39,12 @@ namespace psai = onnxruntime::contrib::packed_sparse_attention_indexer; .MayInplace(psai::kMergedStatus, psai::kIndexShareStatus) \ .MayInplace(psai::kPastKeyState, psai::kPresentKeyState) \ .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 @@ -92,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"); @@ -161,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 { const Tensor* mode = context->InputCount() > psai::kIndexShareMode ? context->Input(psai::kIndexShareMode) : nullptr; @@ -227,9 +234,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); @@ -350,10 +358,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))); @@ -380,6 +402,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 = @@ -389,7 +415,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, @@ -400,23 +426,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); @@ -587,8 +619,8 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, con static_cast(output_capacity), 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); @@ -673,9 +705,10 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c static_cast(batch), static_cast(output_capacity_)); } -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 2c57dd3256b7e..8f82da6379ab5 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 @@ -59,6 +59,7 @@ 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; @@ -129,13 +130,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[]; @@ -187,6 +188,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 = @@ -244,7 +275,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]); } @@ -336,8 +367,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, @@ -377,7 +408,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); } @@ -392,8 +423,8 @@ __global__ void QsaBlockScoreKernel(const T* present_key_state, const float* que // (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. // WarpHeads instead keeps all heads in one warp to reuse keys without CTA-wide barriers. -template -__global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const float* query_rotated, +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, @@ -477,8 +508,8 @@ __global__ void QsaRaggedPrefillScoreKernel(const T* present_key_state, const fl } } -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, @@ -489,8 +520,7 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* 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 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; @@ -516,24 +546,36 @@ __global__ void QsaScoreTileTopKKernel(const T* present_key_state, const float* return; } const int d = lane * 4; - const int64_t query_base = - (static_cast(token) * kQwenNumHeads + head) * kQwenHeadSize; - const float4 query_value = LoadQsaVector4(query_rotated + query_base + d); + 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 += kHierarchicalTileBlocks) { + for (int slot = block_slot; slot < kHierarchicalTopKTileBlocks; + slot += kHierarchicalTileScoreThreads / kWarpSize) { const int block_index = tile * kHierarchicalTopKTileBlocks + slot; - float score = 0.0f; + 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; - 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); + key_value = LoadQsaVector4(present_key_state + key_base + d); } - if (lane == 0) { - head_scores[slot][head] = fmaxf(score, 0.0f); +#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(); @@ -624,6 +666,63 @@ __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* out } } +__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) { + 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(); + } +} + __global__ void QsaEmitHierarchicalTopKKernel(const uint64_t* tile_keys, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, @@ -1320,21 +1419,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; @@ -1355,7 +1483,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( const size_t value_bytes = static_cast(params.head_size) * sizeof(float); 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); @@ -1386,8 +1514,8 @@ Status LaunchQsaPackedSparseAttentionIndexer( uint64_t* merge_output = merge_input + key_count; const int64_t tile_work = static_cast(params.total_tokens) * std::min(tile_count, kHierarchicalScoreMaxBlocksPerRow); - QsaScoreTileTopKKernel<<(std::min(tile_work, kSaiMaxGridDimX)), - kHierarchicalScoreThreads, 0, stream>>>( + 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 = @@ -1397,12 +1525,22 @@ Status LaunchQsaPackedSparseAttentionIndexer( 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; list_width *= 2) { - 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); + 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); } } @@ -1413,18 +1551,18 @@ Status LaunchQsaPackedSparseAttentionIndexer( static const bool warp_heads = ParseEnvironmentVariableWithDefault("ORT_PACKED_SPARSE_INDEXER_WARP_HEADS", false); if (warp_heads) { - QsaRaggedPrefillScoreKernel<<>>( + QsaRaggedPrefillScoreKernel<<>>( present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, overflow_flags, block_scores, params); } else { - QsaRaggedPrefillScoreKernel<<>>( + 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); } @@ -1525,21 +1663,29 @@ Status LaunchCsaPackedSparseAttentionIndexer( return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(T) \ - template Status LaunchQsaPackedSparseAttentionIndexer( \ +#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 T*, const T*, const int32_t*, \ - const int32_t*, bool, int32_t*, int32_t*, T*, T*, int32_t*, T*, float*, int32_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 9e2f812eb450a..ee0f6eb32cbcb 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 @@ -38,6 +38,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; @@ -57,7 +59,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, @@ -70,20 +72,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 93850dffc6d74..e25c593a149df 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -1158,9 +1158,9 @@ 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, 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}, diff --git a/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h b/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h index f34579543a4cb..0fbfd764eb276 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h @@ -32,9 +32,9 @@ class SparseAttentionIndexerQsaConcatProgram final {"sequence_length", ProgramUniformVariableDataType::Uint32}, {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, {"total_sequence_length", ProgramUniformVariableDataType::Uint32}, - {"head_size", ProgramUniformVariableDataType::Uint32}, - {"key_row_stride", ProgramUniformVariableDataType::Uint32}, - {"key_offset", ProgramUniformVariableDataType::Uint32}); + {"head_size", ProgramUniformVariableDataType::Uint32}, + {"key_row_stride", ProgramUniformVariableDataType::Uint32}, + {"key_offset", ProgramUniformVariableDataType::Uint32}); private: bool has_past_; @@ -52,7 +52,7 @@ class SparseAttentionIndexerQsaSelectProgram final {"sequence_length", ProgramUniformVariableDataType::Uint32}, {"num_heads", ProgramUniformVariableDataType::Uint32}, {"head_size", ProgramUniformVariableDataType::Uint32}, - {"query_row_stride", ProgramUniformVariableDataType::Uint32}, + {"query_row_stride", ProgramUniformVariableDataType::Uint32}, {"rotary_width", ProgramUniformVariableDataType::Uint32}, {"max_rotary_length", ProgramUniformVariableDataType::Uint32}, {"rotary_cache_batch_stride", ProgramUniformVariableDataType::Uint32}, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index e3eb187d2966b..cf92e6f11b939 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -487,6 +487,14 @@ 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()) { + fail_shape_inference("num_heads must be a positive int"); + } + if (kv_num_heads <= 0 || kv_num_heads > std::numeric_limits::max()) { + fail_shape_inference("kv_num_heads must be positive and fit in an int"); + } if (hasInputShape(ctx, 0)) { const auto& query_shape = getInputShape(ctx, 0); @@ -496,9 +504,8 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex } auto output_shape = query_shape; - const int64_t num_heads = getAttribute(ctx, "num_heads", 0); - const int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0); if (ctx.getInputType(2) == nullptr) { + output_shape.mutable_dim(2)->Clear(); if (num_heads > 0 && kv_num_heads > 0 && query_dims[2].has_dim_value()) { const int64_t packed_heads = num_heads + 2 * kv_num_heads; const int64_t packed_hidden_size = query_dims[2].dim_value(); @@ -520,7 +527,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; @@ -528,8 +536,6 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex if (hasInputShape(ctx, 0)) { const auto& query_dims = getInputShape(ctx, 0).dim(); - const int64_t num_heads = getAttribute(ctx, "num_heads", 0); - const int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0); if (query_dims.size() == 3 && num_heads > 0 && kv_num_heads > 0 && query_dims[2].has_dim_value()) { const bool is_packed = ctx.getInputType(2) == nullptr; @@ -542,7 +548,7 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex *present_shape.add_dim() = query_dims[0]; present_shape.add_dim()->set_dim_value(kv_num_heads); const auto* total_length_data = ctx.getInputData(10); - if (total_length_data != nullptr) { + if (ctx.getInputType(input_index) == nullptr && total_length_data != nullptr) { const auto total_lengths = ParseData(total_length_data); if (total_lengths.size() != 1) { fail_shape_inference("total_sequence_length must contain exactly one element"); @@ -587,6 +593,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( @@ -623,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", @@ -646,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). @@ -2806,6 +2839,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'"); @@ -2817,9 +2859,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, @@ -2830,8 +2872,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. @@ -3084,6 +3138,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: @@ -3248,8 +3304,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, " @@ -3303,7 +3360,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.", @@ -3327,9 +3384,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 cbf189252e59b..659bc157e9982 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", + 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", + 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)); + ToTensorData(c.past_value)); } else { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); @@ -229,15 +233,29 @@ 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", + 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", + 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)); + 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 +294,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 +745,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 6baeb9fd32725..61c601eaaadcc 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 @@ -794,7 +794,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, int padding = 0, const ConfigOptions* webgpu_options = nullptr) { @@ -810,7 +810,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; @@ -871,8 +871,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 @@ -895,8 +895,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 @@ -912,7 +912,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()); @@ -927,6 +927,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); } @@ -1164,6 +1168,73 @@ TEST(PackedSparseAttentionIndexerTest, QsaHierarchicalPackedRowsWithLargeCapacit } } +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); diff --git a/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc index 3cf1bb9e2ae93..0561c46571879 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc @@ -386,7 +386,9 @@ void QsaReference(const QsaProblem& problem, std::vector& selected, std } else { pooled[static_cast(d)] += problem.key[(static_cast(b) * problem.sequence_length + - position - problem.past_sequence_length) * head_size + d]; + position - problem.past_sequence_length) * + head_size + + d]; } } } @@ -492,7 +494,6 @@ void QsaReference(const QsaProblem& problem, std::vector& selected, std } } } - } struct CsaProblem { From 036b1db49ee100bc5be9ec073df5c23f1d5430b1 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 18:10:52 +0000 Subject: [PATCH 2/2] Fix lint --- .../packed_sparse_attention_indexer_common.h | 6 +- .../cuda/bert/branchwise_rms_norm_impl.cu | 4 +- .../cuda/bert/causal_conv_with_state_impl.cu | 4 +- .../cuda/bert/dynamic_sparse_attention.cc | 30 +-- .../bert/dynamic_sparse_attention_impl.cu | 2 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 46 ++-- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 18 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 28 +- .../cuda/bert/gated_delta_net_impl.cu | 8 +- .../cuda/bert/gated_delta_net_impl.h | 2 +- .../cuda/bert/rotary_embedding_impl.cu | 9 +- .../packed_sparse_attention_indexer_impl.cu | 14 +- .../core/graph/contrib_ops/bert_defs.cc | 242 +++++++++--------- .../cuda/reduction/reduction_functions.cu | 2 +- .../cuda/reduction/reduction_functions.h | 2 +- .../causal_conv_with_state_op_test.cc | 14 +- .../dynamic_sparse_attention_op_test.cc | 40 +-- .../test/contrib_ops/fused_tensor_ops_test.cc | 15 +- .../contrib_ops/hyper_connection_op_test.cc | 10 +- .../contrib_ops/rotary_embedding_op_test.cc | 10 +- ...untime_test_python_symbolic_shape_infer.py | 3 - 21 files changed, 266 insertions(+), 243 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 fdc636bc9626d..72fbac76026f6 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,9 +71,9 @@ 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 + 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, }; diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index c07197fa3bf14..240b057b827e0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -168,9 +168,9 @@ INSTANTIATE(__nv_bfloat16) #undef INSTANTIATE -#define INSTANTIATE_STANDARD(T) \ +#define INSTANTIATE_STANDARD(T) \ template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ - int, int, int, bool, float, float*, float*, float*, void*, size_t, int); + int, int, int, bool, float, float*, float*, float*, void*, size_t, int); INSTANTIATE_STANDARD(float) INSTANTIATE_STANDARD(half) INSTANTIATE_STANDARD(__nv_bfloat16) diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu index c5284935a5e04..eea53e8bbe272 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu @@ -480,8 +480,8 @@ Status LaunchCausalConvProjectionKernel( return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_PROJECTION(T) \ - template Status LaunchCausalConvProjectionKernel( \ +#define INSTANTIATE_PROJECTION(T) \ + template Status LaunchCausalConvProjectionKernel( \ cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, \ int, int, int, int, int, int, int, int, bool, bool); INSTANTIATE_PROJECTION(float) diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index b11283cff7599..f7a65fbad3663 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -19,21 +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()) \ +#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), \ + .TypeConstraint("S", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ + .MayInplace(3, 1) \ + .MayInplace(4, 2) \ + .InputMemoryType(OrtMemTypeCPUInput, 10), \ DynamicSparseAttention); REGISTER_KERNEL_TYPED(float) @@ -142,7 +142,7 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons "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())), + : past_value->IsDataType())), "DynamicSparseAttention: cache types must match activations or both be INT8."); DynamicSparseAttentionParameters parameters; 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 0e3e8746de44d..ccdfdf76bb7ee 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -956,7 +956,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( 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; + : kGroupedDecodeCandidatesPerSplit; const int candidate_begin = split * kCandidatesPerSplit; const int32_t* row_indices = max_selected == 0 ? selected_indices : selected_indices + static_cast(row) * max_selected; diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index 24d2a153bf450..555c98f1fb16d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -13,12 +13,12 @@ namespace onnxruntime::contrib::cuda { -#define REGISTER_ARGMAX_LAST_AXIS(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ +#define REGISTER_ARGMAX_LAST_AXIS(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ ArgMaxLastAxis); REGISTER_ARGMAX_LAST_AXIS(float) @@ -30,38 +30,38 @@ REGISTER_ARGMAX_LAST_AXIS(BFloat16) ONNX_OPERATOR_KERNEL_EX( EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) - .TypeConstraint("F", DataTypeImpl::GetTensorType()), + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()), EmbeddingRMSNorm); ONNX_OPERATOR_KERNEL_EX( RowSelect, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), RowSelect); ONNX_OPERATOR_KERNEL_EX( SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("T", DataTypeImpl::GetTensorType()), SigmoidGateMul); #if !defined(DISABLE_FLOAT8_TYPES) ONNX_OPERATOR_TYPED_KERNEL_EX( DequantizeLinear, kMSDomain, 1, Float8E4M3FN_BFloat16, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) - .TypeConstraint("TS", DataTypeImpl::GetTensorType()), + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("TS", DataTypeImpl::GetTensorType()), DequantizeLinearFp8Bf16); #endif ONNX_OPERATOR_KERNEL_EX( UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("T", DataTypeImpl::GetTensorType()), UntileHyperConnectionWeights); template @@ -153,12 +153,10 @@ SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) const int64_t silu = info.GetAttrOrDefault("silu_gate", 0); ORT_ENFORCE(silu == 0 || silu == 1, "SigmoidGateMul silu_gate must be zero or one"); ORT_ENFORCE(!(broadcast_gate_ && silu_gate_), "SigmoidGateMul broadcast_gate and silu_gate are exclusive"); - ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && - q_stride_ >= 2 * head_dim_ && - tok_stride_ > 0 && q_offset_ >= 0), + ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2 * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0), "SigmoidGateMul received invalid layout attributes"); ORT_ENFORCE(!(broadcast_gate_ || silu_gate_) || (head_dim_ == 0 && heads_ == 0 && q_stride_ == 0 && - tok_stride_ == 0 && q_offset_ == 0), + tok_stride_ == 0 && q_offset_ == 0), "SigmoidGateMul broadcast_gate/silu_gate cannot be combined with packed layout attributes"); } @@ -184,12 +182,12 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { } else { const int64_t row_width = static_cast(heads_) * head_dim_; ORT_RETURN_IF_NOT(count % row_width == 0, - "SigmoidGateMul attention size must be divisible by heads * head_dim"); + "SigmoidGateMul attention size must be divisible by heads * head_dim"); const int64_t tokens = count / row_width; ORT_RETURN_IF_NOT(projection->Shape().Size() >= - (tokens - 1) * tok_stride_ + q_offset_ + - static_cast(heads_) * q_stride_, - "SigmoidGateMul projection is too small for the configured layout"); + (tokens - 1) * tok_stride_ + q_offset_ + + static_cast(heads_) * q_stride_, + "SigmoidGateMul projection is too small for the configured layout"); } Tensor* output = context->Output(0, attention->Shape()); return LaunchSigmoidGateMul( diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 6db5544942e1f..789edc4a26615 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -169,8 +169,8 @@ __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, const BFloat16 a(gate); // Match CUDA Sigmoid's input-type intermediates, including its negative-input subtraction. const BFloat16 value = a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); sigmoid_bf16 = __float2bfloat16_rn(static_cast(value)); } if constexpr (SiluGate) { @@ -254,13 +254,13 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty switch (input_type) { case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); default: return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "ArgMaxLastAxis input type is not supported"); } @@ -299,11 +299,11 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const int64_t block_count = (count - 1) / kThreads + 1; ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), "SigmoidGateMul requires too many CUDA blocks"); -#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ +#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ SigmoidGateMulKernel<<(block_count), kThreads, 0, stream>>>( \ - reinterpret_cast(attention), \ - reinterpret_cast(projection), \ - reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ + reinterpret_cast(attention), \ + reinterpret_cast(projection), \ + reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ token_stride, q_offset, gate_row_width, reinterpret_cast(residual), sigmoid_in_fp32) if (silu_gate) { if (residual != nullptr) { diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index f346bd3f6027d..380e99746af31 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -22,22 +22,22 @@ using onnxruntime::cuda::CudaKernel; using onnxruntime::cuda::OrtToCudaType; namespace gdn = gated_delta_net; -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - GatedDeltaNet, \ - kMSDomain, \ - 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TG", {DataTypeImpl::GetTensorType(), \ +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + GatedDeltaNet, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TG", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}) \ - .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 10) \ - .MayInplace(6, 1), \ + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 10) \ + .MayInplace(6, 1), \ GatedDeltaNet); REGISTER_KERNEL_TYPED(float) diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index f308442329c5d..048faa9c2f3f5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -745,8 +745,8 @@ Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const Vari template static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& plan, const VariantPack& pack, - float scale, int max_threads_per_block, - size_t max_shared_memory_per_block, cudaStream_t stream) { + float scale, int max_threads_per_block, + size_t max_shared_memory_per_block, cudaStream_t stream) { KernelParams p{}; p.total_tokens = desc.total_tokens; p.uniform_len = desc.batch > 0 ? desc.total_tokens / desc.batch : 0; @@ -800,13 +800,13 @@ static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& // The chunked engine only accepts head_size_qk == head_size_v == 128, so the warp // kernel is always the right shape here. return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p, - max_shared_memory_per_block, stream); + max_shared_memory_per_block, stream); } } (void)max_threads_per_block; return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block, - stream); + stream); } template diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h index 2bd82be0b53bf..3035ca02fa083 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h @@ -40,7 +40,7 @@ struct VariantPack { template __device__ __forceinline__ float LoadGate(const void* data, int64_t index, bool gates_in_input_type) { return gates_in_input_type ? static_cast(static_cast(data)[index]) - : static_cast(data)[index]; + : static_cast(data)[index]; } #endif diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu index 9f162783b3a30..55648cc944249 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu @@ -189,8 +189,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int position_ids_format, const bool interleaved, const int max_threads_per_block, int4 in_strides, int4 out_strides, // strides in bnsh coord - const bool use_fp32_arithmetic -) { + const bool use_fp32_arithmetic) { // Note: Current implementation assumes head_size <= max_threads_per_block // because head_size is currently large for LLaMA-2. For smaller head_size // and num_heads values, we can create a block as `block(num_heads, head_size, 1)` @@ -221,8 +220,10 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu interleaved, in_strides, out_strides); } }; - if (use_fp32_arithmetic) launch(std::true_type{}); - else launch(std::false_type{}); + if (use_fp32_arithmetic) + launch(std::true_type{}); + else + launch(std::false_type{}); return CUDA_CALL(cudaGetLastError()); } 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 8f82da6379ab5..4410df08e47cd 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 @@ -667,7 +667,7 @@ __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* out } __global__ void QsaMergeTwoLevelsTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, - int key_stride, int list_width, int rows) { + 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); @@ -1515,7 +1515,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( const int64_t tile_work = static_cast(params.total_tokens) * std::min(tile_count, kHierarchicalScoreMaxBlocksPerRow); QsaScoreTileTopKKernel<<(std::min(tile_work, kSaiMaxGridDimX)), - kHierarchicalTileScoreThreads, 0, stream>>>( + 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 = @@ -1530,7 +1530,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( 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>>>( + kHierarchicalMergeThreads, 0, stream>>>( merge_input, merge_output, key_stride, list_width, params.total_tokens); list_width *= 4; } else { @@ -1663,10 +1663,10 @@ Status LaunchCsaPackedSparseAttentionIndexer( return CUDA_CALL(cudaGetLastError()); } -#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*, \ +#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) \ diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index cf92e6f11b939..7f490a9a1c881 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -3477,55 +3477,55 @@ that are multiplied to query and key before the inner product of query and key i )DOC"; static OpSchema RotaryEmbeddingSchema() { return OpSchema() - .SetDoc(RotaryEmbedding_ver1_doc) - .Attr("scale", - "Custom scale will be used if specified. Default value is 1.0", - AttributeProto::FLOAT, - OPTIONAL_VALUE) - .Attr("interleaved", - "Indicates whether the input has real and imaginary parts interleaved. " - "Default value is 0 (False), meaning the first half of the input consists of real values " - "and the second half consists of imaginary values.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("rotary_embedding_dim", - "Rotary embedding dimension. Default value is 0.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("num_heads", - "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("is_packed_batching", - "ragged batch inputs or not. Default value is 0", - AttributeProto::INT, - OPTIONAL_VALUE) - .Input(0, - "input", - "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", - "T") - .Input(1, - "position_ids", - "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", - "M") - .Input(2, - "cos_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Input(3, - "sin_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Output(0, - "output", - "tensor with same shape as input.", - "T") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") - .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - }); + .SetDoc(RotaryEmbedding_ver1_doc) + .Attr("scale", + "Custom scale will be used if specified. Default value is 1.0", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("interleaved", + "Indicates whether the input has real and imaginary parts interleaved. " + "Default value is 0 (False), meaning the first half of the input consists of real values " + "and the second half consists of imaginary values.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("rotary_embedding_dim", + "Rotary embedding dimension. Default value is 0.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("num_heads", + "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("is_packed_batching", + "ragged batch inputs or not. Default value is 0", + AttributeProto::INT, + OPTIONAL_VALUE) + .Input(0, + "input", + "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", + "T") + .Input(1, + "position_ids", + "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", + "M") + .Input(2, + "cos_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Input(3, + "sin_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Output(0, + "output", + "tensor with same shape as input.", + "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + }); } ONNX_MS_OPERATOR_SET_SCHEMA(RotaryEmbedding, 1, RotaryEmbeddingSchema()); @@ -4689,15 +4689,18 @@ ONNX_MS_OPERATOR_SET_SCHEMA( CausalConvWithState, 1, OpSchema() .SetDoc(CausalConvWithState_ver1_doc) - .Attr("input_channels", "CUDA channels-last input channel slice width; 0 selects all channels. " + .Attr("input_channels", + "CUDA channels-last input channel slice width; 0 selects all channels. " "Nondefault extended attributes require rank-3 input and state_window=0.", AttributeProto::INT, static_cast(0)) .Attr("input_channel_offset", "Starting channel of the channels-last input slice (CUDA only).", AttributeProto::INT, static_cast(0)) - .Attr("state_length", "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. " + .Attr("state_length", + "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. " "An explicit length must cover the receptive field; convolution reads its trailing history.", AttributeProto::INT, static_cast(0)) - .Attr("state_channels_last", "State layout override (CUDA only): -1 follows channels_last, " + .Attr("state_channels_last", + "State layout override (CUDA only): -1 follows channels_last, " "0 uses [B,C,state_length], 1 uses [B,state_length,C].", AttributeProto::INT, static_cast(-1)) .Attr("activation", @@ -4781,9 +4784,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(seq_len - W + j). Extended attributes override the state length and layout " "as for past_state.", "T") - .Output(2, "first_state", "Optional carry state after the first input position, with the same " + .Output(2, "first_state", + "Optional carry state after the first input position, with the same " "layout and length as present_state. CUDA only; requires rank-3 channels-last input " - "and state_window=0.", "T", OpSchema::Optional) + "and state_window=0.", + "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") @@ -5763,20 +5768,26 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("q_stride", "Per-head stride in the packed projection.", AttributeProto::INT, static_cast(0)) .Attr("tok_stride", "Per-token stride in the packed projection.", AttributeProto::INT, static_cast(0)) .Attr("q_offset", "Offset of the query projection within each token.", AttributeProto::INT, static_cast(0)) - .Attr("broadcast_gate", "When one, projection has the attention shape with final dimension one. " - "Packed layout attributes must be zero. CUDA only.", AttributeProto::INT, static_cast(0)) - .Attr("sigmoid_in_fp32", "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " + .Attr("broadcast_gate", + "When one, projection has the attention shape with final dimension one. " + "Packed layout attributes must be zero. CUDA only.", + AttributeProto::INT, static_cast(0)) + .Attr("sigmoid_in_fp32", + "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " "the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.", AttributeProto::INT, static_cast(1)) - .Attr("silu_gate", "When one, projection and attention have identical shapes. Multiply the " + .Attr("silu_gate", + "When one, projection and attention have identical shapes. Multiply the " "BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. " "Cannot be combined with broadcast_gate or packed layout attributes.", AttributeProto::INT, static_cast(0)) .Input(0, "attention", "BF16 attention values.", "T") .Input(1, "projection", "BF16 packed projection containing gate values.", "T") .Input(2, "residual", "Optional BF16 residual with the attention shape.", "T", OpSchema::Optional) - .Output(0, "output", "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " - "before optionally adding residual and rounding again.", "T") + .Output(0, "output", + "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " + "before optionally adding residual and rounding again.", + "T") .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); @@ -5947,64 +5958,66 @@ WebGPU currently supports only the default mode without Y_fp32. static OpSchema BranchwiseRMSNormSchema() { return OpSchema() - .SetDoc(BranchwiseRMSNorm_ver1_doc) - .Attr("epsilon", "Epsilon added before reciprocal square root.", - AttributeProto::FLOAT, 1e-5f) - .Attr("num_branches", - "Number of branches for flattened input. Omit or set to zero for grouped input.", - AttributeProto::INT, static_cast(0)) - .Attr("use_onnx_arithmetic", - "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " - "Supported values are 0 (default) and 1.", - AttributeProto::INT, static_cast(0)) - .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") - .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", - "M", OpSchema::Optional) - .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") - .Output(1, "Y_fp32", "Optional float32 widening of the already T-rounded Y. " - "Requires use_onnx_arithmetic=1.", "F", OpSchema::Optional) - .TypeConstraint("T", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain input and output to floating-point tensors.") - .TypeConstraint("M", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain scale to floating-point tensors.") - .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - const auto start = getAttribute(ctx, "input_start", 0); - const auto width = getAttribute(ctx, "input_width", 0); - if (start < 0 || width < 0 || (width == 0 && start != 0)) { - fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + .SetDoc(BranchwiseRMSNorm_ver1_doc) + .Attr("epsilon", "Epsilon added before reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("num_branches", + "Number of branches for flattened input. Omit or set to zero for grouped input.", + AttributeProto::INT, static_cast(0)) + .Attr("use_onnx_arithmetic", + "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " + "Supported values are 0 (default) and 1.", + AttributeProto::INT, static_cast(0)) + .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") + .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", + "M", OpSchema::Optional) + .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") + .Output(1, "Y_fp32", + "Optional float32 widening of the already T-rounded Y. " + "Requires use_onnx_arithmetic=1.", + "F", OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain input and output to floating-point tensors.") + .TypeConstraint("M", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain scale to floating-point tensors.") + .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const auto start = getAttribute(ctx, "input_start", 0); + const auto width = getAttribute(ctx, "input_width", 0); + if (start < 0 || width < 0 || (width == 0 && start != 0)) { + fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + } + if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || + getAttribute(ctx, "num_branches", 0) != 0)) { + fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + } + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + if (width != 0) { + if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); + auto* last = shape.mutable_dim(shape.dim_size() - 1); + if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { + fail_shape_inference("input slice exceeds the last dimension"); + } + last->set_dim_value(width); } - if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || - getAttribute(ctx, "num_branches", 0) != 0)) { - fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + updateOutputShape(ctx, 0, shape); + } + if (ctx.getNumOutputs() > 1) { + if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { + fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); } + updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); if (hasInputShape(ctx, 0)) { auto shape = getInputShape(ctx, 0); - if (width != 0) { - if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); - auto* last = shape.mutable_dim(shape.dim_size() - 1); - if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { - fail_shape_inference("input slice exceeds the last dimension"); - } - last->set_dim_value(width); - } - updateOutputShape(ctx, 0, shape); + if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); + updateOutputShape(ctx, 1, shape); } - if (ctx.getNumOutputs() > 1) { - if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { - fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); - } - updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - if (hasInputShape(ctx, 0)) { - auto shape = getInputShape(ctx, 0); - if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); - updateOutputShape(ctx, 1, shape); - } - } - }); + } + }); } ONNX_MS_OPERATOR_SET_SCHEMA(BranchwiseRMSNorm, 1, BranchwiseRMSNormSchema()); @@ -6013,7 +6026,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( BranchwiseRMSNormSchema() .Attr("input_start", "CUDA-only starting offset in the last input dimension.", AttributeProto::INT, static_cast(0)) - .Attr("input_width", "CUDA-only width selected from the last input dimension before normalization. " + .Attr("input_width", + "CUDA-only width selected from the last input dimension before normalization. " "Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. " "Outputs retain the leading dimensions and use this last dimension.", AttributeProto::INT, static_cast(0))); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index bf960242503d4..75024a420d3a9 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -539,7 +539,7 @@ struct RoundedSquare { template Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { + int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { return detail::call_reduce_matrix_columns( stream, input, output, m, n, buffer, buffer_size, input_row_stride); } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index e5bb1b03820f0..9c877c0743b31 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -109,7 +109,7 @@ Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output /** Sum FP32-rounded squares across columns, preserving reduce_matrix_columns' summation order. */ template Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); + int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); /** * Computes ReduceSum for an arbitrary set of axes. This is the general fallback diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index 8e015e80d5b32..7cac5d1d53c16 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -365,12 +365,16 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u test.AddAttribute("activation", silu ? "silu" : "none"); test.AddInput("input", {batch, length, width}, input); test.AddInput("weight", {channels, 1, kernel}, weight); - if (use_bias) test.AddInput("bias", {channels}, bias); - else test.AddOptionalInputEdge(); + if (use_bias) + test.AddInput("bias", {channels}, bias); + else + test.AddOptionalInputEdge(); const std::vector state_shape = state_last ? std::vector{batch, history, channels} - : std::vector{batch, channels, history}; - if (use_past) test.AddInput("past", state_shape, past); - else test.AddOptionalInputEdge(); + : std::vector{batch, channels, history}; + if (use_past) + test.AddInput("past", state_shape, past); + else + test.AddOptionalInputEdge(); test.AddOutput("output", {batch, length, channels}, output); test.AddOutput("present", state_shape, present); if (first_output) test.AddOutput("first", state_shape, first); 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 659bc157e9982..bc648f4f5ba37 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -168,11 +168,11 @@ void RunDynamicSparseAttentionCase( } 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)); + {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)); + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.past_value)); } else { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); @@ -239,20 +239,24 @@ void RunDynamicSparseAttentionCase( 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); + 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)); + {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)); + {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}); } @@ -420,11 +424,11 @@ TEST(DynamicSparseAttentionTest, Int8CacheRejectsMissingOrMalformedScales_CUDA) } 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"; + 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); + OpTester::ExpectResult::kExpectFailure, error); } } @@ -517,14 +521,14 @@ DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool au 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]; + : 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]; + : c.expected_present_value[offset + d] * c.value_scale[kv * c.head_size + d]; numerator[d] += static_cast(weight * value); } } diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc index 441a7bf62ccce..2173434746ce9 100644 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -146,8 +146,8 @@ void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_ const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat()); const BFloat16 reciprocal(1.0f / denominator.ToFloat()); const BFloat16 rounded_sigmoid = sigmoid_in_fp32 - ? BFloat16(sigmoid) - : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); + ? BFloat16(sigmoid) + : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); const BFloat16 activation = silu_gate ? BFloat16(gate * rounded_sigmoid.ToFloat()) : rounded_sigmoid; const BFloat16 product(values[i].ToFloat() * activation.ToFloat()); expected[i] = residual ? BFloat16(product.ToFloat() + residuals[i].ToFloat()) : product; @@ -186,7 +186,7 @@ TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) { test.AddAttribute("broadcast_gate", 1); test.AddInput("attention", {2, 3}, ToBFloat16(std::vector(6, 1.0f))); test.AddInput("projection", bad_residual ? std::vector{2, 1} : std::vector{1, 2}, - ToBFloat16({1.0f, 1.0f})); + ToBFloat16({1.0f, 1.0f})); if (bad_residual) test.AddInput("residual", {1}, ToBFloat16({1.0f})); test.AddOutput("output", {2, 3}, ToBFloat16(std::vector(6, 0.0f))); test.Run(OpTester::ExpectResult::kExpectFailure, @@ -208,8 +208,8 @@ TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { test.AddInput("residual", {0, 3}, {}); test.AddOutput("output", {0, 3}, {}); const std::vector errors = {"", "broadcast_gate must be zero or one", - "sigmoid_in_fp32 must be zero or one", - "cannot be combined with packed layout attributes"}; + "sigmoid_in_fp32 must be zero or one", + "cannot be combined with packed layout attributes"}; test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, errors[mode], {}, nullptr, &providers); } @@ -227,14 +227,13 @@ TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { test.AddInput("projection", {0, mode == 4 ? 1 : 3}, {}); test.AddOutput("output", {0, 3}, {}); const std::vector errors = {"", "silu_gate must be zero or one", "are exclusive", - "cannot be combined with packed layout attributes", - "silu gate must have the same shape"}; + "cannot be combined with packed layout attributes", + "silu gate must have the same shape"}; test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, errors[mode], {}, nullptr, &providers); } } - TEST(FusedTensorOpsTest, UntileHyperConnectionWeights) { constexpr int64_t kInputElements = 160ll * 20480; constexpr int64_t kVectorCount = kInputElements / 16; diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index 7942e8fc22425..d7e1bc7acdd72 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc @@ -113,7 +113,10 @@ class StandardRMSChainTester : public OpTester { public: StandardRMSChainTester(int output_type, float epsilon, int64_t start = 0, int64_t width = 0) : OpTester("StandardRMSChain", 17, kOnnxDomain, false), - output_type_(output_type), epsilon_(epsilon), start_(start), width_(width) {} + output_type_(output_type), + epsilon_(epsilon), + start_(start), + width_(width) {} void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, std::vector>&) override { @@ -132,7 +135,8 @@ class StandardRMSChainTester : public OpTester { source = value("selected"); graph.AddNode("slice", "Slice", "", {inputs[0], constant("starts", start_), constant("ends", start_ + width_), - constant("axes", -1)}, {source}); + constant("axes", -1)}, + {source}); } auto* x = value("float_x"); graph.AddNode("cast", "Cast", "", {source}, {x}) @@ -200,7 +204,7 @@ void TestStandardRMSArithmetic(int type, bool packed = false) { std::vector expected; for (OpTester* tester : {static_cast(&reference), &fused}) { const std::vector input_shape = rows == 48 ? std::vector{1, 2, 24, stride} - : std::vector{1, rows, stride}; + : std::vector{1, rows, stride}; auto output_shape = input_shape; output_shape.back() = width; tester->AddInput("X", input_shape, input); diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc index 9ebbc10c97bf2..a6957e18bc194 100644 --- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc +++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc @@ -27,7 +27,10 @@ class FloatRotaryChainTester : public OpTester { public: FloatRotaryChainTester(int type, int heads, int rotary_dim, bool interleaved) : OpTester("RotaryEmbedding", 1, kMSDomain, false), - type_(type), heads_(heads), rotary_dim_(rotary_dim), interleaved_(interleaved) {} + type_(type), + heads_(heads), + rotary_dim_(rotary_dim), + interleaved_(interleaved) {} void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, std::vector>&) override { @@ -74,7 +77,7 @@ void TestFloatRotaryArithmetic(int type) { const int rotary_dim = 64; const int cache_rows = batch * sequence + 3; std::vector shape = transposed ? std::vector{batch, heads, sequence, width} - : std::vector{batch, sequence, heads * width}; + : std::vector{batch, sequence, heads * width}; std::vector input(batch * sequence * heads * width), cosine(cache_rows * rotary_dim / 2), sine(cosine.size()); const bool scalar_position = !transposed; @@ -94,8 +97,7 @@ void TestFloatRotaryArithmetic(int type) { std::vector expected; for (OpTester* tester : {static_cast(&reference), &fused}) { tester->AddInput("input", shape, input); - tester->AddInput("positions", scalar_position ? std::vector{1} - : std::vector{batch, sequence}, positions); + tester->AddInput("positions", scalar_position ? std::vector{1} : std::vector{batch, sequence}, positions); tester->AddInput("cosine", {cache_rows, rotary_dim / 2}, cosine); tester->AddInput("sine", {cache_rows, rotary_dim / 2}, sine); tester->AddOutput("output", shape, std::vector(input.size())); diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 96bb35672a6cf..6ac581d72cab8 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -882,9 +882,6 @@ def infer(query_width, packed): self.assertEqual(separate_graph.output[0].type.tensor_type.shape.dim[2].dim_value, 33) self.assertFalse(separate_graph.output[1].type.tensor_type.HasField("shape")) - - - def test_unsqueeze_opset_11(self): graph = helper.make_graph( [