Validate SparseAttention CSR indices and key lengths element values

onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc

@@ -36,6 +36,74 @@ namespace contrib {
 REGISTER_KERNEL_TYPED(float)
 REGISTER_KERNEL_TYPED(MLFloat16)
 
+namespace {
+
+// Validate CSR row-pointer monotonicity and column-index range.
+// Must be called after CheckInputs has populated the parameters struct.
+Status ValidateCSRIndices(const SparseAttentionParameters& parameters,
+                          const Tensor& block_row_indices,
+                          const Tensor& block_col_indices) {
+  const int num_layout = parameters.num_sparse_layout;
+  const int max_blocks = parameters.stride_row_indices - 1;
+  const int col_count = parameters.stride_col_indices;
+
+  const int32_t* row_data = block_row_indices.Data<int32_t>();
+  const int32_t* col_data = block_col_indices.Data<int32_t>();
+  for (int l = 0; l < num_layout; ++l) {
+    const int32_t* r = row_data + l * (max_blocks + 1);
+    if (r[0] != 0) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "block_row_indices[", l, "][0] must be 0, got ", r[0]);
+    }
+    for (int i = 0; i < max_blocks; ++i) {
+      if (r[i] < 0 || r[i] > r[i + 1] || r[i + 1] > col_count) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "block_row_indices values are not monotonically non-decreasing or exceed "
+                               "block_col_indices columns at layout ",
+                               l, " row ", i,
+                               ": r[", i, "]=", r[i], ", r[", i + 1, "]=", r[i + 1],
+                               ", col_count=", col_count);
+      }
+    }
+    const int32_t* c = col_data + l * col_count;
+    const int nnz = r[max_blocks];
+    for (int k = 0; k < nnz; ++k) {
+      if (c[k] < 0 || c[k] >= max_blocks) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "block_col_indices[", l, "][", k, "]=", c[k],
+                               " is out of valid range [0, ", max_blocks, ")");
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+// Validate total_key_lengths element values.
+Status ValidateKeyLengths(const SparseAttentionParameters& parameters,
+                          const Tensor& total_key_lengths) {
+  const int batch_size = parameters.batch_size;
+  const int sequence_length = parameters.sequence_length;
+  const int total_sequence_length = parameters.total_sequence_length;
+
+  const auto* key_len_data = total_key_lengths.Data<int32_t>();
+  const bool is_prompt = (sequence_length == total_sequence_length);
+  const int min_key_length = is_prompt ? 1 : sequence_length;
+  for (int i = 0; i < batch_size; ++i) {
+    const int key_length = key_len_data[i];
+    if (key_length < min_key_length || key_length > total_sequence_length) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "key_total_sequence_lengths value ", key_length,
+                             " at batch index ", i,
+                             " is out of range [", min_key_length, ", ", total_sequence_length, "].");
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace
+
 template <typename T>
 SparseAttention<T>::SparseAttention(const OpKernelInfo& info) : OpKernel(info), SparseAttentionBase(info) {
 }
@@ -75,6 +143,11 @@ Status SparseAttention<T>::Compute(OpKernelContext* context) const {
                                                            block_col_indices,
                                                            total_key_lengths,
                                                            total_seq_len));
+  ORT_RETURN_IF_ERROR(ValidateCSRIndices(parameters,
+                                         *block_row_indices,
+                                         *block_col_indices));
+  ORT_RETURN_IF_ERROR(ValidateKeyLengths(parameters,
+                                         *total_key_lengths));
 
   const int batch_size = parameters.batch_size;
   const int sequence_length = parameters.sequence_length;

onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h

@@ -198,26 +198,13 @@ Status CheckInputs(void* params,
                            past_key_dims[3]);
   }
 
-  // Check the shape and values of total_key_sequence_lengths.
+  // Check the shape of total_key_sequence_lengths.
   const auto& k_len_dim = total_key_lengths->Shape().GetDims();
   if (k_len_dim.size() != 1 || k_len_dim[0] != batch_size) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                            "key_total_sequence_lengths must have shape (batch_size).");
   }
 
