[CUDA] Support per-batch committed_length in DynamicSparseAttention

onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc

@@ -165,8 +165,9 @@ Status DynamicSparseAttention<T>::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<T>::ComputeInternal(OpKernelContext* context) cons
             parameters, GetDeviceProp().sharedMemPerBlock);
     auto validation_workspace =
         GetScratchBuffer<uint32_t>(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<int32_t>()));
+        committed_length == nullptr ? nullptr : committed_length->Data<int32_t>(),
+        committed_length_stride));
   }
 
   const size_t attention_workspace_elements =

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<int>(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<int64_t>(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) {

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 <typename T>
 Status LaunchDynamicSparseAttention(

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)