diff --git a/onnxruntime/contrib_ops/cuda/bert/bert_padding.cu b/onnxruntime/contrib_ops/cuda/bert/bert_padding.cu index 7d4db2ca4370a..4bfc280b916b5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/bert_padding.cu +++ b/onnxruntime/contrib_ops/cuda/bert/bert_padding.cu @@ -55,7 +55,9 @@ __global__ void getTokenOffset(int* token_count_buffer, int index = 0; cumulated_token_count[0] = 0; for (int i = 0; i < batch_size; i++) { - const int count = sequence_token_count[i]; + // token_offset holds exactly batch_size * sequence_length entries, so a per-row count outside + // [0, sequence_length] would drive both loops below past the end of the allocation. + const int count = min(max(sequence_token_count[i], 0), sequence_length); if (count > max_tokens) { max_tokens = count; } @@ -70,7 +72,7 @@ __global__ void getTokenOffset(int* token_count_buffer, // Offset of paddings for (int i = 0; i < batch_size; i++) { - const int count = sequence_token_count[i]; + const int count = min(max(sequence_token_count[i], 0), sequence_length); for (int j = 0; j < sequence_length - count; j++) { token_offset[index] = i * sequence_length + count + j; index++; diff --git a/onnxruntime/contrib_ops/cuda/bert/remove_padding.cc b/onnxruntime/contrib_ops/cuda/bert/remove_padding.cc index eba4c48301cf3..2522717b1063c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/remove_padding.cc +++ b/onnxruntime/contrib_ops/cuda/bert/remove_padding.cc @@ -53,6 +53,13 @@ Status RemovePadding::ComputeInternal(OpKernelContext* context) const { int64_t sequence_length = dims[1]; int64_t hidden_size = dims[2]; + const auto& sequence_token_count_dims = sequence_token_count->Shape().GetDims(); + if (sequence_token_count_dims.size() != 1 || sequence_token_count_dims[0] != batch_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'sequence_token_count' is expected to have shape (batch_size), got shape ", + sequence_token_count->Shape(), " for batch_size ", batch_size); + } + auto token_count_buffer = GetScratchBuffer(2, GetComputeStream(context)); TensorShapeVector token_offset_shape(2); diff --git a/onnxruntime/test/contrib_ops/remove_padding_op_test.cc b/onnxruntime/test/contrib_ops/remove_padding_op_test.cc index fe415e09fde62..d163d8ceddb74 100644 --- a/onnxruntime/test/contrib_ops/remove_padding_op_test.cc +++ b/onnxruntime/test/contrib_ops/remove_padding_op_test.cc @@ -229,5 +229,79 @@ TEST(RemovePaddingTest, RemovePaddingBatch3_AllWithPadding) { total_tokens); } +// Out-of-range sequence_token_count values must be clamped to [0, sequence_length] so the kernel +// stays inside the token_offset allocation of batch_size * sequence_length entries. +TEST(RemovePaddingTest, RemovePaddingBatch1_TokenCountAboveSequenceLength) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + int total_tokens = 2; + int max_token_count = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + // Far beyond sequence_length; clamped to sequence_length. + std::vector sequence_token_count_data = {0x40000000}; + + std::vector output_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector token_offset_data = {0, 1}; + + std::vector cumulated_seq_len_data = {0, 2}; + + RunRemovePaddingTests( + input_data, + sequence_token_count_data, + output_data, + token_offset_data, + cumulated_seq_len_data, + max_token_count, + batch_size, + sequence_length, + hidden_size, + total_tokens); +} + +TEST(RemovePaddingTest, RemovePaddingBatch2_NegativeTokenCount) { + int batch_size = 2; + int sequence_length = 2; + int hidden_size = 4; + int total_tokens = 2; + int max_token_count = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f, + 0.1f, 0.2f, 0.3f, 0.4f, + 0.5f, 0.6f, 0.7f, 0.8f}; + + // Negative count is clamped to 0, so the first sequence contributes no tokens. + std::vector sequence_token_count_data = {-5, 2}; + + std::vector output_data = { + 0.1f, 0.2f, 0.3f, 0.4f, + 0.5f, 0.6f, 0.7f, 0.8f}; + + std::vector token_offset_data = {2, 3, 0, 1}; + + std::vector cumulated_seq_len_data = {0, 0, 2}; + + RunRemovePaddingTests( + input_data, + sequence_token_count_data, + output_data, + token_offset_data, + cumulated_seq_len_data, + max_token_count, + batch_size, + sequence_length, + hidden_size, + total_tokens); +} + } // namespace test } // namespace onnxruntime