-  const auto* key_len_data = total_key_lengths->Data<int32_t>();
-  const bool is_prompt = (sequence_length == total_sequence_length);
-  const int min_key_length = is_prompt ? 1 : sequence_length;
-  for (int i = 0; i < batch_size; ++i) {
-    const int key_length = key_len_data[i];
-    if (key_length < min_key_length || key_length > total_sequence_length) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "key_total_sequence_lengths value ", key_length,
-                             " at batch index ", i,
-                             " is out of range [", min_key_length, ", ", total_sequence_length, "].");
-    }
-  }
-
   int rotary_dim = 0;
   int max_rotary_sequence_length = 0;
   if (do_rotary) {

onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc

@@ -215,6 +215,23 @@ Status SparseAttention<T>::ComputeInternal(OpKernelContext* context) const {
   data.kernel_layout.csr_col_indices = block_col_indices->Data<int32_t>();
   data.kernel_layout.csr_row_indices = block_row_indices->Data<int32_t>();
 
+  // Validate CSR indices and key lengths on device to prevent out-of-bounds access.
+  {
+    auto csr_error_buffer = GetScratchBuffer<int32_t>(1, GetComputeStream(context));
+    ORT_RETURN_IF_ERROR(ValidateCSRIndicesOnDevice(
+        cuda_stream,
+        data.kernel_layout.csr_row_indices,
+        data.kernel_layout.csr_col_indices,
+        seqlens_k_total->Data<int32_t>(),
+        parameters.num_sparse_layout,
+        parameters.stride_row_indices - 1,  // max_blocks
+        parameters.stride_col_indices,      // col_count
+        parameters.batch_size,
+        parameters.sequence_length,
+        parameters.total_sequence_length,
+        csr_error_buffer.get()));
+  }
+
   size_t rotary_buffer_bytes = 0;
   if (do_rotary_) {
     rotary_buffer_bytes = 2 * sizeof(T) * parameters.batch_size * parameters.num_heads *

onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu

@@ -333,6 +333,120 @@ template Status QkvToContext<BFloat16>(
     contrib::SparseAttentionParameters& parameters,
     SparseAttentionData<BFloat16>& data);
 
+// Validation kernel for CSR sparse layout indices and key sequence lengths.
+// Each block handles one layout (blocks [0, num_layout)) or key lengths (block num_layout).
+// All threads in a warp cooperate via strided iteration over elements.
+// Writes a CSRValidationError code to *error_flag if any check fails.
+__global__ void ValidateCSRIndicesKernel(
+    const int32_t* csr_row_indices,
+    const int32_t* csr_col_indices,
+    const int32_t* seqlens_k_total,
+    int max_blocks,
+    int col_count,
+    int num_layout,
+    int batch_size,
+    int sequence_length,
+    int total_sequence_length,
+    int32_t* error_flag) {
+  int block_id = blockIdx.x;
+  int tid = threadIdx.x;
+  int num_threads = blockDim.x;
+
+  if (block_id < num_layout) {
+    // Validate CSR indices for this layout.
+    const int stride_row = max_blocks + 1;
+    const int32_t* r = csr_row_indices + block_id * stride_row;
+
+    // Phase 1: thread 0 validates all row pointers sequentially.
+    // Row arrays are small (max_blocks+1 elements), so single-thread scan is sufficient.
+    // All threads must reach __syncthreads before proceeding to col validation.
+    __shared__ int row_valid;
+    if (tid == 0) {
+      row_valid = 1;
+      if (r[0] != 0) {
+        atomicCAS(error_flag, kCSRValidationOk, kCSRValidationRowFirstNotZero);
+        row_valid = 0;
+      } else {
+        for (int i = 0; i < max_blocks; ++i) {
+          if (r[i] < 0 || r[i] > r[i + 1] || r[i + 1] > col_count) {
+            atomicCAS(error_flag, kCSRValidationOk, kCSRValidationRowNonMonotonic);
+            row_valid = 0;
+            break;
+          }
+        }
+      }
+    }
+    __syncthreads();
+    if (!row_valid) return;
+
+    // Phase 2: row pointers are validated, r[max_blocks] is safe to use as NNZ bound.
+    // All threads cooperate on the potentially larger col-index array.
+    const int nnz = r[max_blocks];
+    const int32_t* c = csr_col_indices + block_id * col_count;
+    for (int i = tid; i < nnz; i += num_threads) {
+      if (c[i] < 0 || c[i] >= max_blocks) {
+        atomicCAS(error_flag, kCSRValidationOk, kCSRValidationColOutOfRange);
+        return;
+      }
+    }
+  } else if (block_id == num_layout) {
+    // Validate key lengths. All threads cooperate in strided fashion.
+    bool is_prompt = (sequence_length == total_sequence_length);
+    int min_key_length = is_prompt ? 1 : sequence_length;
+    for (int i = tid; i < batch_size; i += num_threads) {
+      int key_length = seqlens_k_total[i];
+      if (key_length < min_key_length || key_length > total_sequence_length) {
+        atomicCAS(error_flag, kCSRValidationOk, kCSRValidationKeyLengthOutOfRange);
+        return;
+      }
+    }
+  }
+}
+
+Status ValidateCSRIndicesOnDevice(
+    cudaStream_t stream,
+    const int32_t* csr_row_indices,
+    const int32_t* csr_col_indices,
+    const int32_t* seqlens_k_total,
+    int num_layout,
+    int max_blocks,
+    int col_count,
+    int batch_size,
+    int sequence_length,
+    int total_sequence_length,
+    int32_t* d_error_flag) {
+  CUDA_RETURN_IF_ERROR(cudaMemsetAsync(d_error_flag, 0, sizeof(int32_t), stream));
+
+  // Launch num_layout blocks for CSR validation + 1 block for key-length validation.
+  // Each block uses a full warp (32 threads) with strided iteration over elements.
+  ValidateCSRIndicesKernel<<<num_layout + 1, 32, 0, stream>>>(
+      csr_row_indices, csr_col_indices, seqlens_k_total,
+      max_blocks, col_count, num_layout,
+      batch_size, sequence_length, total_sequence_length,
+      d_error_flag);
+  CUDA_RETURN_IF_ERROR(cudaGetLastError());
+
+  // Copy error flag back to host.
+  int32_t h_error_flag = 0;
+  CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(&h_error_flag, d_error_flag, sizeof(int32_t),
+                                       cudaMemcpyDeviceToHost, stream));
+  CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream));
+
+  if (h_error_flag != kCSRValidationOk) {
+    const char* msg = (h_error_flag == kCSRValidationRowFirstNotZero)
+                          ? "block_row_indices first element must be 0 for all layouts"
+                      : (h_error_flag == kCSRValidationRowNonMonotonic)
+                          ? "block_row_indices values are not monotonically non-decreasing or exceed "
+                            "block_col_indices columns"
+                      : (h_error_flag == kCSRValidationColOutOfRange)
+                          ? "block_col_indices value is out of valid range"
+                          : "key_total_sequence_lengths value is out of valid range";
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, msg);
+  }
+
+  return Status::OK();
+}
+
 }  // namespace cuda
 }  // namespace contrib
 }  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.h

