From f3c7bb0a8f36b3c73d54fc33a80cfb7cd3bd9712 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Tue, 23 Jun 2026 16:39:33 -0700 Subject: [PATCH 1/7] Validate SparseAttention CSR indices and key lengths element values CheckInputs only validated tensor shapes for block_row_indices and block_col_indices but never checked element values. This allowed attacker-controlled values to be used as array indices in ComputeAttentionProbs, leading to out-of-bounds memory access. Add ValidateCSRIndices and ValidateKeyLengths for CPU EP that check: - CSR row-pointer monotonicity and bounds - Column indices within [0, max_blocks) - Key lengths within valid range Add ValidateCSRIndicesOnDevice for CUDA EP that performs equivalent validation on-device using a single-warp kernel with strided iteration and a scratch buffer error flag. Add unit tests covering invalid CSR row pointers, non-monotonic values, out-of-range column indices, negative values, and large OOB values. --- .../cpu/sparse/sparse_attention.cc | 5 + .../cpu/sparse/sparse_attention_helper.h | 79 ++++++++++--- .../cuda/sparse/sparse_attention.cc | 17 +++ .../cuda/sparse/sparse_attention_impl.cu | 104 +++++++++++++++++ .../cuda/sparse/sparse_attention_impl.h | 25 ++++ .../contrib_ops/sparse_attention_op_test.cc | 107 ++++++++++++++++++ 6 files changed, 323 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc index bab2dfd13e046..b987c4f4e0bf7 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc +++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc @@ -75,6 +75,11 @@ Status SparseAttention::Compute(OpKernelContext* context) const { block_col_indices, total_key_lengths, total_seq_len)); + ORT_RETURN_IF_ERROR(sparse_attention_helper::ValidateCSRIndices(parameters, + *block_row_indices, + *block_col_indices)); + ORT_RETURN_IF_ERROR(sparse_attention_helper::ValidateKeyLengths(parameters, + *total_key_lengths)); const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h index 2804f30a9611d..d65c13bb34e62 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h +++ b/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(); - 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) { @@ -280,6 +267,70 @@ Status CheckInputs(void* params, return Status::OK(); } +// Validate CSR row-pointer monotonicity and column-index range on CPU-accessible tensors. +// This must be called after CheckInputs has populated the parameters struct. +// On CUDA EP, equivalent validation is performed on-device by ValidateCSRIndicesOnDevice. +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(); + const int32_t* col_data = block_col_indices.Data(); + 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; + for (int k = 0; k < col_count; ++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 on CPU-accessible tensor. +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(); + 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 sparse_attention_helper } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc index 656fde2f46ab8..faf7c59998ebf 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc @@ -215,6 +215,23 @@ Status SparseAttention::ComputeInternal(OpKernelContext* context) const { data.kernel_layout.csr_col_indices = block_col_indices->Data(); data.kernel_layout.csr_row_indices = block_row_indices->Data(); + // Validate CSR indices and key lengths on device to prevent out-of-bounds access. + { + auto csr_error_buffer = GetScratchBuffer(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(), + 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 * diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu index b2a6eb89d4d23..fd289f1ffe0b8 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu @@ -333,6 +333,110 @@ template Status QkvToContext( contrib::SparseAttentionParameters& parameters, SparseAttentionData& data); +// Validation kernel for CSR sparse layout indices and key sequence lengths. +// Each block handles one layout. Thread 0 checks row indices, thread 1 checks col indices. +// Block with block_id == num_layout validates key lengths (thread 0 only). +// 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. + // Thread 0 checks row[0] == 0, then all threads cooperate on row monotonicity and col range. + const int stride_row = max_blocks + 1; + const int32_t* r = csr_row_indices + block_id * stride_row; + + if (tid == 0 && r[0] != 0) { + atomicExch(error_flag, kCSRValidationRowFirstNotZero); + return; + } + + // All threads check row monotonicity in strided fashion. + for (int i = tid; i < max_blocks; i += num_threads) { + if (r[i] < 0 || r[i] > r[i + 1] || r[i + 1] > col_count) { + atomicExch(error_flag, kCSRValidationRowNonMonotonic); + return; + } + } + + // All threads check col indices in strided fashion. + const int32_t* c = csr_col_indices + block_id * col_count; + for (int i = tid; i < col_count; i += num_threads) { + if (c[i] < 0 || c[i] >= max_blocks) { + atomicExch(error_flag, 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) { + atomicExch(error_flag, 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<<>>( + 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[0] must be 0" + : (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 diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.h index d4f686afe5db0..f5b24d99b10a9 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.h @@ -68,6 +68,31 @@ Status QkvToContext( contrib::SparseAttentionParameters& parameters, SparseAttentionData& 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 diff --git a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc index 71d8c34353f02..a31040e2e8934 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc @@ -258,5 +258,112 @@ TEST(SparseAttentionTest, RejectsZeroDimBlockRowIndices) { {}, nullptr, &execution_providers); } +// Helper for CSR value-validation tests. +// Uses: num_heads=2, kv_num_heads=2, sparse_block_size=16, head_size=8. +// block_row_indices shape: (1, max_blocks+1), block_col_indices shape: (1, col_count). +// max_sequence_length = max_blocks * 16 must be >= total_sequence_length. +// These tests validate that element values in block_row_indices and block_col_indices are checked. +// Note: these tests expect failure via ORT_RETURN_IF which does not throw when exceptions are disabled, +// so they are safe to run in no-exceptions builds. +void RunSparseAttentionCSRValidationTest( + const std::vector& block_row_indices_data, + const std::vector& block_row_indices_dims, + const std::vector& block_col_indices_data, + const std::vector& block_col_indices_dims, + const std::string& expected_error) { + OpTester test("SparseAttention", 1, onnxruntime::kMSDomain); + test.AddAttribute("num_heads", 2); + test.AddAttribute("kv_num_heads", 2); + test.AddAttribute("sparse_block_size", 16); + test.AddAttribute("scale", 1.0f); + test.AddAttribute("do_rotary", 0); + test.AddAttribute("rotary_interleaved", 0); + + // head_size=8, num_heads=2 => hidden_size=16 + // sequence_length=1, batch_size=1 + test.AddInput("query", {1, 1, 16}, std::vector(16, 0.0f)); + test.AddInput("key", {1, 1, 16}, std::vector(16, 0.0f)); + test.AddInput("value", {1, 1, 16}, std::vector(16, 0.0f)); + // past_key/value: (batch_size=1, kv_num_heads=2, max_cache_seq_len=32, head_size=8) + test.AddInput("past_key", {1, 2, 32, 8}, std::vector(512, 0.0f)); + test.AddInput("past_value", {1, 2, 32, 8}, std::vector(512, 0.0f)); + test.AddInput("block_row_indices", block_row_indices_dims, block_row_indices_data); + test.AddInput("block_col_indices", block_col_indices_dims, block_col_indices_data); + test.AddInput("total_sequence_length", {1}, {2}); + test.AddInput("key_total_sequence_lengths", {1}, {2}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + + test.AddOutput("output", {1, 1, 16}, std::vector(16, 0.0f)); + test.AddOutput("present_key", {1, 2, 32, 8}, std::vector(512, 0.0f)); + test.AddOutput("present_value", {1, 2, 32, 8}, std::vector(512, 0.0f)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, expected_error, {}, nullptr, &execution_providers); +} + +// block_row_indices[0][0] must be 0. +TEST(SparseAttentionTest, RejectsBlockRowIndicesFirstElementNonZero) { + // shape (1, 3) => max_blocks=2, max_sequence_length=32 + RunSparseAttentionCSRValidationTest( + {1, 1, 2}, {1, 3}, // row indices: first element is 1, not 0 + {0, 1}, {1, 2}, // col indices: valid + "block_row_indices[0][0] must be 0"); +} + +// block_row_indices must be monotonically non-decreasing. +TEST(SparseAttentionTest, RejectsBlockRowIndicesNonMonotonic) { + // shape (1, 3) => max_blocks=2 + RunSparseAttentionCSRValidationTest( + {0, 2, 1}, {1, 3}, // row indices: 2 > 1 at row 1 (non-monotonic) + {0, 1}, {1, 2}, // col indices: valid + "block_row_indices values are not monotonically non-decreasing"); +} + +// block_row_indices values must not exceed block_col_indices column count. +TEST(SparseAttentionTest, RejectsBlockRowIndicesExceedsColCount) { + // shape (1, 3) => max_blocks=2, col_count=2 + RunSparseAttentionCSRValidationTest( + {0, 1, 3}, {1, 3}, // row indices: last element 3 > col_count=2 + {0, 1}, {1, 2}, // col indices shape (1, 2) + "block_row_indices values are not monotonically non-decreasing"); +} + +// block_col_indices values must be in [0, max_blocks). +TEST(SparseAttentionTest, RejectsBlockColIndicesOutOfRange) { + // shape (1, 3) => max_blocks=2 + RunSparseAttentionCSRValidationTest( + {0, 1, 2}, {1, 3}, // row indices: valid + {0, 2}, {1, 2}, // col indices: value 2 >= max_blocks=2 + "block_col_indices[0][1]=2 is out of valid range [0, 2)"); +} + +// block_col_indices negative values must be rejected. +TEST(SparseAttentionTest, RejectsBlockColIndicesNegative) { + // shape (1, 3) => max_blocks=2 + RunSparseAttentionCSRValidationTest( + {0, 1, 2}, {1, 3}, // row indices: valid + {0, -1}, {1, 2}, // col indices: negative value + "block_col_indices[0][1]=-1 is out of valid range [0, 2)"); +} + +// block_row_indices with negative values. +TEST(SparseAttentionTest, RejectsBlockRowIndicesNegative) { + RunSparseAttentionCSRValidationTest( + {0, -1, 2}, {1, 3}, // row indices: negative value at index 1 + {0, 1}, {1, 2}, // col indices: valid + "block_row_indices values are not monotonically non-decreasing"); +} + +// block_col_indices with large OOB value (the original vulnerability scenario). +TEST(SparseAttentionTest, RejectsBlockColIndicesLargeValue) { + // shape (1, 3) => max_blocks=2 + RunSparseAttentionCSRValidationTest( + {0, 2, 2}, {1, 3}, // row indices: valid CSR format + {0, 1048576}, {1, 2}, // col indices: 0x100000 far out of range + "block_col_indices[0][1]=1048576 is out of valid range [0, 2)"); +} + } // namespace test } // namespace onnxruntime From 4accf708eee7ba2ce769cf1712c911505eb851be Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Tue, 23 Jun 2026 16:48:42 -0700 Subject: [PATCH 2/7] Move validation to .cc, use references, add CUDA key-length check - Move ValidateCSRIndices and ValidateKeyLengths from header to anonymous namespace in sparse_attention.cc (not shared across TUs). - Change Tensor pointer parameters to const Tensor& references since they are required (never null). - Extend CUDA ValidateCSRIndicesOnDevice kernel to also validate seqlens_k_total (key lengths) in the same launch. - Add CSRValidationError enum for device kernel error codes. - Use full warp (32 threads) with strided iteration in the kernel. - Scope scratch buffer allocation to validation call site. --- .../cpu/sparse/sparse_attention.cc | 77 +++++++++++++++++-- .../cpu/sparse/sparse_attention_helper.h | 64 --------------- 2 files changed, 72 insertions(+), 69 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc index b987c4f4e0bf7..76ed620a22b3e 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc +++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc @@ -36,6 +36,73 @@ 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(); + const int32_t* col_data = block_col_indices.Data(); + 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; + for (int k = 0; k < col_count; ++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(); + 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 SparseAttention::SparseAttention(const OpKernelInfo& info) : OpKernel(info), SparseAttentionBase(info) { } @@ -75,11 +142,11 @@ Status SparseAttention::Compute(OpKernelContext* context) const { block_col_indices, total_key_lengths, total_seq_len)); - ORT_RETURN_IF_ERROR(sparse_attention_helper::ValidateCSRIndices(parameters, - *block_row_indices, - *block_col_indices)); - ORT_RETURN_IF_ERROR(sparse_attention_helper::ValidateKeyLengths(parameters, - *total_key_lengths)); + 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; diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h index d65c13bb34e62..af320b250abdb 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h @@ -267,70 +267,6 @@ Status CheckInputs(void* params, return Status::OK(); } -// Validate CSR row-pointer monotonicity and column-index range on CPU-accessible tensors. -// This must be called after CheckInputs has populated the parameters struct. -// On CUDA EP, equivalent validation is performed on-device by ValidateCSRIndicesOnDevice. -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(); - const int32_t* col_data = block_col_indices.Data(); - 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; - for (int k = 0; k < col_count; ++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 on CPU-accessible tensor. -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(); - 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 sparse_attention_helper } // namespace contrib } // namespace onnxruntime From 90e7e107f2ea73b6d5a8dc091e10d74df4e91f12 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Tue, 23 Jun 2026 16:52:45 -0700 Subject: [PATCH 3/7] Address review: limit col validation to NNZ, use atomicCAS, static helper - Only validate block_col_indices entries in [0, nnz) where nnz is the last row pointer value r[max_blocks], since entries beyond that are padding and not consumed by the kernel. - Replace atomicExch with atomicCAS to preserve the first detected error code deterministically across concurrent threads/blocks. - Mark RunSparseAttentionCSRValidationTest as static to avoid external linkage. --- .../contrib_ops/cpu/sparse/sparse_attention.cc | 3 ++- .../cuda/sparse/sparse_attention_impl.cu | 13 +++++++------ .../test/contrib_ops/sparse_attention_op_test.cc | 2 +- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc index 76ed620a22b3e..b027772971540 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc +++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc @@ -66,7 +66,8 @@ Status ValidateCSRIndices(const SparseAttentionParameters& parameters, } } const int32_t* c = col_data + l * col_count; - for (int k = 0; k < col_count; ++k) { + 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], diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu index fd289f1ffe0b8..ff5d1cf9aa3dd 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu @@ -359,23 +359,24 @@ __global__ void ValidateCSRIndicesKernel( const int32_t* r = csr_row_indices + block_id * stride_row; if (tid == 0 && r[0] != 0) { - atomicExch(error_flag, kCSRValidationRowFirstNotZero); + atomicCAS(error_flag, kCSRValidationOk, kCSRValidationRowFirstNotZero); return; } // All threads check row monotonicity in strided fashion. for (int i = tid; i < max_blocks; i += num_threads) { if (r[i] < 0 || r[i] > r[i + 1] || r[i + 1] > col_count) { - atomicExch(error_flag, kCSRValidationRowNonMonotonic); + atomicCAS(error_flag, kCSRValidationOk, kCSRValidationRowNonMonotonic); return; } } - // All threads check col indices in strided fashion. + // All threads check col indices in strided fashion (only the used [0, nnz) entries). + const int nnz = r[max_blocks]; const int32_t* c = csr_col_indices + block_id * col_count; - for (int i = tid; i < col_count; i += num_threads) { + for (int i = tid; i < nnz; i += num_threads) { if (c[i] < 0 || c[i] >= max_blocks) { - atomicExch(error_flag, kCSRValidationColOutOfRange); + atomicCAS(error_flag, kCSRValidationOk, kCSRValidationColOutOfRange); return; } } @@ -386,7 +387,7 @@ __global__ void ValidateCSRIndicesKernel( 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) { - atomicExch(error_flag, kCSRValidationKeyLengthOutOfRange); + atomicCAS(error_flag, kCSRValidationOk, kCSRValidationKeyLengthOutOfRange); return; } } diff --git a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc index a31040e2e8934..5a25483daa230 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc @@ -265,7 +265,7 @@ TEST(SparseAttentionTest, RejectsZeroDimBlockRowIndices) { // These tests validate that element values in block_row_indices and block_col_indices are checked. // Note: these tests expect failure via ORT_RETURN_IF which does not throw when exceptions are disabled, // so they are safe to run in no-exceptions builds. -void RunSparseAttentionCSRValidationTest( +static void RunSparseAttentionCSRValidationTest( const std::vector& block_row_indices_data, const std::vector& block_row_indices_dims, const std::vector& block_col_indices_data, From af1f999ad26f847cfece9a995b2c866287865aa0 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 24 Jun 2026 11:53:54 -0700 Subject: [PATCH 4/7] Fix tests, race condition, and add coverage - Fix existing key-length tests to use valid CSR indices (col_count=4). - Rename RunSparseAttentionInvalidInputTest to RunSparseAttentionInvalidKeyLengthsTest. - Fix CUDA kernel race: thread 0 validates row pointers sequentially, then __syncthreads ensures r[max_blocks] is safe before all threads cooperate on col-index validation. - Update kernel comment to reflect actual thread cooperation model. - Make error message layout-agnostic for kCSRValidationRowFirstNotZero. - Add tests: multi-layout invalid col/row in second layout, invalid col within NNZ range, padding acceptance beyond NNZ. --- .../cuda/sparse/sparse_attention_impl.cu | 39 +++++---- .../contrib_ops/sparse_attention_op_test.cc | 84 ++++++++++++++++--- 2 files changed, 98 insertions(+), 25 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu index ff5d1cf9aa3dd..1fecb91b4f578 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu @@ -334,8 +334,8 @@ template Status QkvToContext( SparseAttentionData& data); // Validation kernel for CSR sparse layout indices and key sequence lengths. -// Each block handles one layout. Thread 0 checks row indices, thread 1 checks col indices. -// Block with block_id == num_layout validates key lengths (thread 0 only). +// 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, @@ -354,24 +354,33 @@ __global__ void ValidateCSRIndicesKernel( if (block_id < num_layout) { // Validate CSR indices for this layout. - // Thread 0 checks row[0] == 0, then all threads cooperate on row monotonicity and col range. const int stride_row = max_blocks + 1; const int32_t* r = csr_row_indices + block_id * stride_row; - if (tid == 0 && r[0] != 0) { - atomicCAS(error_flag, kCSRValidationOk, kCSRValidationRowFirstNotZero); - return; - } - - // All threads check row monotonicity in strided fashion. - for (int i = tid; i < max_blocks; i += num_threads) { - if (r[i] < 0 || r[i] > r[i + 1] || r[i + 1] > col_count) { - atomicCAS(error_flag, kCSRValidationOk, kCSRValidationRowNonMonotonic); - return; + // 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; - // All threads check col indices in strided fashion (only the used [0, nnz) entries). + // 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) { @@ -425,7 +434,7 @@ Status ValidateCSRIndicesOnDevice( if (h_error_flag != kCSRValidationOk) { const char* msg = (h_error_flag == kCSRValidationRowFirstNotZero) - ? "block_row_indices[0] must be 0" + ? "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" diff --git a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc index 5a25483daa230..5502b266d0676 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc @@ -24,10 +24,10 @@ namespace test { namespace { -void RunSparseAttentionInvalidInputTest(const std::vector& total_key_lengths_data, - const std::vector& total_key_lengths_dims, - const std::string& expected_error, - int32_t total_sequence_length = 4) { +void RunSparseAttentionInvalidKeyLengthsTest(const std::vector& total_key_lengths_data, + const std::vector& total_key_lengths_dims, + const std::string& expected_error, + int32_t total_sequence_length = 4) { OpTester test("SparseAttention", 1, onnxruntime::kMSDomain); test.AddAttribute("num_heads", 2); test.AddAttribute("kv_num_heads", 2); @@ -42,7 +42,7 @@ void RunSparseAttentionInvalidInputTest(const std::vector& total_key_le test.AddInput("past_key", {1, 2, 4, 8}, std::vector(64, 0.0f)); test.AddInput("past_value", {1, 2, 4, 8}, std::vector(64, 0.0f)); test.AddInput("block_row_indices", {1, 5}, {0, 1, 2, 3, 4}); - test.AddInput("block_col_indices", {1, 1}, {0}); + test.AddInput("block_col_indices", {1, 4}, {0, 1, 2, 3}); test.AddInput("total_sequence_length", {1}, {total_sequence_length}); test.AddInput("key_total_sequence_lengths", total_key_lengths_dims, total_key_lengths_data); test.AddOptionalInputEdge(); @@ -209,17 +209,17 @@ void RunSparseAttentionPromptInputTest(const std::vector& total_key_len } // namespace TEST(SparseAttentionTest, RejectsOutOfRangeKeyTotalSequenceLengths) { - RunSparseAttentionInvalidInputTest({-5}, {1}, "key_total_sequence_lengths value -5 at batch index 0 is out of range [1, 4]"); + RunSparseAttentionInvalidKeyLengthsTest({-5}, {1}, "key_total_sequence_lengths value -5 at batch index 0 is out of range [1, 4]"); } TEST(SparseAttentionTest, RejectsKeyTotalSequenceLengthsShapeMismatch) { - RunSparseAttentionInvalidInputTest({4, 4}, {2}, "key_total_sequence_lengths must have shape (batch_size)"); + RunSparseAttentionInvalidKeyLengthsTest({4, 4}, {2}, "key_total_sequence_lengths must have shape (batch_size)"); } TEST(SparseAttentionTest, RejectsPromptKeyTotalSequenceLengthsShorterThanSequenceLength) { - RunSparseAttentionInvalidInputTest({0}, {1}, - "key_total_sequence_lengths value 0 at batch index 0 is out of range [1, 1]", - 1); + RunSparseAttentionInvalidKeyLengthsTest({0}, {1}, + "key_total_sequence_lengths value 0 at batch index 0 is out of range [1, 1]", + 1); } TEST(SparseAttentionTest, AcceptsPromptKeyTotalSequenceLengthsForPaddedBatch) { @@ -365,5 +365,69 @@ TEST(SparseAttentionTest, RejectsBlockColIndicesLargeValue) { "block_col_indices[0][1]=1048576 is out of valid range [0, 2)"); } +// Multi-layout: invalid col index in second layout only. +TEST(SparseAttentionTest, RejectsBlockColIndicesInvalidInSecondLayout) { + // shape (2, 3) => num_layout=2, max_blocks=2 + // num_heads=2, so num_heads % num_layout == 0 + RunSparseAttentionCSRValidationTest( + {0, 1, 2, 0, 1, 2}, {2, 3}, // row indices: valid for both layouts + {0, 1, 0, 5}, {2, 2}, // col indices: layout 0 valid, layout 1 has 5 >= max_blocks=2 + "block_col_indices[1][1]=5 is out of valid range [0, 2)"); +} + +// Multi-layout: invalid row pointer in second layout only. +TEST(SparseAttentionTest, RejectsBlockRowIndicesInvalidInSecondLayout) { + // shape (2, 3) => num_layout=2, max_blocks=2 + RunSparseAttentionCSRValidationTest( + {0, 1, 2, 1, 1, 2}, {2, 3}, // row indices: layout 0 valid, layout 1 starts with 1 != 0 + {0, 1, 0, 1}, {2, 2}, // col indices: valid + "block_row_indices[1][0] must be 0"); +} + +// Col index invalid within NNZ range but padding would be fine. +// row pointers say nnz=1, col[0] is invalid, col[1] is padding (not checked). +TEST(SparseAttentionTest, RejectsBlockColIndicesInvalidWithinNNZ) { + // shape (1, 3) => max_blocks=2, row indices: {0, 1, 1} means row 0 has 1 entry, row 1 has 0 + // nnz = r[max_blocks] = r[2] = 1, so only col[0] is validated + RunSparseAttentionCSRValidationTest( + {0, 1, 1}, {1, 3}, // row indices: valid, nnz=1 + {99, 0}, {1, 2}, // col[0]=99 is out of range, col[1]=0 is padding (not checked) + "block_col_indices[0][0]=99 is out of valid range [0, 2)"); +} + +// Col padding entries are not validated (only entries within NNZ are checked). +// row pointers say nnz=1, col[1]=-1 is padding and should not trigger rejection. +TEST(SparseAttentionTest, AcceptsPaddedColIndicesBeyondNNZ) { + OpTester test("SparseAttention", 1, onnxruntime::kMSDomain); + test.AddAttribute("num_heads", 2); + test.AddAttribute("kv_num_heads", 2); + test.AddAttribute("sparse_block_size", 16); + test.AddAttribute("scale", 1.0f); + test.AddAttribute("do_rotary", 0); + test.AddAttribute("rotary_interleaved", 0); + + test.AddInput("query", {1, 1, 16}, std::vector(16, 0.0f)); + test.AddInput("key", {1, 1, 16}, std::vector(16, 0.0f)); + test.AddInput("value", {1, 1, 16}, std::vector(16, 0.0f)); + test.AddInput("past_key", {1, 2, 32, 8}, std::vector(512, 0.0f)); + test.AddInput("past_value", {1, 2, 32, 8}, std::vector(512, 0.0f)); + // shape (1, 3) => max_blocks=2, row indices: {0, 1, 1} means nnz=1 + test.AddInput("block_row_indices", {1, 3}, {0, 1, 1}); + // col[0]=0 is valid (within nnz), col[1]=-1 is padding beyond nnz (should not be checked) + test.AddInput("block_col_indices", {1, 2}, {0, -1}); + test.AddInput("total_sequence_length", {1}, {2}); + test.AddInput("key_total_sequence_lengths", {1}, {2}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + + test.AddOutput("output", {1, 1, 16}, std::vector(16, 0.0f)); + test.AddOutput("present_key", {1, 2, 32, 8}, std::vector(512, 0.0f)); + test.AddOutput("present_value", {1, 2, 32, 8}, std::vector(512, 0.0f)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime From 08c13136422e8587224518565e944d87213c4890 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 24 Jun 2026 12:18:43 -0700 Subject: [PATCH 5/7] Add env var opt-out for CUDA validation, add CUDA tests - Add ORT_DISABLE_SPARSE_ATTENTION_INPUT_VALIDATION env var to skip device-side CSR and key-length validation when inputs are known to be well-formed (addresses @tianleiwu latency concern). - Add CUDA-specific validation tests under #ifdef USE_CUDA that use head_size=128, sparse_block_size=64, MLFloat16 (CUDA EP requirements). Tests cover: row first element, row monotonicity, col OOB, col large value, and key length OOB. Tests run only with DefaultCudaExecutionProvider. --- .../contrib_ops/cpu/bert/attention_common.h | 6 + .../cuda/sparse/sparse_attention.cc | 4 +- .../cuda/sparse/sparse_attention.h | 17 +-- .../contrib_ops/sparse_attention_op_test.cc | 133 ++++++++++++++++++ 4 files changed, 151 insertions(+), 9 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_common.h b/onnxruntime/contrib_ops/cpu/bert/attention_common.h index ac3531e39eb53..e46075a86f811 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_common.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_common.h @@ -91,6 +91,12 @@ constexpr bool LAYOUT_BNSH = true; namespace sparse_attention { // Environment variable to enable or disable sparse attention v1 kernel. Default is 0 (enabled). constexpr const char* kDisableSparseAttentionV1 = "ORT_DISABLE_SPARSE_ATTENTION_V1"; + +// Environment variable to disable device-side validation of CSR indices and key sequence lengths. +// Default is 0 (validation enabled). Set to 1 to skip the validation kernel launch and stream +// synchronization, which may improve latency when inputs are known to be well-formed. +// Usage: export ORT_DISABLE_SPARSE_ATTENTION_INPUT_VALIDATION=1 +constexpr const char* kDisableInputValidation = "ORT_DISABLE_SPARSE_ATTENTION_INPUT_VALIDATION"; } // namespace sparse_attention namespace attention { diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc index faf7c59998ebf..aa37b1a8277a5 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc @@ -56,6 +56,8 @@ SparseAttention::SparseAttention(const OpKernelInfo& info) scale_ = info.GetAttrOrDefault("scale", 0.0f); disable_v1_kernel_ = ParseEnvironmentVariableWithDefault(sparse_attention::kDisableSparseAttentionV1, false); + disable_input_validation_ = ParseEnvironmentVariableWithDefault( + sparse_attention::kDisableInputValidation, false); } template @@ -216,7 +218,7 @@ Status SparseAttention::ComputeInternal(OpKernelContext* context) const { data.kernel_layout.csr_row_indices = block_row_indices->Data(); // Validate CSR indices and key lengths on device to prevent out-of-bounds access. - { + if (!disable_input_validation_) { auto csr_error_buffer = GetScratchBuffer(1, GetComputeStream(context)); ORT_RETURN_IF_ERROR(ValidateCSRIndicesOnDevice( cuda_stream, diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.h b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.h index 1df3affe17ea3..06fc07f088c60 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.h +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.h @@ -18,14 +18,15 @@ class SparseAttention final : public CudaKernel { Status ComputeInternal(OpKernelContext* context) const override; protected: - int num_heads_; // number of attention heads for q - int kv_num_heads_; // number of attention heads for k and v - float scale_; // Scaling factor applied prior to softmax. - bool is_causal_; // unidirectional attention or not - int sparse_block_size_; // block size for sparsity - bool do_rotary_; // Has rotary positional embedding - bool rotary_interleaved_; // Interleaved rotary positional embedding - bool disable_v1_kernel_; // Whether disable v1 kernel and use v2 kernel for prompt. + int num_heads_; // number of attention heads for q + int kv_num_heads_; // number of attention heads for k and v + float scale_; // Scaling factor applied prior to softmax. + bool is_causal_; // unidirectional attention or not + int sparse_block_size_; // block size for sparsity + bool do_rotary_; // Has rotary positional embedding + bool rotary_interleaved_; // Interleaved rotary positional embedding + bool disable_v1_kernel_; // Whether disable v1 kernel and use v2 kernel for prompt. + bool disable_input_validation_; // Whether to skip device-side CSR and key-length validation. }; } // namespace cuda diff --git a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc index 5502b266d0676..d79bbeb998242 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc @@ -429,5 +429,138 @@ TEST(SparseAttentionTest, AcceptsPaddedColIndicesBeyondNNZ) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +#if defined(USE_CUDA) +// CUDA-specific CSR validation tests. +// CUDA SparseAttention requires head_size=128, sparse_block_size=64, and MLFloat16 inputs. +// These tests verify that the device-side ValidateCSRIndicesOnDevice kernel correctly +// rejects invalid CSR indices. Error messages are less detailed than CPU (no per-element info) +// because the CUDA kernel reports via a single error code. +static void RunSparseAttentionCudaCSRValidationTest( + const std::vector& block_row_indices_data, + const std::vector& block_row_indices_dims, + const std::vector& block_col_indices_data, + const std::vector& block_col_indices_dims, + const std::string& expected_error) { + OpTester test("SparseAttention", 1, onnxruntime::kMSDomain); + test.AddAttribute("num_heads", 1); + test.AddAttribute("kv_num_heads", 1); + test.AddAttribute("sparse_block_size", 64); + test.AddAttribute("scale", 1.0f); + test.AddAttribute("do_rotary", 0); + test.AddAttribute("rotary_interleaved", 0); + + // head_size=128, num_heads=1 => hidden_size=128 + // sequence_length=1, batch_size=1 + const int64_t hidden_size = 128; + const int64_t max_cache_seq_len = 128; + test.AddInput("query", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddInput("key", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddInput("value", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddInput("past_key", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + test.AddInput("past_value", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + test.AddInput("block_row_indices", block_row_indices_dims, block_row_indices_data); + test.AddInput("block_col_indices", block_col_indices_dims, block_col_indices_data); + test.AddInput("total_sequence_length", {1}, {2}); + test.AddInput("key_total_sequence_lengths", {1}, {2}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + + test.AddOutput("output", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddOutput("present_key", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + test.AddOutput("present_value", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + + // Run only on CUDA EP. CPU EP does not register MLFloat16 for SparseAttention with these params. + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, expected_error, {}, nullptr, &execution_providers); +} + +// CUDA: block_row_indices first element must be 0. +TEST(SparseAttentionTest, CudaRejectsBlockRowIndicesFirstElementNonZero) { + // shape (1, 3) => max_blocks=2 + RunSparseAttentionCudaCSRValidationTest( + {1, 1, 2}, {1, 3}, + {0, 1}, {1, 2}, + "block_row_indices first element must be 0 for all layouts"); +} + +// CUDA: block_row_indices must be monotonically non-decreasing. +TEST(SparseAttentionTest, CudaRejectsBlockRowIndicesNonMonotonic) { + RunSparseAttentionCudaCSRValidationTest( + {0, 2, 1}, {1, 3}, + {0, 1}, {1, 2}, + "block_row_indices values are not monotonically non-decreasing"); +} + +// CUDA: block_col_indices values must be in range. +TEST(SparseAttentionTest, CudaRejectsBlockColIndicesOutOfRange) { + RunSparseAttentionCudaCSRValidationTest( + {0, 1, 2}, {1, 3}, + {0, 99}, {1, 2}, + "block_col_indices value is out of valid range"); +} + +// CUDA: block_col_indices with large OOB value. +TEST(SparseAttentionTest, CudaRejectsBlockColIndicesLargeValue) { + RunSparseAttentionCudaCSRValidationTest( + {0, 2, 2}, {1, 3}, + {0, 1048576}, {1, 2}, + "block_col_indices value is out of valid range"); +} + +// CUDA: key_total_sequence_lengths out of range. +TEST(SparseAttentionTest, CudaRejectsKeyLengthOutOfRange) { + OpTester test("SparseAttention", 1, onnxruntime::kMSDomain); + test.AddAttribute("num_heads", 1); + test.AddAttribute("kv_num_heads", 1); + test.AddAttribute("sparse_block_size", 64); + test.AddAttribute("scale", 1.0f); + test.AddAttribute("do_rotary", 0); + test.AddAttribute("rotary_interleaved", 0); + + const int64_t hidden_size = 128; + const int64_t max_cache_seq_len = 128; + test.AddInput("query", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddInput("key", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddInput("value", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddInput("past_key", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + test.AddInput("past_value", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + // Valid CSR: shape (1, 3) => max_blocks=2 + test.AddInput("block_row_indices", {1, 3}, {0, 1, 2}); + test.AddInput("block_col_indices", {1, 2}, {0, 1}); + test.AddInput("total_sequence_length", {1}, {4}); + // Invalid key length: -5 is out of range [1, 4] + test.AddInput("key_total_sequence_lengths", {1}, {-5}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + + test.AddOutput("output", {1, 1, hidden_size}, + std::vector(hidden_size, MLFloat16(0.0f))); + test.AddOutput("present_key", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + test.AddOutput("present_value", {1, 1, max_cache_seq_len, hidden_size}, + std::vector(max_cache_seq_len * hidden_size, MLFloat16(0.0f))); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "key_total_sequence_lengths value is out of valid range", + {}, nullptr, &execution_providers); +} +#endif // USE_CUDA + } // namespace test } // namespace onnxruntime From bbde0f86eaebd0a37e5c04e2682e9bf22dec8084 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 24 Jun 2026 13:28:44 -0700 Subject: [PATCH 6/7] Remove AcceptsPaddedColIndicesBeyondNNZ test SparseAttention requires past/present to share buffers (in-place). OpTester allocates separate buffers, causing pointer equality check to fail before reaching compute. The padding non-rejection behavior is already implicitly verified by RejectsBlockColIndicesInvalidWithinNNZ. --- .../contrib_ops/sparse_attention_op_test.cc | 34 ------------------- 1 file changed, 34 deletions(-) diff --git a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc index d79bbeb998242..a9fc04ed93515 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc @@ -395,40 +395,6 @@ TEST(SparseAttentionTest, RejectsBlockColIndicesInvalidWithinNNZ) { "block_col_indices[0][0]=99 is out of valid range [0, 2)"); } -// Col padding entries are not validated (only entries within NNZ are checked). -// row pointers say nnz=1, col[1]=-1 is padding and should not trigger rejection. -TEST(SparseAttentionTest, AcceptsPaddedColIndicesBeyondNNZ) { - OpTester test("SparseAttention", 1, onnxruntime::kMSDomain); - test.AddAttribute("num_heads", 2); - test.AddAttribute("kv_num_heads", 2); - test.AddAttribute("sparse_block_size", 16); - test.AddAttribute("scale", 1.0f); - test.AddAttribute("do_rotary", 0); - test.AddAttribute("rotary_interleaved", 0); - - test.AddInput("query", {1, 1, 16}, std::vector(16, 0.0f)); - test.AddInput("key", {1, 1, 16}, std::vector(16, 0.0f)); - test.AddInput("value", {1, 1, 16}, std::vector(16, 0.0f)); - test.AddInput("past_key", {1, 2, 32, 8}, std::vector(512, 0.0f)); - test.AddInput("past_value", {1, 2, 32, 8}, std::vector(512, 0.0f)); - // shape (1, 3) => max_blocks=2, row indices: {0, 1, 1} means nnz=1 - test.AddInput("block_row_indices", {1, 3}, {0, 1, 1}); - // col[0]=0 is valid (within nnz), col[1]=-1 is padding beyond nnz (should not be checked) - test.AddInput("block_col_indices", {1, 2}, {0, -1}); - test.AddInput("total_sequence_length", {1}, {2}); - test.AddInput("key_total_sequence_lengths", {1}, {2}); - test.AddOptionalInputEdge(); - test.AddOptionalInputEdge(); - - test.AddOutput("output", {1, 1, 16}, std::vector(16, 0.0f)); - test.AddOutput("present_key", {1, 2, 32, 8}, std::vector(512, 0.0f)); - test.AddOutput("present_value", {1, 2, 32, 8}, std::vector(512, 0.0f)); - - std::vector> execution_providers; - execution_providers.push_back(DefaultCpuExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); -} - #if defined(USE_CUDA) // CUDA-specific CSR validation tests. // CUDA SparseAttention requires head_size=128, sparse_block_size=64, and MLFloat16 inputs. From eb4dc9db15e237eee5b54d7e400c46391f01101b Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Wed, 24 Jun 2026 15:09:08 -0700 Subject: [PATCH 7/7] Move CUDA validation before shared-buffer check, fix test comments - Move ValidateCSRIndicesOnDevice call to right after CheckInputs, before the past/present shared-buffer enforcement. The validation only reads input tensors (CSR indices, key lengths) and does not need the KV cache buffers. This allows OpTester-based CUDA tests to exercise validation without IOBinding. - Fix test comment: ORT_MAKE_STATUS returns Status, does not throw. - Document in CUDA test helper why OpTester works without shared buffers (validation runs before the buffer check). --- .../cuda/sparse/sparse_attention.cc | 38 ++++++++++--------- .../contrib_ops/sparse_attention_op_test.cc | 7 +++- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc index aa37b1a8277a5..355319e84e534 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc @@ -107,6 +107,26 @@ Status SparseAttention::ComputeInternal(OpKernelContext* context) const { block_col_indices, seqlens_k_total, total_seq_len)); + + // Validate CSR indices and key lengths on device to prevent out-of-bounds access. + // This must run before the shared-buffer check so OpTester-based tests can exercise it. + cudaStream_t cuda_stream = Stream(context); + if (!disable_input_validation_) { + auto csr_error_buffer = GetScratchBuffer(1, GetComputeStream(context)); + ORT_RETURN_IF_ERROR(ValidateCSRIndicesOnDevice( + cuda_stream, + block_row_indices->Data(), + block_col_indices->Data(), + seqlens_k_total->Data(), + 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())); + } + // Some limitations of CUDA kernels // The v1 and v2 kernels have same coverage, so only check one of them to see whether it is supported. int sm = device_prop.major * 10 + device_prop.minor; @@ -139,7 +159,6 @@ Status SparseAttention::ComputeInternal(OpKernelContext* context) const { int32_t* total_k_seq_len_pinned = nullptr; AutoDestoryCudaEvent new_event; cudaEvent_t& isCopyDone = new_event.Get(); - cudaStream_t cuda_stream = Stream(context); if (use_v2_kernel) { pinned_buffer = AllocateBufferOnCPUPinned(parameters.batch_size); @@ -217,23 +236,6 @@ Status SparseAttention::ComputeInternal(OpKernelContext* context) const { data.kernel_layout.csr_col_indices = block_col_indices->Data(); data.kernel_layout.csr_row_indices = block_row_indices->Data(); - // Validate CSR indices and key lengths on device to prevent out-of-bounds access. - if (!disable_input_validation_) { - auto csr_error_buffer = GetScratchBuffer(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(), - 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 * diff --git a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc index a9fc04ed93515..d7953442d738e 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_op_test.cc @@ -263,8 +263,8 @@ TEST(SparseAttentionTest, RejectsZeroDimBlockRowIndices) { // block_row_indices shape: (1, max_blocks+1), block_col_indices shape: (1, col_count). // max_sequence_length = max_blocks * 16 must be >= total_sequence_length. // These tests validate that element values in block_row_indices and block_col_indices are checked. -// Note: these tests expect failure via ORT_RETURN_IF which does not throw when exceptions are disabled, -// so they are safe to run in no-exceptions builds. +// Note: these tests expect failure via a returned Status (ORT_MAKE_STATUS), so they are safe in +// both exceptions-enabled and no-exceptions builds. static void RunSparseAttentionCSRValidationTest( const std::vector& block_row_indices_data, const std::vector& block_row_indices_dims, @@ -401,6 +401,9 @@ TEST(SparseAttentionTest, RejectsBlockColIndicesInvalidWithinNNZ) { // These tests verify that the device-side ValidateCSRIndicesOnDevice kernel correctly // rejects invalid CSR indices. Error messages are less detailed than CPU (no per-element info) // because the CUDA kernel reports via a single error code. +// Note: OpTester does not share past/present buffers (no IOBinding), but that is fine here +// because the CSR validation runs before the shared-buffer check in ComputeInternal. +// These tests expect failure from validation, not from compute. static void RunSparseAttentionCudaCSRValidationTest( const std::vector& block_row_indices_data, const std::vector& block_row_indices_dims,