diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index 9555b9f919cd6..d3e020eb16e2e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -165,8 +165,9 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons ORT_RETURN_IF_NOT(parameters.head_size <= GetDeviceProp().maxThreadsPerBlock, "DynamicSparseAttention: head_size exceeds the CUDA device thread-block limit."); ORT_RETURN_IF(committed_length != nullptr && - (parameters.batch_size != 1 || committed_length->Shape().Size() != 1), - "DynamicSparseAttention: committed_length requires batch_size=1 and one int32 element."); + committed_length->Shape().Size() != 1 && + committed_length->Shape().Size() != parameters.batch_size, + "DynamicSparseAttention: committed_length must have 1 or batch_size int32 elements."); TensorShape output_shape({parameters.batch_size, parameters.sequence_length, parameters.query_hidden_size}); @@ -253,11 +254,17 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons parameters, GetDeviceProp().sharedMemPerBlock); auto validation_workspace = GetScratchBuffer(validation_workspace_elements, GetComputeStream(context)); + const int committed_length_stride = + (committed_length != nullptr && committed_length->Shape().Size() == parameters.batch_size && + parameters.batch_size > 1) + ? 1 + : 0; ORT_RETURN_IF_ERROR(ValidateDynamicSparseAttentionOnDevice( stream, data.selected_indices, data.selected_counts, data.seqlens_k, data.position_ids, parameters, validation_error, validation_workspace.get(), GetDeviceProp().sharedMemPerBlock, strict_validation_, data.key_scale, data.value_scale, - committed_length == nullptr ? nullptr : committed_length->Data())); + committed_length == nullptr ? nullptr : committed_length->Data(), + committed_length_stride)); } const size_t attention_workspace_elements = diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index 83d3536f0b443..d489526959ea7 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -105,7 +105,8 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, const float* key_scale, const float* value_scale, int cache_channels, - const int32_t* committed_length) { + const int32_t* committed_length, + int committed_length_stride) { const int row = static_cast(blockIdx.x); const int b = row / sequence_length; const int s = row - b * sequence_length; @@ -232,10 +233,10 @@ __global__ void ValidateInputsKernel(const int32_t* selected_indices, return; } - if (row == 0 && threadIdx.x == 0 && committed_length != nullptr) { - const int64_t committed = *committed_length; + if (s == 0 && threadIdx.x == 0 && committed_length != nullptr) { + const int64_t committed = committed_length[b * committed_length_stride]; const int64_t total = committed + sequence_length; - if (committed < 0 || total > main_capacity || seqlens_k[0] != total - 1) { + if (committed < 0 || total > main_capacity || seqlens_k[b] != total - 1) { atomicCAS(error_flag, kDynamicSparseAttentionValidationOk, 9); } } @@ -1865,7 +1866,8 @@ Status ValidateDynamicSparseAttentionOnDevice( bool copy_result_to_host, const float* key_scale, const float* value_scale, - const int32_t* committed_length) { + const int32_t* committed_length, + int committed_length_stride) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(error_flag, 0, sizeof(int32_t), stream)); const int64_t row_count = static_cast(parameters.batch_size) * parameters.sequence_length; @@ -1895,7 +1897,7 @@ Status ValidateDynamicSparseAttentionOnDevice( parameters.selected_kv_source == DynamicSparseAttentionKvSource::kAuxiliary, validation_bitmap, validation_hash_entries, use_shared_hash_table, error_flag, key_scale, value_scale, parameters.kv_num_heads * parameters.head_size, - committed_length); + committed_length, committed_length_stride); CUDA_RETURN_IF_ERROR(cudaGetLastError()); if (!copy_result_to_host) { diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index 6050386bd9c71..ed71674276a95 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -76,7 +76,8 @@ Status ValidateDynamicSparseAttentionOnDevice( bool copy_result_to_host, const float* key_scale = nullptr, const float* value_scale = nullptr, - const int32_t* committed_length = nullptr); + const int32_t* committed_length = nullptr, + int committed_length_stride = 0); template Status LaunchDynamicSparseAttention( diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 768f88cde070f..d31d2fdfc01d9 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -659,8 +659,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Optional sink logit [num_heads]. Participates in the shared softmax denominator in both attention modes.", "T", OpSchema::Optional) .Input(17, "committed_length", - "Optional single-request committed token count for cursor validation. When supplied, validation_status " - "is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent. CUDA only.", + "Optional per-request committed token count ([1] or [batch]) for cursor validation. When supplied, " + "validation_status is 9 if seqlens_k, sequence_length, or cache_capacity is inconsistent. CUDA only.", "M", OpSchema::Optional) .Input(18, "key_scale", "Positive finite per-channel INT8 key-cache scales [kv_num_heads * head_size].", "S", OpSchema::Optional)