@@ -68,6 +68,31 @@ Status QkvToContext(
     contrib::SparseAttentionParameters& parameters,
     SparseAttentionData<T>& data);
 
+// Error codes returned by the CSR validation kernel via the device error flag.
+enum CSRValidationError : int32_t {
+  kCSRValidationOk = 0,
+  kCSRValidationRowFirstNotZero = 1,
+  kCSRValidationRowNonMonotonic = 2,
+  kCSRValidationColOutOfRange = 3,
+  kCSRValidationKeyLengthOutOfRange = 4,
+};
+
+// Validate CSR row-pointer monotonicity, column-index range, and key lengths on device.
+// Returns Status::OK() if valid, or INVALID_ARGUMENT with a description of the failure.
+// d_error_flag must point to a device-allocated int32_t scratch buffer (1 element).
+Status ValidateCSRIndicesOnDevice(
+    cudaStream_t stream,
+    const int32_t* csr_row_indices,  // device pointer, shape [num_layout, max_blocks + 1]
+    const int32_t* csr_col_indices,  // device pointer, shape [num_layout, col_count]
+    const int32_t* seqlens_k_total,  // device pointer, shape [batch_size]
+    int num_layout,
+    int max_blocks,
+    int col_count,
+    int batch_size,
+    int sequence_length,
+    int total_sequence_length,
+    int32_t* d_error_flag);  // device scratch buffer (1 int32_t)
+
 }  // namespace cuda
 }  // namespace contrib
 }  // namespace onnxruntime