From 066ef98502b48c65a26b136965577360b8bcaeb2 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 11 Aug 2026 15:58:25 -0700 Subject: [PATCH 1/3] Compute SparseAttention CUDA buffer sizes and offsets in size_t The scratch allocation sizes in SparseAttention::ComputeInternal and the Q/K/V and rotary offsets in QkvToContext were evaluated as products of int shape fields and only widened afterwards, so a large batch_size or sequence_length could wrap the product before it reached GetScratchBuffer or the pointer arithmetic. Use SafeInt so the products are computed at full width and overflow throws instead of wrapping. The Triton kernel parameter structs take 32-bit strides, so also bound the corresponding element counts in CheckInputs to keep those strides representable. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/sparse/sparse_attention_helper.h | 23 +++++++++++++++++++ .../cuda/sparse/sparse_attention.cc | 11 +++++---- .../cuda/sparse/sparse_attention_impl.cu | 7 +++--- 3 files changed, 33 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h index af320b250abdb..8ceb83c3c6bfc 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h @@ -3,6 +3,8 @@ #pragma once +#include + #include "core/common/common.h" #include "core/providers/common.h" #include "contrib_ops/cpu/bert/attention_common.h" @@ -264,6 +266,27 @@ Status CheckInputs(void* params, parameters->stride_row_indices = static_cast(block_row_indices_dim[1]); parameters->stride_col_indices = static_cast(block_col_indices_dim[1]); + // Buffer sizes and kernel strides are computed as products of these dimensions, and the Triton kernels + // take 32-bit strides. Reject shapes whose products do not fit in int32 so that the sizes and offsets + // derived from them cannot wrap. + constexpr int64_t max_int32 = static_cast(std::numeric_limits::max()); + const int64_t q_elements = static_cast(batch_size) * sequence_length * + (num_heads + 2 * static_cast(kv_num_heads)) * head_size; + if (q_elements > max_int32) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "batch_size * sequence_length * (num_heads + 2 * kv_num_heads) * head_size shall not " + "exceed ", + max_int32, ". Got ", q_elements); + } + + const int64_t kv_cache_elements = static_cast(batch_size) * kv_num_heads * + max_cache_sequence_length * head_size; + if (kv_cache_elements > max_int32) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "batch_size * kv_num_heads * max_cache_sequence_length * head_size shall not exceed ", + max_int32, ". Got ", kv_cache_elements); + } + return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc index 355319e84e534..30062ba8b8a30 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc @@ -6,6 +6,7 @@ #include "contrib_ops/cpu/sparse/sparse_attention_helper.h" #include "contrib_ops/cuda/sparse/sparse_attention_v1/sparse_attention_v1_api.h" #include "contrib_ops/cuda/sparse/sparse_attention_v2/sparse_attention_v2_api.h" +#include "core/common/safeint.h" #include "core/platform/env_var_utils.h" #include "contrib_ops/cuda/bert/transformer_cuda_common.h" @@ -238,16 +239,16 @@ Status SparseAttention::ComputeInternal(OpKernelContext* context) const { size_t rotary_buffer_bytes = 0; if (do_rotary_) { - rotary_buffer_bytes = 2 * sizeof(T) * parameters.batch_size * parameters.num_heads * + rotary_buffer_bytes = 2 * sizeof(T) * SafeInt(parameters.batch_size) * parameters.num_heads * parameters.sequence_length * parameters.head_size; - rotary_buffer_bytes += sizeof(int64_t) * parameters.batch_size * parameters.sequence_length; + rotary_buffer_bytes += sizeof(int64_t) * SafeInt(parameters.batch_size) * parameters.sequence_length; } auto rotary_buffer = GetScratchBuffer(rotary_buffer_bytes, GetComputeStream(context)); data.rotary_buffer = reinterpret_cast(rotary_buffer.get()); size_t transposed_q_bytes = 0; if (!parameters.is_packed_qkv) { - transposed_q_bytes = parameters.batch_size * parameters.sequence_length * + transposed_q_bytes = SafeInt(parameters.batch_size) * parameters.sequence_length * parameters.num_heads * parameters.head_size * sizeof(T); } auto transposed_q_buffer = GetScratchBuffer(transposed_q_bytes, GetComputeStream(context)); @@ -257,8 +258,8 @@ Status SparseAttention::ComputeInternal(OpKernelContext* context) const { size_t unpacked_qkv_bytes = 0; if (parameters.is_packed_qkv) { - unpacked_qkv_bytes = (parameters.batch_size * parameters.sequence_length * - (parameters.num_heads + 2 * parameters.kv_num_heads) * + unpacked_qkv_bytes = (SafeInt(parameters.batch_size) * parameters.sequence_length * + (SafeInt(parameters.num_heads) + 2 * parameters.kv_num_heads) * parameters.head_size * sizeof(T)); } auto unpacked_qkv_buffer = GetScratchBuffer(unpacked_qkv_bytes, GetComputeStream(context)); diff --git a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu index 1fecb91b4f578..5501c01bc71e8 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "contrib_ops/cuda/sparse/sparse_attention_impl.h" +#include "core/common/safeint.h" #include "contrib_ops/cuda/utils/dump_cuda_tensor.h" #include "contrib_ops/cuda/bert/rotary_embedding_impl.h" #include "contrib_ops/cuda/bert/group_query_attention_impl.h" @@ -132,8 +133,8 @@ Status QkvToContext( key = reinterpret_cast(data.key); value = reinterpret_cast(data.value); } else { - size_t q_size = static_cast(batch_size * sequence_length * num_heads * head_size); - size_t k_size = static_cast(batch_size * sequence_length * kv_num_heads * head_size); + size_t q_size = SafeInt(batch_size) * sequence_length * num_heads * head_size; + size_t k_size = SafeInt(batch_size) * sequence_length * kv_num_heads * head_size; auto q = reinterpret_cast(data.unpacked_qkv_buffer); auto k = reinterpret_cast(data.unpacked_qkv_buffer + q_size); auto v = reinterpret_cast(data.unpacked_qkv_buffer + q_size + k_size); @@ -165,7 +166,7 @@ Status QkvToContext( #endif if (parameters.do_rotary) { - size_t bsh = static_cast(parameters.batch_size * parameters.sequence_length * parameters.head_size); + size_t bsh = SafeInt(parameters.batch_size) * parameters.sequence_length * parameters.head_size; size_t q_size = bsh * static_cast(parameters.num_heads); size_t k_size = bsh * static_cast(parameters.kv_num_heads); auto q_buffer = reinterpret_cast(data.rotary_buffer); From 74653f72edeff13ea3eaf9dbb8806011b4d67a3e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 11 Aug 2026 22:06:56 -0700 Subject: [PATCH 2/3] Avoid overflow in SparseAttention shape checks --- .../cpu/sparse/sparse_attention_helper.h | 30 ++++++++++++------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h index 8ceb83c3c6bfc..136b5d724a49c 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include "core/common/common.h" @@ -14,6 +15,18 @@ namespace onnxruntime { namespace contrib { namespace sparse_attention_helper { +inline bool ProductExceedsInt32Max(std::initializer_list factors) { + constexpr int64_t max_int32 = static_cast(std::numeric_limits::max()); + int64_t product = 1; + for (int64_t factor : factors) { + if (factor < 0) return true; + if (factor == 0) return false; + if (product > max_int32 / factor) return true; + product *= factor; + } + return false; +} + Status CheckInputs(void* params, const Tensor* query, const Tensor* key, @@ -269,22 +282,17 @@ Status CheckInputs(void* params, // Buffer sizes and kernel strides are computed as products of these dimensions, and the Triton kernels // take 32-bit strides. Reject shapes whose products do not fit in int32 so that the sizes and offsets // derived from them cannot wrap. - constexpr int64_t max_int32 = static_cast(std::numeric_limits::max()); - const int64_t q_elements = static_cast(batch_size) * sequence_length * - (num_heads + 2 * static_cast(kv_num_heads)) * head_size; - if (q_elements > max_int32) { + if (ProductExceedsInt32Max({batch_size, sequence_length, + static_cast(num_heads) + 2 * static_cast(kv_num_heads), head_size})) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "batch_size * sequence_length * (num_heads + 2 * kv_num_heads) * head_size shall not " - "exceed ", - max_int32, ". Got ", q_elements); + "exceed int32 max"); } - const int64_t kv_cache_elements = static_cast(batch_size) * kv_num_heads * - max_cache_sequence_length * head_size; - if (kv_cache_elements > max_int32) { + if (ProductExceedsInt32Max({batch_size, kv_num_heads, max_cache_sequence_length, head_size})) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "batch_size * kv_num_heads * max_cache_sequence_length * head_size shall not exceed ", - max_int32, ". Got ", kv_cache_elements); + "batch_size * kv_num_heads * max_cache_sequence_length * head_size shall not exceed " + "int32 max"); } return Status::OK(); From 8acbf33ab2e6000667d5b1ed616d3bc153a2aa92 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 11 Aug 2026 22:09:18 -0700 Subject: [PATCH 3/3] Simplify SparseAttention product bounds --- .../contrib_ops/cpu/sparse/sparse_attention_helper.h | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h index 136b5d724a49c..5c334d27912ca 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h @@ -4,9 +4,9 @@ #pragma once #include -#include #include "core/common/common.h" +#include "core/common/safeint.h" #include "core/providers/common.h" #include "contrib_ops/cpu/bert/attention_common.h" #include "contrib_ops/cpu/bert/attention_parameters.h" @@ -16,13 +16,11 @@ namespace contrib { namespace sparse_attention_helper { inline bool ProductExceedsInt32Max(std::initializer_list factors) { - constexpr int64_t max_int32 = static_cast(std::numeric_limits::max()); - int64_t product = 1; + int32_t product = 1; for (int64_t factor : factors) { - if (factor < 0) return true; - if (factor == 0) return false; - if (product > max_int32 / factor) return true; - product *= factor; + int32_t next_product; + if (factor < 0 || !SafeMultiply(product, factor, next_product)) return true; + product = next_product; } return false; }