From 4892e091d10b1a8eaf300749887d24932ced90d7 Mon Sep 17 00:00:00 2001 From: vraspar Date: Fri, 20 Mar 2026 22:36:59 +0000 Subject: [PATCH 1/3] Add validation for non-negative end_position in Attention mask_index --- .../contrib_ops/cpu/bert/attention_base.h | 11 ++++ .../contrib_ops/cpu/bert/attention_helper.h | 4 +- .../cuda/bert/attention_softmax.cu | 4 +- .../test/contrib_ops/attention_op_test.cc | 50 +++++++++++++++++++ 4 files changed, 65 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_base.h index 2872fcfda5bbf..9477fbf765808 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_base.h @@ -121,6 +121,17 @@ inline Status AttentionBase::CheckMask(const Tensor* mask_index, } mask_type = (mask_dims[0] == batch_size ? AttentionMaskType::MASK_1D_KEY_SEQ_LEN : mask_dims[0] == 2 * batch_size ? AttentionMaskType::MASK_1D_END_START : AttentionMaskType::MASK_1D_KEY_SEQ_LEN_START); + + // Validate that end_position values (first batch_size elements) are non-negative. + // Negative end_position causes out-of-bounds writes in PrepareMask. + const int32_t* mask_data = mask_index->Data(); + for (int64_t i = 0; i < batch_size; i++) { + if (mask_data[i] < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "mask_index value ", mask_data[i], " at index ", i, + " is negative. mask_index end_position values must be non-negative."); + } + } } else if (mask_dims.size() == 2) { if (mask_dims[0] == batch_size && mask_dims[1] == total_sequence_length) { mask_type = AttentionMaskType::MASK_2D_KEY_PADDING; diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h index aef47edd5fcd2..73ed43b64f49e 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h @@ -95,14 +95,14 @@ void PrepareMask(const int32_t* mask_index, // mask_index is 1D: (B) or (2B) => (Bx)T // Handle right-side padding: mask value at or after the end position will be mask_filter_value - int end_position = mask_index[b_i]; + int end_position = std::max(0, std::min(static_cast(mask_index[b_i]), all_sequence_length)); for (int m_i = end_position; m_i < all_sequence_length; m_i++) { p_mask[m_i] = static_cast(mask_filter_value); } // Handle left-side padding: mask value before the start position will be mask_filter_value if (has_mask_start_position) { - int start_position = std::min(mask_index[b_i + batch_size], all_sequence_length); + int start_position = std::max(0, std::min(static_cast(mask_index[b_i + batch_size]), all_sequence_length)); for (int m_i = 0; m_i < start_position; m_i++) { p_mask[m_i] = static_cast(mask_filter_value); } diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu index ff7ac67852427..2c67ad342e9e1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu @@ -603,7 +603,7 @@ __global__ void MaskedSoftmaxKernelSmall(const int total_sequence_length, if (threadIdx.x == 0) { const int batch = blockIdx.y; start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0; - end_position = min(total_sequence_length, mask_end[batch]); + end_position = max(0, min(total_sequence_length, mask_end[batch])); // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result. if (start_position >= end_position) { @@ -735,7 +735,7 @@ __global__ void MaskedSoftmaxKernel(const int total_sequence_length, if (threadIdx.x == 0) { const int batch = blockIdx.y; start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0; - end_position = min(total_sequence_length, mask_end[batch]); + end_position = max(0, min(total_sequence_length, mask_end[batch])); // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result. if (start_position >= end_position) { diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 411629535254d..6268e425ebb61 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -2009,6 +2009,56 @@ TEST(ContribOpAttentionTest, AttentionMaskIndexOutOfRange) { AttentionMaskType::MASK_1D_END_START); } +TEST(ContribOpAttentionTest, AttentionMaskIndexNegativeEndPosition) { + // Test that negative end_position in mask_index is rejected (heap underflow prevention). + int batch_size = 2; + int sequence_length = 2; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f, + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + // Negative end_position values in 1D mask_index + std::vector mask_index_data = {-10, -1}; + + OpTester tester("Attention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(number_of_heads)); + tester.AddAttribute("mask_filter_value", static_cast(-10000.0f)); + + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector weights_dims = {hidden_size, 3 * hidden_size}; + std::vector bias_dims = {3 * hidden_size}; + std::vector mask_index_dims = {batch_size}; + std::vector output_dims = {batch_size, sequence_length, hidden_size}; + + tester.AddInput("input", input_dims, input_data); + tester.AddInput("weight", weights_dims, weight_data); + tester.AddInput("bias", bias_dims, bias_data); + tester.AddInput("mask_index", mask_index_dims, mask_index_data); + tester.AddOptionalInputEdge(); // past + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // past_sequence_length + + tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_length * hidden_size, 0.0f)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, "mask_index value", {}, nullptr, &execution_providers); +} + #if !defined(__wasm__) // TODO: fix in web assembly TEST(ContribOpAttentionTest, AttentionPastState_dynamic) { From b446cfcc5b45dc37341edee44a2d9b2e8cd295c2 Mon Sep 17 00:00:00 2001 From: vraspar Date: Mon, 23 Mar 2026 19:40:35 +0000 Subject: [PATCH 2/3] Add missing include for algorithm in attention_helper.h --- onnxruntime/contrib_ops/cpu/bert/attention_helper.h | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h index 73ed43b64f49e..71da9eeff1f13 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_helper.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include "core/util/math.h" #include "core/util/math_cpuonly.h" From 39aa32d944f10eb7d0631f66d023a6c9f87db8d8 Mon Sep 17 00:00:00 2001 From: vraspar Date: Thu, 2 Apr 2026 21:05:19 +0000 Subject: [PATCH 3/3] Guard mask_index validation to CPU tensors only to fix CUDA segfault Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cpu/bert/attention_base.h | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_base.h b/onnxruntime/contrib_ops/cpu/bert/attention_base.h index 9477fbf765808..fad8d9275c555 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_base.h @@ -124,12 +124,15 @@ inline Status AttentionBase::CheckMask(const Tensor* mask_index, // Validate that end_position values (first batch_size elements) are non-negative. // Negative end_position causes out-of-bounds writes in PrepareMask. - const int32_t* mask_data = mask_index->Data(); - for (int64_t i = 0; i < batch_size; i++) { - if (mask_data[i] < 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "mask_index value ", mask_data[i], " at index ", i, - " is negative. mask_index end_position values must be non-negative."); + // Only validate when mask_index is on CPU; GPU tensors are clamped in the CUDA kernel. + if (mask_index->Location().device.Type() == OrtDevice::CPU) { + const int32_t* mask_data = mask_index->Data(); + for (int64_t i = 0; i < batch_size; i++) { + if (mask_data[i] < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "mask_index value ", mask_data[i], " at index ", i, + " is negative. mask_index end_position values must be non-negative."); + } } } } else if (mask_dims.size() == 2) {