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..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,7 +71,11 @@ enum OutputIndex : int { kPresentStateLengths = 5, // [batch_size, 2], int32 kStateUpdate = 6, // qsa only, optional: [batch_size, state_update_capacity, head_size] kIndexShareStatus = 7, + kMidKvBuffer = 8, // qsa only, optional: buffer snapshot after the first request token + kMidStateLengths = 9, // qsa only, optional: state lengths after the first request token + kValidationStatus = 10, // qsa only, optional: per-request overflow/metadata status kOutputCount = 8, + kMaxOutputCount = 11, }; // Every PackedSparseAttentionIndexer node declares all 6 fixed outputs; present_gate_buffer is an diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index 878659302a19a..f7a65fbad3663 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -19,19 +19,21 @@ namespace cuda { using namespace onnxruntime::cuda; -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - DynamicSparseAttention, \ - kMSDomain, \ - 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ - .MayInplace(3, 1) \ - .MayInplace(4, 2) \ - .InputMemoryType(OrtMemTypeCPUInput, 10), \ +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + DynamicSparseAttention, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TCache", BuildKernelDefConstraints()) \ + .TypeConstraint("S", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ + .MayInplace(3, 1) \ + .MayInplace(4, 2) \ + .InputMemoryType(OrtMemTypeCPUInput, 10), \ DynamicSparseAttention); REGISTER_KERNEL_TYPED(float) @@ -129,6 +131,19 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons const Tensor* q_norm_weight = context->Input(14); const Tensor* k_norm_weight = context->Input(15); const Tensor* head_sink = context->Input(16); + const Tensor* committed_length = + context->InputCount() > 17 ? context->Input(17) : nullptr; + const Tensor* key_scale = context->InputCount() > 18 ? context->Input(18) : nullptr; + const Tensor* value_scale = context->InputCount() > 19 ? context->Input(19) : nullptr; + const bool quantized_cache = past_key != nullptr && past_key->IsDataType(); + ORT_RETURN_IF_NOT((key_scale != nullptr) == (value_scale != nullptr), + "DynamicSparseAttention: key_scale and value_scale must be provided together."); + ORT_RETURN_IF_NOT(quantized_cache == (key_scale != nullptr), + "DynamicSparseAttention: INT8 caches require per-channel key_scale and value_scale."); + ORT_RETURN_IF_NOT((past_key == nullptr || past_key->IsDataType() || quantized_cache) && + (past_value == nullptr || (quantized_cache ? past_value->IsDataType() + : past_value->IsDataType())), + "DynamicSparseAttention: cache types must match activations or both be INT8."); DynamicSparseAttentionParameters parameters; ORT_RETURN_IF_ERROR(dynamic_sparse_attention_helper::CheckInputs( @@ -140,9 +155,18 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons qk_norm_epsilon_, parameters)); parameters.rotary_interleaved = rotary_interleaved_; parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr; + if (quantized_cache) { + ORT_RETURN_IF_NOT(past_value != nullptr && + key_scale->Shape() == TensorShape({parameters.kv_num_heads * parameters.head_size}) && + value_scale->Shape() == key_scale->Shape(), + "DynamicSparseAttention: INT8 cache scales must have shape [kv_num_heads * head_size]."); + } ORT_RETURN_IF_NOT(parameters.head_size <= GetDeviceProp().maxThreadsPerBlock, "DynamicSparseAttention: head_size exceeds the CUDA device thread-block limit."); + ORT_RETURN_IF(committed_length != nullptr && + (parameters.batch_size != 1 || committed_length->Shape().Size() != 1), + "DynamicSparseAttention: committed_length requires batch_size=1 and one int32 element."); TensorShape output_shape({parameters.batch_size, parameters.sequence_length, parameters.query_hidden_size}); @@ -153,6 +177,8 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons parameters.cache_capacity, parameters.head_size}); Tensor* present_key_output = context->OutputCount() > 1 ? context->Output(1, present_shape) : nullptr; Tensor* present_value_output = context->OutputCount() > 2 ? context->Output(2, present_shape) : nullptr; + Tensor* validation_status_output = + context->OutputCount() > 3 ? context->Output(3, TensorShape({1})) : nullptr; const size_t query_elements = SafeInt(parameters.batch_size) * parameters.sequence_length * parameters.query_hidden_size; auto prepared_query = GetScratchBuffer(query_elements, GetComputeStream(context)); @@ -160,16 +186,17 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons const size_t cache_elements = SafeInt(parameters.batch_size) * parameters.kv_num_heads * parameters.cache_capacity * parameters.head_size; - IAllocatorUniquePtr key_cache_scratch; - IAllocatorUniquePtr value_cache_scratch; - T* key_cache = present_key_output == nullptr ? nullptr : present_key_output->MutableData(); - T* value_cache = present_value_output == nullptr ? nullptr : present_value_output->MutableData(); + const size_t cache_bytes = SafeInt(cache_elements) * (quantized_cache ? 1 : sizeof(T)); + IAllocatorUniquePtr key_cache_scratch; + IAllocatorUniquePtr value_cache_scratch; + void* key_cache = present_key_output == nullptr ? nullptr : present_key_output->MutableDataRaw(); + void* value_cache = present_value_output == nullptr ? nullptr : present_value_output->MutableDataRaw(); if (key_cache == nullptr) { - key_cache_scratch = GetScratchBuffer(cache_elements, GetComputeStream(context)); + key_cache_scratch = GetScratchBuffer(cache_bytes, GetComputeStream(context)); key_cache = key_cache_scratch.get(); } if (value_cache == nullptr) { - value_cache_scratch = GetScratchBuffer(cache_elements, GetComputeStream(context)); + value_cache_scratch = GetScratchBuffer(cache_bytes, GetComputeStream(context)); value_cache = value_cache_scratch.get(); } @@ -178,8 +205,11 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons data.query = reinterpret_cast(query->Data()); data.key = key == nullptr ? nullptr : reinterpret_cast(key->Data()); data.value = value == nullptr ? nullptr : reinterpret_cast(value->Data()); - data.past_key = past_key == nullptr ? nullptr : reinterpret_cast(past_key->Data()); - data.past_value = past_value == nullptr ? nullptr : reinterpret_cast(past_value->Data()); + data.past_key = past_key == nullptr ? nullptr : reinterpret_cast(past_key->DataRaw()); + data.past_value = past_value == nullptr ? nullptr : reinterpret_cast(past_value->DataRaw()); + data.quantized_cache = quantized_cache; + data.key_scale = key_scale == nullptr ? nullptr : key_scale->Data(); + data.value_scale = value_scale == nullptr ? nullptr : value_scale->Data(); data.auxiliary_key = auxiliary_key == nullptr ? nullptr : reinterpret_cast(auxiliary_key->Data()); data.auxiliary_value = auxiliary_value == nullptr @@ -207,8 +237,17 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(stream, &capture_status)); ORT_RETURN_IF_NOT(capture_status == cudaStreamCaptureStatusNone, "DynamicSparseAttention strict validation is not supported during CUDA graph capture."); + } - auto validation_error = GetScratchBuffer(1, GetComputeStream(context)); + IAllocatorUniquePtr validation_error_scratch; + int32_t* validation_error = nullptr; + if (validation_status_output != nullptr) { + validation_error = validation_status_output->MutableData(); + } else if (strict_validation_) { + validation_error_scratch = GetScratchBuffer(1, GetComputeStream(context)); + validation_error = validation_error_scratch.get(); + } + if (validation_error != nullptr) { const size_t validation_workspace_elements = GetDynamicSparseAttentionValidationWorkspaceSize( parameters, GetDeviceProp().sharedMemPerBlock); @@ -216,13 +255,18 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons GetScratchBuffer(validation_workspace_elements, GetComputeStream(context)); ORT_RETURN_IF_ERROR(ValidateDynamicSparseAttentionOnDevice( stream, data.selected_indices, data.selected_counts, data.seqlens_k, - data.position_ids, parameters, validation_error.get(), validation_workspace.get(), - GetDeviceProp().sharedMemPerBlock, true)); + data.position_ids, parameters, validation_error, validation_workspace.get(), + GetDeviceProp().sharedMemPerBlock, strict_validation_, data.key_scale, data.value_scale)); + if (committed_length != nullptr) { + ORT_RETURN_IF_ERROR(ValidateDynamicSparseAttentionCursor( + stream, data.seqlens_k, committed_length->Data(), + parameters.sequence_length, parameters.cache_capacity, validation_error)); + } } const size_t attention_workspace_elements = GetDynamicSparseAttentionWorkspaceSize( - parameters, sizeof(CudaT), GetDeviceProp().sharedMemPerBlock); + parameters, sizeof(CudaT), GetDeviceProp().sharedMemPerBlock, quantized_cache); auto attention_workspace = GetScratchBuffer(attention_workspace_elements, GetComputeStream(context)); data.attention_workspace = attention_workspace.get(); diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index 1b5a2b0404f53..ccdfdf76bb7ee 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -6,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 8e9c6fde2ad20..eaedf62a4ffca 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.cc @@ -7,6 +7,7 @@ #include #include #include +#include #include "contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h" #include "core/providers/cuda/cuda_common.h" @@ -19,15 +20,16 @@ namespace cuda { using namespace onnxruntime::cuda; namespace psai = onnxruntime::contrib::packed_sparse_attention_indexer; -#define REGISTER_KERNEL_TYPED(T) \ +#define REGISTER_KERNEL_TYPED(K, T, TState) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ PackedSparseAttentionIndexer, \ kMSDomain, \ 1, \ - T, \ + K, \ kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_KEY", DataTypeImpl::GetTensorType()) \ .TypeConstraint("I", DataTypeImpl::GetTensorType()) \ .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ .InputMemoryType(OrtMemTypeCPUInput, psai::kStateUpdateActive) \ @@ -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 { bool reuse = false; const Tensor* mode = context->InputCount() > psai::kIndexShareMode ? context->Input(psai::kIndexShareMode) @@ -231,9 +238,10 @@ Status PackedSparseAttentionIndexer::ComputeInternal(OpKernelContext* context return is_qsa ? ComputeQsa(context) : ComputeCsa(context); } -template -Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) const { +template +Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) const { using CudaT = typename OrtToCudaType::type; + using CudaStateT = typename OrtToCudaType::type; const Tensor* query = context->Input(psai::kQuery); const Tensor* key = context->Input(psai::kKey); @@ -354,10 +362,24 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con ? context->Output( psai::kStateUpdate, TensorShape({batch_size, state_update_capacity_, head_size})) : nullptr; + Tensor* mid_kv_buffer = + context->OutputCount() > psai::kMidKvBuffer + ? context->Output(psai::kMidKvBuffer, TensorShape({batch_size, buffer_capacity, head_size})) + : nullptr; + Tensor* mid_state_lengths = + context->OutputCount() > psai::kMidStateLengths + ? context->Output(psai::kMidStateLengths, TensorShape({batch_size, psai::kStateLengthColumns})) + : nullptr; + Tensor* validation_status = + context->OutputCount() > psai::kValidationStatus + ? context->Output(psai::kValidationStatus, TensorShape({batch_size})) + : nullptr; ORT_RETURN_IF(selected_indices == nullptr || selected_counts == nullptr || present_key_state == nullptr || present_kv_buffer == nullptr || present_state_lengths == nullptr, "PackedSparseAttentionIndexer: policy_mode 'qsa' requires selected_indices, selected_counts, " "present_key_state, present_kv_buffer and present_state_lengths outputs"); + ORT_RETURN_IF((mid_kv_buffer == nullptr) != (mid_state_lengths == nullptr), + "PackedSparseAttentionIndexer: mid_kv_buffer and mid_state_lengths must be connected together"); if (state_update != nullptr && state_update->SizeInBytes() > 0) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(state_update->MutableDataRaw(), 0, state_update->SizeInBytes(), Stream(context))); @@ -383,6 +405,10 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con params.scale = has_scale_ ? scale_ : 1.0f / std::sqrt(static_cast(head_size)); params.block_topk = static_cast(token_budget_ / compress_ratio_); params.state_update_capacity = static_cast(state_update_capacity_); + params.mid_kv_buffer = + mid_kv_buffer == nullptr ? nullptr : reinterpret_cast(mid_kv_buffer->MutableData()); + params.mid_state_lengths = + mid_state_lengths == nullptr ? nullptr : mid_state_lengths->MutableData(); auto float_workspace = GetScratchBuffer(GetQsaPackedWorkspaceFloatCount(params), GetComputeStream(context)); auto overflow_flags = @@ -392,7 +418,7 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con const CudaT* key_data = has_separate_key ? reinterpret_cast(key->Data()) : query_data + query_width; - return LaunchQsaPackedSparseAttentionIndexer( + ORT_RETURN_IF_ERROR((LaunchQsaPackedSparseAttentionIndexer( Stream(context), params, query_data, key_data, @@ -403,23 +429,29 @@ Status PackedSparseAttentionIndexer::ComputeQsa(OpKernelContext* context) con cumulative_sequence_lengths->Data(), past_sequence_lengths->Data(), position_ids != nullptr ? position_ids->Data() : nullptr, - reinterpret_cast(past_key_state->Data()), + reinterpret_cast(past_key_state->Data()), reinterpret_cast(past_kv_buffer->Data()), past_state_lengths->Data(), state_update_capture_count != nullptr ? state_update_capture_count->Data() : nullptr, state_update_active == nullptr || state_update_active->Data()[0] != 0, selected_indices->MutableData(), selected_counts->MutableData(), - reinterpret_cast(present_key_state->MutableData()), + reinterpret_cast(present_key_state->MutableData()), reinterpret_cast(present_kv_buffer->MutableData()), present_state_lengths->MutableData(), state_update != nullptr ? reinterpret_cast(state_update->MutableData()) : nullptr, float_workspace.get(), - overflow_flags.get()); + overflow_flags.get()))); + if (validation_status != nullptr) { + return LaunchMapQsaValidationStatus( + Stream(context), cumulative_sequence_lengths->Data(), overflow_flags.get(), + static_cast(batch_size), validation_status->MutableData()); + } + return Status::OK(); } -template -Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, const bool state_only) const { +template +Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, const bool state_only) const { using CudaT = typename OrtToCudaType::type; const Tensor* query = context->Input(psai::kQuery); @@ -612,8 +644,8 @@ Status PackedSparseAttentionIndexer::ComputeCsa(OpKernelContext* context, con static_cast(compress_ratio_), false, overflow_flags.get()); } -template -Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { +template +Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) const { const Tensor* query = context->Input(psai::kQuery); const Tensor* lengths = context->Input(psai::kPastSequenceLengths); const Tensor* indices = context->Input(psai::kMergedIndices); @@ -715,9 +747,10 @@ Status PackedSparseAttentionIndexer::ComputeReuse(OpKernelContext* context) c static_cast(state_capacity_), static_cast(compress_ratio_), is_qsa); } -template class PackedSparseAttentionIndexer; -template class PackedSparseAttentionIndexer; -template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; +template class PackedSparseAttentionIndexer; } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h index 9330edc01b0a8..1ce56f0fd229f 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer.h @@ -11,7 +11,7 @@ namespace onnxruntime { namespace contrib { namespace cuda { -template +template class PackedSparseAttentionIndexer final : public onnxruntime::cuda::CudaKernel { public: explicit PackedSparseAttentionIndexer(const OpKernelInfo& info); diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 6efc9be09c9c5..b61b2d8c80b47 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -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; @@ -139,13 +140,13 @@ __global__ void ElementwiseCopyKernel(const T* src, T* dst, int64_t count) { // leading RoPE -> append into fixed-capacity key_state) and publishes the raw trailing buffer. // Reads only past_kv_buffer / key (never present_kv_buffer / present_key_state), so it is correct // whether or not the present/past tensors are the same aliased allocation. -template +template __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, const T* cos_cache, const T* sin_cache, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const T* past_kv_buffer, const int32_t* past_state_lengths, const int32_t* state_update_capture_count, bool state_update_active, - T* present_key_state, T* present_kv_buffer, + TState* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths, T* state_update, int32_t* overflow_flags, PackedSparseAttentionIndexerParams params) { extern __shared__ float shared[]; @@ -197,6 +198,36 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con overflow_flags[b] = rejected ? 1 : 0; } + if (params.mid_kv_buffer != nullptr && params.mid_state_lengths != nullptr) { + const int mid_req = min(req_len, 1); + const int mid_pending = old_buf_len + mid_req; + const int mid_blocks = rejected ? 0 : mid_pending / params.compress_ratio; + const int mid_buf_len = rejected ? old_buf_len : mid_pending % params.compress_ratio; + if (threadIdx.x == 0) { + params.mid_state_lengths[b * 2 + psai::kKeyStateLength] = old_key_len + mid_blocks; + params.mid_state_lengths[b * 2 + psai::kBufferLength] = mid_buf_len; + } + T* mid_buffer = static_cast(params.mid_kv_buffer); + const int mid_total = params.buffer_capacity * params.head_size; + for (int i = static_cast(threadIdx.x); i < mid_total; i += static_cast(blockDim.x)) { + const int t = i / params.head_size; + const int d = i - t * params.head_size; + const int virtual_pos = mid_blocks * params.compress_ratio + t; + const int64_t slot = static_cast(b) * params.buffer_capacity * params.head_size + i; + if (t >= mid_buf_len || rejected) { + mid_buffer[slot] = past_kv_buffer[slot]; + } else { + mid_buffer[slot] = + virtual_pos < old_buf_len + ? past_kv_buffer[(static_cast(b) * params.buffer_capacity + virtual_pos) * + params.head_size + + d] + : key[(static_cast(req_start) + (virtual_pos - old_buf_len)) * params.key_row_stride + + d]; + } + } + } + for (int t = static_cast(threadIdx.x); t < capture_count; t += static_cast(blockDim.x)) { if ((old_buf_len + t + 1) % params.compress_ratio != 0) { const int64_t capture_base = @@ -254,7 +285,7 @@ __global__ void QsaUpdateStateKernel(const T* key, const T* key_norm_weight, con const int64_t capture_base = (static_cast(b) * params.state_update_capacity + completion_token) * params.head_size; for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { - present_key_state[out_base + d] = from_float(rotated[d]); + present_key_state[out_base + d] = from_float(rotated[d]); if (completion_token < capture_count) { state_update[capture_base + d] = from_float(rotated[d]); } @@ -346,8 +377,8 @@ __global__ void PackedRotateQueryKernel(const T* query, const T* query_norm_weig // One block per (token, key_state slot). Scores are directly dotted against the already-prepared // present_key_state entry (unlike the dense op, no per-query pooling/normalize/rotate is repeated // here because the packed contract stores fully-prepared blocks in key_state). -template -__global__ void QsaBlockScoreKernel(const T* present_key_state, const float* query_rotated, +template +__global__ void QsaBlockScoreKernel(const TState* present_key_state, const float* query_rotated, const int32_t* cumulative_sequence_lengths, const int32_t* past_sequence_lengths, const int64_t* position_ids, const int32_t* present_state_lengths, float* block_scores, @@ -387,7 +418,7 @@ __global__ void QsaBlockScoreKernel(const T* present_key_state, const float* que query_rotated + (static_cast(token) * params.num_heads + head) * params.head_size; float partial = 0.0f; for (int d = static_cast(threadIdx.x); d < params.head_size; d += static_cast(blockDim.x)) { - partial += query_head[d] * to_float(present_key_state[key_base + d]); + partial += query_head[d] * to_float(present_key_state[key_base + d]); } score += fmaxf(SaiBlockSum(partial, reduction), 0.0f); } @@ -402,8 +433,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, @@ -487,8 +518,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, @@ -499,8 +530,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; @@ -526,24 +556,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(); @@ -777,6 +819,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, @@ -1495,21 +1594,50 @@ size_t GetCsaPackedWorkspaceFloatCount(const PackedSparseAttentionIndexerParams& return rows * params.num_heads * params.head_size + rows * static_cast(std::max(params.state_capacity, 1)); } -template +__global__ void MapQsaValidationStatusKernel( + const int32_t* cumulative_sequence_lengths, + const int32_t* overflow_flags, + int batch_size, + int32_t* validation_status) { + const int batch = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (batch >= batch_size) { + return; + } + const int request_length = cumulative_sequence_lengths[batch + 1] - cumulative_sequence_lengths[batch]; + validation_status[batch] = overflow_flags[batch] == 0 ? 0 : (request_length >= 2 ? 0x99 : 9); +} + +Status LaunchMapQsaValidationStatus( + cudaStream_t stream, + const int32_t* cumulative_sequence_lengths, + const int32_t* overflow_flags, + int batch_size, + int32_t* validation_status) { + if (batch_size <= 0) { + return Status::OK(); + } + constexpr int kStatusThreads = 128; + const int blocks = (batch_size + kStatusThreads - 1) / kStatusThreads; + MapQsaValidationStatusKernel<<>>( + cumulative_sequence_lengths, overflow_flags, batch_size, validation_status); + return CUDA_CALL(cudaGetLastError()); +} + +template Status LaunchQsaPackedSparseAttentionIndexer( cudaStream_t stream, const PackedSparseAttentionIndexerParams& params, const T* query, const T* key, const T* query_norm_weight, const T* key_norm_weight, const T* cos_cache, const T* sin_cache, const int32_t* cumulative_sequence_lengths, - const int32_t* past_sequence_lengths, const int64_t* position_ids, const T* past_key_state, + const int32_t* past_sequence_lengths, const int64_t* position_ids, const TState* past_key_state, const T* past_kv_buffer, const int32_t* past_state_lengths, const int32_t* state_update_capture_count, bool state_update_active, int32_t* selected_indices, - int32_t* selected_counts, T* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths, + int32_t* selected_counts, TState* present_key_state, T* present_kv_buffer, int32_t* present_state_lengths, T* state_update, float* float_workspace, int32_t* overflow_flags) { if (params.batch_size > 0) { const int64_t key_state_elems = static_cast(params.batch_size) * params.state_capacity * params.head_size; if (key_state_elems > 0 && present_key_state != past_key_state) { - ElementwiseCopyKernel<<>>( + ElementwiseCopyKernel<<>>( past_key_state, present_key_state, key_state_elems); } const int64_t buffer_elems = static_cast(params.batch_size) * params.buffer_capacity * params.head_size; @@ -1530,9 +1658,9 @@ Status LaunchQsaPackedSparseAttentionIndexer( const size_t value_bytes = static_cast(params.head_size) * sizeof(float); ORT_RETURN_IF_ERROR(ConfigureDynamicSharedMemory( - QsaUpdateStateKernel, 2 * value_bytes + kThreads * sizeof(float), "QsaUpdateStateKernel")); + QsaUpdateStateKernel, 2 * value_bytes + kThreads * sizeof(float), "QsaUpdateStateKernel")); const int state_blocks = static_cast(std::min(params.batch_size, kSaiMaxGridDimX)); - QsaUpdateStateKernel<<>>( + QsaUpdateStateKernel<<>>( key, key_norm_weight, cos_cache, sin_cache, cumulative_sequence_lengths, past_sequence_lengths, past_kv_buffer, past_state_lengths, state_update_capture_count, state_update_active, present_key_state, present_kv_buffer, present_state_lengths, state_update, overflow_flags, params); @@ -1565,8 +1693,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 = @@ -1576,12 +1704,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); } } @@ -1592,18 +1730,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); } @@ -1708,21 +1846,29 @@ Status LaunchCsaPackedSparseAttentionIndexer( return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(T) \ - template Status LaunchQsaPackedSparseAttentionIndexer( \ - cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ - const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const T*, const T*, const int32_t*, \ - const int32_t*, bool, int32_t*, int32_t*, T*, T*, int32_t*, T*, float*, int32_t*); \ +#define INSTANTIATE_QSA(T, TState) \ + template Status LaunchQsaPackedSparseAttentionIndexer( \ + cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ + const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const TState*, const T*, \ + const int32_t*, const int32_t*, bool, int32_t*, int32_t*, TState*, T*, int32_t*, T*, float*, int32_t*); + +#define INSTANTIATE_CSA(T) \ template Status LaunchCsaPackedSparseAttentionIndexer( \ cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ const T*, const T*, const T*, const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const T*, \ const T*, const T*, const int32_t*, int32_t*, int32_t*, T*, T*, T*, int32_t*, float*, int32_t*); -INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(float) -INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(half) -INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER(__nv_bfloat16) +INSTANTIATE_QSA(float, float) +INSTANTIATE_QSA(half, half) +INSTANTIATE_QSA(__nv_bfloat16, __nv_bfloat16) +INSTANTIATE_QSA(__nv_bfloat16, float) + +INSTANTIATE_CSA(float) +INSTANTIATE_CSA(half) +INSTANTIATE_CSA(__nv_bfloat16) -#undef INSTANTIATE_PACKED_SPARSE_ATTENTION_INDEXER +#undef INSTANTIATE_QSA +#undef INSTANTIATE_CSA } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.h index ab4dcd4849116..9d4f5e4e9a3eb 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; @@ -59,7 +61,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, @@ -72,20 +74,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/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 5b88eae6b079a..ca05c3a2cbe74 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -483,7 +483,6 @@ void GroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); - const int64_t num_heads = getAttribute(ctx, "num_heads", 0); const int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0); if (num_heads <= 0 || num_heads > std::numeric_limits::max()) { @@ -504,7 +503,7 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex auto output_shape = query_shape; if (ctx.getInputType(2) == nullptr) { output_shape.mutable_dim(2)->Clear(); - if (query_dims[2].has_dim_value()) { + 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(); if (packed_hidden_size % packed_heads != 0) { @@ -524,7 +523,8 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex continue; } - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, output_index); + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput( + ctx, ctx.getInputType(input_index) != nullptr ? input_index : 0, output_index); if (hasInputShape(ctx, input_index)) { ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, input_index, output_index); continue; @@ -532,7 +532,8 @@ void DynamicSparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContex if (hasInputShape(ctx, 0)) { const auto& query_dims = getInputShape(ctx, 0).dim(); - if (query_dims.size() == 3 && query_dims[2].has_dim_value()) { + 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; const int64_t divisor = is_packed ? num_heads + 2 * kv_num_heads : num_heads; if (query_dims[2].dim_value() % divisor != 0) { @@ -589,6 +590,12 @@ The main cache uses BNSH layout and ordinary contiguous append semantics. Select Auxiliary K/V use BNSH layout and are read-only. When `auxiliary_kv_shared=1`, auxiliary_value may be omitted and auxiliary_key is used as both K and V. DeepSeek-style post-attention inverse/conjugate RoPE and output projection remain the exporter's responsibility. + +Main caches may instead use INT8 with positive finite float32 key_scale and value_scale vectors, one scale per +KV-head channel. Newly appended keys (after optional normalization and RoPE) and values are divided by their scales, +rounded to nearest with ties to even, and saturated to [-128, 127]. Attention reads the quantized representation, +including newly appended rows. Auxiliary caches remain floating point. Invalid scale values produce validation_status=10 +when validation is requested; CUDA strict validation also reports them as an error. )DOC"; ONNX_MS_OPERATOR_SET_SCHEMA( @@ -629,8 +636,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(2, "value", "Current main value [batch, sequence, kv_num_heads * head_size].", "T", OpSchema::Optional) - .Input(3, "past_key", "Main key cache in BNSH layout.", "T", OpSchema::Optional) - .Input(4, "past_value", "Main value cache in BNSH layout.", "T", OpSchema::Optional) + .Input(3, "past_key", "Main key cache in BNSH layout, same type as query or INT8.", "TCache", OpSchema::Optional) + .Input(4, "past_value", "Main value cache in BNSH layout, same type as past_key.", "TCache", OpSchema::Optional) .Input(5, "auxiliary_key", "Read-only auxiliary key sequence in BNSH layout.", "T", OpSchema::Optional) .Input(6, "auxiliary_value", "Read-only auxiliary value sequence in BNSH layout.", "T", @@ -652,13 +659,34 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Input(16, "head_sink", "Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.", "T", OpSchema::Optional) + .Input(17, "committed_length", + "Optional single-request committed token count for cursor validation. When supplied, validation_status " + "is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent.", + "M", OpSchema::Optional) + .Input(18, "key_scale", "Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].", + "S", OpSchema::Optional) + .Input(19, "value_scale", "Positive finite per-channel INT8 value-cache scales [kv_num_heads * head_size].", + "S", OpSchema::Optional) .Output(0, "output", "Attention output [batch, sequence, num_heads * head_size].", "T") - .Output(1, "present_key", "Updated main key cache in BNSH layout.", "T", OpSchema::Optional) - .Output(2, "present_value", "Updated main value cache in BNSH layout.", "T", OpSchema::Optional) + .Output(1, "present_key", "Updated main key cache in BNSH layout.", "TCache", OpSchema::Optional) + .Output(2, "present_value", "Updated main value cache in BNSH layout.", "TCache", OpSchema::Optional) + .Output(3, "validation_status", "Optional int32 status: 0 when valid, a validation code when invalid.", + "M", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain all floating-point inputs and outputs to one element type.") + .TypeConstraint("TCache", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)", "tensor(int8)"}, + "Cache type must match T unless INT8 per-channel scales are supplied.") + .TypeConstraint("S", {"tensor(float)"}, "INT8 cache scales use float32.") .TypeConstraint("M", {"tensor(int32)"}, "Constrain selection and sequence metadata to int32.") - .TypeAndShapeInferenceFunction(DynamicSparseAttentionTypeAndShapeInference)); + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + DynamicSparseAttentionTypeAndShapeInference(ctx); + if (ctx.getNumOutputs() > 3 && ctx.hasOutput(3)) { + updateOutputElemType(ctx, 3, ONNX_NAMESPACE::TensorProto_DataType_INT32); + ONNX_NAMESPACE::TensorShapeProto status_shape; + status_shape.add_dim()->set_dim_value(1); + updateOutputShape(ctx, 3, status_shape); + } + })); constexpr const char* Attention_ver1_doc = R"DOC( Multi-Head Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT). @@ -2861,6 +2889,15 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference const bool has_capture_count = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateCaptureCount); const bool has_state_update_active = PackedSparseAttentionIndexerHasInput(ctx, psai::kStateUpdateActive); const bool has_state_update_output = ctx.hasOutput(psai::kStateUpdate); + const bool has_mid_kv_buffer = ctx.hasOutput(psai::kMidKvBuffer); + const bool has_mid_state_lengths = ctx.hasOutput(psai::kMidStateLengths); + const bool has_validation_status = ctx.hasOutput(psai::kValidationStatus); + if (has_mid_kv_buffer != has_mid_state_lengths) { + fail_shape_inference("PackedSparseAttentionIndexer: mid_kv_buffer and mid_state_lengths must be connected together"); + } + if (!is_qsa && (has_mid_kv_buffer || has_validation_status)) { + fail_shape_inference("PackedSparseAttentionIndexer: mid-state and validation outputs are only supported for qsa"); + } if (!is_qsa && (state_update_capacity > 0 || has_capture_count || has_state_update_active || has_state_update_output)) { fail_shape_inference("PackedSparseAttentionIndexer: state update capture is only valid for policy_mode 'qsa'"); @@ -2872,9 +2909,9 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference } if (ctx.getNumOutputs() < static_cast(psai::kFixedOutputCount) || - ctx.getNumOutputs() > static_cast(has_indexshare ? psai::kOutputCount : psai::kStateUpdate + 1)) { - fail_shape_inference("PackedSparseAttentionIndexer: expected ", psai::kFixedOutputCount, " or ", - psai::kOutputCount, " declared outputs, got ", ctx.getNumOutputs()); + ctx.getNumOutputs() > static_cast(psai::kMaxOutputCount)) { + fail_shape_inference("PackedSparseAttentionIndexer: expected ", psai::kFixedOutputCount, " to ", + psai::kMaxOutputCount, " declared outputs, got ", ctx.getNumOutputs()); } if (ctx.hasOutput(psai::kPresentGateBuffer) == is_qsa) { fail_shape_inference("PackedSparseAttentionIndexer: output ", psai::kPresentGateBuffer, @@ -2885,8 +2922,20 @@ void PackedSparseAttentionIndexerTypeAndShapeInference(ONNX_NAMESPACE::Inference updateOutputElemType(ctx, psai::kSelectedCounts, ONNX_NAMESPACE::TensorProto_DataType_INT32); if (has_indexshare) updateOutputElemType(ctx, psai::kIndexShareStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); updateOutputElemType(ctx, psai::kPresentStateLengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); + if (has_mid_state_lengths) { + updateOutputElemType(ctx, psai::kMidStateLengths, ONNX_NAMESPACE::TensorProto_DataType_INT32); + propagateShapeFromInputToOutput(ctx, psai::kPastStateLengths, psai::kMidStateLengths); + } + if (has_validation_status) { + updateOutputElemType(ctx, psai::kValidationStatus, ONNX_NAMESPACE::TensorProto_DataType_INT32); + propagateShapeFromInputToOutput(ctx, psai::kPastSequenceLengths, psai::kValidationStatus); + } propagateElemTypeFromInputToOutput(ctx, psai::kPastKeyState, psai::kPresentKeyState); propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kPresentKvBuffer); + if (has_mid_kv_buffer) { + propagateElemTypeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); + propagateShapeFromInputToOutput(ctx, psai::kPastKvBuffer, psai::kMidKvBuffer); + } if (!is_qsa) { // present_gate_buffer keeps its fixed positional slot (with an empty name) for policy_mode // 'qsa'; only propagate its type/shape when it is actually produced. @@ -3154,6 +3203,8 @@ Unlike SparseAttentionIndexer, this operator: rejected as a deterministic no-op on state rather than truncated or allowed to corrupt memory; * additionally emits selected_counts, the exact number of active (non -1) entries per query, so that no downstream consumer needs to scan selected_indices for its query's true count. + * for QSA, optionally emits the pending-state snapshot after the first request token and a + per-request validation status for overflow or inconsistent state metadata. Both policy_mode values keep the semantics of SparseAttentionIndexer, applied independently to each request's own packed token range and fixed-capacity state slice: @@ -3318,8 +3369,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "past_key_state", "Generic fixed-capacity state: policy_mode 'qsa' stores prepared complete-block keys; " "policy_mode 'csa' stores compressed keys. Shape is (batch_size, state_capacity, head_size) and " - "never changes across calls.", - "T") + "never changes across calls. For QSA, this state may use FLOAT while query/key tensors use " + "FLOAT16 or BFLOAT16.", + "T_KEY") .Input(13, "past_kv_buffer", "Generic fixed-capacity pending-token buffer. Shape is (batch_size, 2 * compress_ratio - 1, " @@ -3373,7 +3425,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Output(2, "present_key_state", "Updated generic key state, with the same fixed shape as past_key_state.", - "T") + "T_KEY") .Output(3, "present_kv_buffer", "Updated generic pending-token buffer, with the same fixed shape as past_kv_buffer.", @@ -3397,9 +3449,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..bc648f4f5ba37 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -88,6 +88,10 @@ struct DynamicSparseAttentionCase { std::vector q_norm_weight; std::vector k_norm_weight; std::vector head_sink; + std::vector key_scale; + std::vector value_scale; + std::optional committed_length; + std::optional expected_validation_status; std::vector expected_output; std::vector expected_present_key; @@ -108,7 +112,7 @@ std::vector ToTensorData(const std::vector& values) { } } -template +template void RunDynamicSparseAttentionCase( const DynamicSparseAttentionCase& c, std::unique_ptr execution_provider, @@ -163,12 +167,12 @@ void RunDynamicSparseAttentionCase( ToTensorData(c.value)); } if (c.has_past) { - tester.AddInput("past_key", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.past_key)); - tester.AddInput("past_value", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.past_value)); + tester.AddInput("past_key", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.past_key)); + tester.AddInput("past_value", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.past_value)); } else { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); @@ -229,15 +233,33 @@ void RunDynamicSparseAttentionCase( } else { tester.AddInput("head_sink", {c.num_heads}, ToTensorData(c.head_sink)); } + if (c.committed_length.has_value()) { + tester.AddInput("committed_length", {1}, {*c.committed_length}); + } else { + tester.AddOptionalInputEdge(); + } + if (!c.key_scale.empty() || !c.value_scale.empty()) { + if (c.key_scale.empty()) + tester.AddOptionalInputEdge(); + else + tester.AddInput("key_scale", {int64_t(c.key_scale.size())}, c.key_scale); + if (c.value_scale.empty()) + tester.AddOptionalInputEdge(); + else + tester.AddInput("value_scale", {int64_t(c.value_scale.size())}, c.value_scale); + } tester.AddOutput("output", {c.batch_size, c.sequence_length, hidden_size}, ToTensorData(c.expected_output)); - tester.AddOutput("present_key", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.expected_present_key)); - tester.AddOutput("present_value", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.expected_present_value)); + tester.AddOutput("present_key", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.expected_present_key)); + tester.AddOutput("present_value", + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.expected_present_value)); + if (c.expected_validation_status.has_value()) { + tester.AddOutput("validation_status", {1}, {*c.expected_validation_status}); + } tester.SetOutputTolerance(0.005f); std::vector> execution_providers; @@ -276,6 +298,283 @@ DynamicSparseAttentionCase MakeSingleTokenSelectedValueCase(float value = 9.0f) } #ifdef USE_CUDA +DynamicSparseAttentionCase MakeInt8CacheCase() { + DynamicSparseAttentionCase c; + c.sequence_length = 2; + c.num_heads = 4; + c.kv_num_heads = 2; + c.cache_sequence_length = 4; + c.max_selected = 3; + c.query.assign(64, 0.0f); + c.key.resize(32); + c.value.resize(32); + c.key_scale.resize(16); + c.value_scale.resize(16); + c.past_key.resize(64); + c.past_value.resize(64); + for (size_t i = 0; i < 64; ++i) { + c.past_key[i] = float(int(i % 13) - 6); + c.past_value[i] = float(int(i % 17) - 8); + } + for (size_t i = 0; i < 16; ++i) { + c.key_scale[i] = i < 8 ? 0.5f : 0.25f; + c.value_scale[i] = i < 8 ? 0.25f : 0.5f; + } + for (size_t i = 0; i < 32; ++i) { + c.key[i] = i % 3 == 0 ? -300.0f : 0.25f * float(i); + c.value[i] = i % 3 == 0 ? 300.0f : 0.125f * float(i); + } + c.key[0] = -1e20f; + c.value[0] = 1e20f; + c.selected_indices = {0, 2, -1, 1, 2, 3}; + c.selected_counts = {2, 3}; + c.seqlens_k = {3}; + c.total_sequence_length = 4; + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + for (int s = 0; s < 2; ++s) { + for (int h = 0; h < 2; ++h) { + for (int d = 0; d < 8; ++d) { + const int channel = h * 8 + d; + const int src = s * 16 + channel; + const int dst = (h * 4 + 2 + s) * 8 + d; + c.expected_present_key[dst] = std::clamp(std::nearbyint(c.key[src] / c.key_scale[channel]), -128.0f, 127.0f); + c.expected_present_value[dst] = std::clamp(std::nearbyint(c.value[src] / c.value_scale[channel]), -128.0f, 127.0f); + } + } + } + c.expected_output.assign(64, 0.0f); + for (int s = 0; s < 2; ++s) { + for (int h = 0; h < 4; ++h) { + const int kv = h / 2; + for (int d = 0; d < 8; ++d) { + float sum = 0.0f; + for (int i = 0; i < c.selected_counts[s]; ++i) { + const int pos = c.selected_indices[s * 3 + i]; + sum += c.expected_present_value[(kv * 4 + pos) * 8 + d] * c.value_scale[kv * 8 + d]; + } + c.expected_output[(s * 4 + h) * 8 + d] = sum / c.selected_counts[s]; + } + } + } + return c; +} + +TEST(DynamicSparseAttentionTest, Int8CachePerChannelGqaAppendAndAttention_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase(MakeInt8CacheCase(), std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheBFloat16_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase(MakeInt8CacheCase(), std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheSplitAttention_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + auto c = MakeInt8CacheCase(); + c.cache_sequence_length = 1024; + c.total_sequence_length = 1024; + c.seqlens_k = {1023}; + c.max_selected = 512; + c.selected_counts = {512, 512}; + c.selected_indices.resize(1024); + std::iota(c.selected_indices.begin(), c.selected_indices.end(), 0); + c.past_key.assign(2 * 1024 * 8, 0.0f); + c.past_value.assign(2 * 1024 * 8, 3.0f); + auto key = c.past_key; + auto value = c.past_value; + for (int h = 0; h < 2; ++h) { + for (int s = 0; s < 2; ++s) { + for (int d = 0; d < 8; ++d) { + key[(h * 1024 + 1022 + s) * 8 + d] = c.expected_present_key[(h * 4 + 2 + s) * 8 + d]; + value[(h * 1024 + 1022 + s) * 8 + d] = c.expected_present_value[(h * 4 + 2 + s) * 8 + d]; + } + } + } + c.expected_present_key = key; + c.expected_present_value = value; + for (int s = 0; s < 2; ++s) { + for (int h = 0; h < 4; ++h) { + for (int d = 0; d < 8; ++d) { + float sum = 0.0f; + for (int i = 0; i < 512; ++i) { + sum += value[((h / 2) * 1024 + s * 512 + i) * 8 + d] * c.value_scale[(h / 2) * 8 + d]; + } + c.expected_output[(s * 4 + h) * 8 + d] = sum / 512.0f; + } + } + } + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheRejectsMissingOrMalformedScales_CUDA) { + for (int variant = 0; variant < 3; ++variant) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + auto c = MakeInt8CacheCase(); + if (variant == 0) { + c.key_scale.clear(); + c.value_scale.clear(); + } else if (variant == 1) { + c.value_scale.clear(); + } else { + c.key_scale.resize(8); + } + const char* error = variant == 0 ? "INT8 caches require per-channel" + : variant == 1 ? "must be provided together" + : "INT8 cache scales must have shape"; + RunDynamicSparseAttentionCase(c, std::move(cuda_ep), + OpTester::ExpectResult::kExpectFailure, error); + } +} + +DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool auxiliary) { + DynamicSparseAttentionCase c; + c.batch_size = 2; + c.sequence_length = sequence_length; + c.num_heads = 20; + c.kv_num_heads = 2; + c.head_size = 256; + c.cache_sequence_length = 520; + c.max_selected = 512; + c.attention_mode = auxiliary ? "local_plus_selected" : "selected_only"; + c.local_window_size = auxiliary ? 16 : -1; + c.smooth_softmax = 1; + c.scale = 0.0625f; + c.head_sink.assign(c.num_heads, 0.25f); + c.seqlens_k = {519, 518}; + c.total_sequence_length = 520; + const int channels = static_cast(c.kv_num_heads * c.head_size); + const int rows = static_cast(c.batch_size * c.sequence_length); + c.query.resize(rows * c.num_heads * c.head_size); + c.key.resize(rows * channels); + c.value.resize(rows * channels); + c.key_scale.resize(channels); + c.value_scale.resize(channels); + for (int d = 0; d < channels; ++d) { + c.key_scale[d] = d % 2 ? 0.125f : 0.0625f; + c.value_scale[d] = d % 2 ? 0.25f : 0.125f; + } + for (size_t i = 0; i < c.query.size(); ++i) { + c.query[i] = float(int(i % 7) - 3) * 0.125f; + } + for (size_t i = 0; i < c.key.size(); ++i) { + c.key[i] = float(int(i % 11) - 5) * c.key_scale[i % channels]; + c.value[i] = float(int(i % 13) - 6) * c.value_scale[i % channels]; + } + c.past_key.resize(c.batch_size * c.cache_sequence_length * channels); + c.past_value.resize(c.past_key.size()); + for (size_t i = 0; i < c.past_key.size(); ++i) { + c.past_key[i] = float(int(i % 17) - 8); + c.past_value[i] = float(int(i % 19) - 9); + } + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + for (int row = 0; row < rows; ++row) { + const int b = row / sequence_length; + const int pos = c.seqlens_k[b] + 1 - sequence_length + row % sequence_length; + for (int d = 0; d < channels; ++d) { + const int64_t dst = + ((b * c.kv_num_heads + d / c.head_size) * c.cache_sequence_length + pos) * c.head_size + + d % c.head_size; + c.expected_present_key[dst] = c.key[row * channels + d] / c.key_scale[d]; + c.expected_present_value[dst] = c.value[row * channels + d] / c.value_scale[d]; + } + } + if (auxiliary) { + c.selected_kv_source = "auxiliary"; + c.auxiliary_sequence_length = 512; + c.auxiliary_key.resize(c.batch_size * c.auxiliary_sequence_length * channels); + c.auxiliary_value.resize(c.auxiliary_key.size()); + for (size_t i = 0; i < c.auxiliary_key.size(); ++i) { + c.auxiliary_key[i] = float(int(i % 17) - 8) * 0.125f; + c.auxiliary_value[i] = float(int(i % 19) - 9) * 0.25f; + } + } + c.selected_counts.resize(rows); + c.selected_indices.assign(rows * c.max_selected, -1); + c.expected_output.assign(c.query.size(), 0.0f); + for (int row = 0; row < rows; ++row) { + const int b = row / sequence_length; + const int pos = c.seqlens_k[b] + 1 - sequence_length + row % sequence_length; + c.selected_counts[row] = row == 0 ? 0 : 509; + std::iota(c.selected_indices.begin() + row * c.max_selected, + c.selected_indices.begin() + row * c.max_selected + c.selected_counts[row], 0); + for (int h = 0; h < c.num_heads; ++h) { + const int kv = h / (c.num_heads / c.kv_num_heads); + const int64_t out = (row * c.num_heads + h) * c.head_size; + std::vector numerator(c.head_size, 0.0f); + double denominator = std::exp(double(c.head_sink[h])); + const int local_count = auxiliary ? static_cast(c.local_window_size) : 0; + const int local_start = pos - local_count + 1; + for (int i = 0; i < local_count + c.selected_counts[row]; ++i) { + const bool selected = i >= local_count; + const int index = selected ? i - local_count : local_start + i; + if (selected && !auxiliary && index >= local_start) continue; + const bool from_auxiliary = selected && auxiliary; + const int64_t length = from_auxiliary ? c.auxiliary_sequence_length : c.cache_sequence_length; + const int64_t offset = ((b * c.kv_num_heads + kv) * length + index) * c.head_size; + double dot = 0.0; + for (int d = 0; d < c.head_size; ++d) { + const float key = from_auxiliary ? c.auxiliary_key[offset + d] + : c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d]; + dot += c.query[out + d] * key; + } + const double weight = std::exp(dot * c.scale); + denominator += weight; + for (int d = 0; d < c.head_size; ++d) { + const float value = from_auxiliary ? c.auxiliary_value[offset + d] + : c.expected_present_value[offset + d] * c.value_scale[kv * c.head_size + d]; + numerator[d] += static_cast(weight * value); + } + } + for (int d = 0; d < c.head_size; ++d) { + c.expected_output[out + d] = static_cast(numerator[d] / denominator); + } + } + } + return c; +} + +TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecode_CUDA) { + for (int sequence_length : {1, 2}) { + for (bool auxiliary : {false, true}) { + SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + + ", auxiliary=" + std::to_string(auxiliary)); + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase( + MakeGroupedInt8CacheCase(sequence_length, auxiliary), std::move(cuda_ep)); + } + } +} + +TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeFloat16_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + RunDynamicSparseAttentionCase( + MakeGroupedInt8CacheCase(2, false), std::move(cuda_ep)); +} + +TEST(DynamicSparseAttentionTest, Int8CacheRejectsNonpositiveOrNonfiniteScales_CUDA) { + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DYNAMIC_SPARSE_ATTENTION_STRICT_VALIDATION", "1"}}}; + for (float invalid : {0.0f, -1.0f, std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) GTEST_SKIP(); + auto c = MakeInt8CacheCase(); + c.value_scale[15] = invalid; + RunDynamicSparseAttentionCase( + c, std::move(cuda_ep), OpTester::ExpectResult::kExpectFailure, + "INT8 cache scales must be positive and finite"); + } +} + Status ResolveDynamicSparseAttentionGraph(size_t output_count, int64_t query_width = 8, int64_t num_heads = 1, int64_t kv_num_heads = 1, bool packed_qkv = false, const char* query_width_symbol = nullptr, @@ -450,6 +749,32 @@ TEST(DynamicSparseAttentionTest, SelectedOnlyMainVariableCountsGqaAndCacheAppend RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); } +TEST(DynamicSparseAttentionTest, ValidationStatusReportsCommittedCursor_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available."; + } + + auto c = MakeSingleTokenSelectedOnlyCase(); + c.cache_sequence_length = 2; + c.total_sequence_length = 2; + c.past_key.assign(16, 0.0f); + c.past_value.assign(16, 0.0f); + c.key.assign(8, 0.0f); + c.value.assign(8, 0.0f); + c.expected_present_key.assign(16, 0.0f); + c.expected_present_value.assign(16, 0.0f); + c.committed_length = 0; + c.expected_validation_status = 0; + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); + + cuda_ep = DefaultCudaExecutionProvider(); + ASSERT_NE(cuda_ep, nullptr); + c.committed_length = 1; + c.expected_validation_status = 9; + RunDynamicSparseAttentionCase(c, std::move(cuda_ep)); +} + TEST(DynamicSparseAttentionTest, LocalPlusSelectedAuxiliaryJointSoftmaxSinkSharedKv_CUDA) { auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) { diff --git a/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/packed_sparse_attention_indexer_op_test.cc index 27406a7a76cac..8aa9025d726c1 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 @@ -905,7 +905,7 @@ QsaPackedProblem MakeQsaPackedProblem(QsaPackedProblem problem = {}) { return problem; } -template +template void RunQsaPackedTest(float tolerance, QsaPackedProblem problem = MakeQsaPackedProblem(), ProviderKind provider_kind = ProviderKind::Cuda, QsaPackedResult* actual = nullptr, bool packed_qk = false, const ConfigOptions* webgpu_options = nullptr, int padding = 0) { @@ -921,7 +921,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; @@ -982,8 +982,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 @@ -1006,8 +1006,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 @@ -1023,7 +1023,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()); @@ -1038,6 +1038,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); } @@ -1345,6 +1349,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);