From 2818327cd2981fb954a41e5f92149004c8c679dc Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 29 Apr 2026 11:58:25 -0700 Subject: [PATCH 1/9] Validate seqlens_k against cos_cache bounds in GroupQueryAttention to prevent rotary embedding OOB read --- .../contrib_ops/cpu/bert/group_query_attention.cc | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 5698bcb659f20..6ff6bd8deeeff 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -144,6 +144,19 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); if (do_rotary_) { ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); + // Validate seqlens_k values against cos_cache size to prevent OOB in rotary embedding lookup. + { + const int cos_cache_max_seq = static_cast(cos_cache->Shape().GetDims()[0]); + const int32_t* seqlens_k_data = seqlens_k->Data(); + for (int b = 0; b < batch_size; b++) { + // position_id = seqlens_k[b] (in token generation), must be < cos_cache rows + if (seqlens_k_data[b] >= cos_cache_max_seq) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "seqlens_k[", b, "] = ", seqlens_k_data[b], + " exceeds cos_cache dimension 0 (", cos_cache_max_seq, ")"); + } + } + } // Initialize rotary parameters rotary_embedding_helper::RotaryParameters rotary_params = {}; rotary_params.batch_size = batch_size; From caca09b95318d3d66cf9ef768113d6878b89c8bc Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 29 Apr 2026 16:21:03 -0700 Subject: [PATCH 2/9] Address comments --- .../cpu/bert/group_query_attention.cc | 12 ++-- .../group_query_attention_op_test.cc | 56 +++++++++++++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 6ff6bd8deeeff..234e8af6e45c5 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -144,16 +144,18 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); if (do_rotary_) { ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); - // Validate seqlens_k values against cos_cache size to prevent OOB in rotary embedding lookup. + // Validate seqlens_k values against cos/sin cache size to prevent OOB in rotary embedding lookup. + // Use the minimum of cos_cache and sin_cache dim-0 since CheckRotaryCaches does not enforce equality. { - const int cos_cache_max_seq = static_cast(cos_cache->Shape().GetDims()[0]); + const int rotary_cache_max_seq = static_cast(std::min(cos_cache->Shape().GetDims()[0], + sin_cache->Shape().GetDims()[0])); const int32_t* seqlens_k_data = seqlens_k->Data(); for (int b = 0; b < batch_size; b++) { - // position_id = seqlens_k[b] (in token generation), must be < cos_cache rows - if (seqlens_k_data[b] >= cos_cache_max_seq) { + // position_id = seqlens_k[b] (in token generation), must be < cache rows + if (seqlens_k_data[b] >= rotary_cache_max_seq) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "seqlens_k[", b, "] = ", seqlens_k_data[b], - " exceeds cos_cache dimension 0 (", cos_cache_max_seq, ")"); + " exceeds rotary cache dimension 0 (", rotary_cache_max_seq, ")"); } } } diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 0690094031bb8..1895f51a1e945 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -307,5 +307,61 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { {}, nullptr, &execution_providers); } +// Regression: seqlens_k valid for KV cache but exceeding cos_cache.shape[0] must be rejected +// when do_rotary is enabled. Without this check, the position ID derived from seqlens_k +// would index out of bounds in the cos/sin cache, leaking heap memory into output. +TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_OOB) { + constexpr int num_heads = 1; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int rotary_half_dim = head_size / 2; // cos/sin cache dim-1 + + constexpr int cos_cache_max_seq = 4; // small rotary cache + constexpr int past_seq_len = 16; // large KV cache + constexpr int seqlens_k_val = 10; // valid for KV (10 < 16) but OOB for cos (10 >= 4) + constexpr int total_seq_len = 11; // seqlens_k + 1 + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + tester.AddAttribute("do_rotary", static_cast(1)); + + tester.AddInput("query", {1, 1, hidden_size}, std::vector(hidden_size, 1.0f)); + tester.AddInput("key", {1, 1, kv_hidden_size}, std::vector(kv_hidden_size, 1.0f)); + tester.AddInput("value", {1, 1, kv_hidden_size}, std::vector(kv_hidden_size, 1.0f)); + + // Past KV cache is large enough for seqlens_k=10 + tester.AddInput("past_key", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.5f)); + tester.AddInput("past_value", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.5f)); + + tester.AddInput("seqlens_k", {1}, {seqlens_k_val}); + tester.AddInput("total_sequence_length", {1}, {total_seq_len}); + + // cos/sin cache with only 4 rows — seqlens_k=10 exceeds this + tester.AddInput("cos_cache", {cos_cache_max_seq, rotary_half_dim}, + std::vector(cos_cache_max_seq * rotary_half_dim, 1.0f)); + tester.AddInput("sin_cache", {cos_cache_max_seq, rotary_half_dim}, + std::vector(cos_cache_max_seq * rotary_half_dim, 0.0f)); + + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + tester.AddOutput("output", {1, 1, hidden_size}, std::vector(hidden_size, 0.0f)); + tester.AddOutput("present_key", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.0f)); + tester.AddOutput("present_value", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.0f)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, "exceeds rotary cache dimension 0", + {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime From e2d30d1a9c607ac449996e27587a45a8478be02b Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 29 Apr 2026 17:27:10 -0700 Subject: [PATCH 3/9] Fix unit tests --- onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 1895f51a1e945..c3494aacd28b9 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -313,10 +313,10 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_OOB) { constexpr int num_heads = 1; constexpr int kv_num_heads = 1; - constexpr int head_size = 8; + constexpr int head_size = 16; // must be multiple of 16 for rotary constexpr int hidden_size = num_heads * head_size; constexpr int kv_hidden_size = kv_num_heads * head_size; - constexpr int rotary_half_dim = head_size / 2; // cos/sin cache dim-1 + constexpr int rotary_half_dim = head_size / 2; // cos/sin cache dim-1 = 8 constexpr int cos_cache_max_seq = 4; // small rotary cache constexpr int past_seq_len = 16; // large KV cache From d049874001db593f36aa105014831c612464a621 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 29 Apr 2026 17:28:23 -0700 Subject: [PATCH 4/9] Fix lint errors --- onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc | 2 +- .../test/contrib_ops/group_query_attention_op_test.cc | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 234e8af6e45c5..c370a178db2fe 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -148,7 +148,7 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { // Use the minimum of cos_cache and sin_cache dim-0 since CheckRotaryCaches does not enforce equality. { const int rotary_cache_max_seq = static_cast(std::min(cos_cache->Shape().GetDims()[0], - sin_cache->Shape().GetDims()[0])); + sin_cache->Shape().GetDims()[0])); const int32_t* seqlens_k_data = seqlens_k->Data(); for (int b = 0; b < batch_size; b++) { // position_id = seqlens_k[b] (in token generation), must be < cache rows diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index c3494aacd28b9..060bea2f4f167 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -318,10 +318,10 @@ TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_OOB) { constexpr int kv_hidden_size = kv_num_heads * head_size; constexpr int rotary_half_dim = head_size / 2; // cos/sin cache dim-1 = 8 - constexpr int cos_cache_max_seq = 4; // small rotary cache - constexpr int past_seq_len = 16; // large KV cache - constexpr int seqlens_k_val = 10; // valid for KV (10 < 16) but OOB for cos (10 >= 4) - constexpr int total_seq_len = 11; // seqlens_k + 1 + constexpr int cos_cache_max_seq = 4; // small rotary cache + constexpr int past_seq_len = 16; // large KV cache + constexpr int seqlens_k_val = 10; // valid for KV (10 < 16) but OOB for cos (10 >= 4) + constexpr int total_seq_len = 11; // seqlens_k + 1 OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); From a8f6fd12472c7aa4f07d0be621d4ffad610c1eab Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 30 Apr 2026 10:58:50 -0700 Subject: [PATCH 5/9] Fix tests --- .../group_query_attention_op_test.cc | 56 ++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 060bea2f4f167..e3c6825c1eb5e 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -321,7 +321,7 @@ TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_OOB) { constexpr int cos_cache_max_seq = 4; // small rotary cache constexpr int past_seq_len = 16; // large KV cache constexpr int seqlens_k_val = 10; // valid for KV (10 < 16) but OOB for cos (10 >= 4) - constexpr int total_seq_len = 11; // seqlens_k + 1 + constexpr int total_seq_len = 4; // passes CheckRotaryCaches (4 <= cos_cache_max_seq) OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); @@ -363,5 +363,59 @@ TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_OOB) { {}, nullptr, &execution_providers); } +// Positive test: seqlens_k within cos/sin cache bounds with do_rotary enabled should succeed. +TEST(GroupQueryAttentionTest, SeqlensKWithinCosCache_Rotary) { + constexpr int num_heads = 1; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; // must be multiple of 16 for rotary + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int rotary_half_dim = head_size / 2; + + constexpr int cos_cache_max_seq = 16; // rotary cache large enough + constexpr int past_seq_len = 16; + constexpr int seqlens_k_val = 3; // valid: 3 < 16 (cos cache) and 3 < 16 (KV cache) + constexpr int total_seq_len = 4; // seqlens_k + 1 + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + tester.AddAttribute("do_rotary", static_cast(1)); + + tester.AddInput("query", {1, 1, hidden_size}, std::vector(hidden_size, 1.0f)); + tester.AddInput("key", {1, 1, kv_hidden_size}, std::vector(kv_hidden_size, 1.0f)); + tester.AddInput("value", {1, 1, kv_hidden_size}, std::vector(kv_hidden_size, 1.0f)); + + tester.AddInput("past_key", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.5f)); + tester.AddInput("past_value", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.5f)); + + tester.AddInput("seqlens_k", {1}, {seqlens_k_val}); + tester.AddInput("total_sequence_length", {1}, {total_seq_len}); + + tester.AddInput("cos_cache", {cos_cache_max_seq, rotary_half_dim}, + std::vector(cos_cache_max_seq * rotary_half_dim, 1.0f)); + tester.AddInput("sin_cache", {cos_cache_max_seq, rotary_half_dim}, + std::vector(cos_cache_max_seq * rotary_half_dim, 0.0f)); + + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + tester.AddOutput("output", {1, 1, hidden_size}, std::vector(hidden_size, 0.0f)); + tester.AddOutput("present_key", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.0f)); + tester.AddOutput("present_value", {1, kv_num_heads, past_seq_len, head_size}, + std::vector(kv_num_heads * past_seq_len * head_size, 0.0f)); + + tester.SetOutputTolerance(1e6f); // shape acceptance test, not numerical correctness + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", + {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime From 8b96eca2aea4d239bb5ae13396e632319604096a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 30 Apr 2026 11:14:08 -0700 Subject: [PATCH 6/9] Fix lint --- onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index e3c6825c1eb5e..e750c84679798 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -374,8 +374,8 @@ TEST(GroupQueryAttentionTest, SeqlensKWithinCosCache_Rotary) { constexpr int cos_cache_max_seq = 16; // rotary cache large enough constexpr int past_seq_len = 16; - constexpr int seqlens_k_val = 3; // valid: 3 < 16 (cos cache) and 3 < 16 (KV cache) - constexpr int total_seq_len = 4; // seqlens_k + 1 + constexpr int seqlens_k_val = 3; // valid: 3 < 16 (cos cache) and 3 < 16 (KV cache) + constexpr int total_seq_len = 4; // seqlens_k + 1 OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); From aebe7dd69ccf36145db3299a26ab929b165b095a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 1 May 2026 13:58:03 -0700 Subject: [PATCH 7/9] Address comments --- .../cpu/bert/group_query_attention.cc | 4 +- .../group_query_attention_op_test.cc | 61 ++++++++++++++++++- 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index c370a178db2fe..8aa6e87f1d0c3 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -152,10 +152,10 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { const int32_t* seqlens_k_data = seqlens_k->Data(); for (int b = 0; b < batch_size; b++) { // position_id = seqlens_k[b] (in token generation), must be < cache rows - if (seqlens_k_data[b] >= rotary_cache_max_seq) { + if (seqlens_k_data[b] < 0 || seqlens_k_data[b] >= rotary_cache_max_seq) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "seqlens_k[", b, "] = ", seqlens_k_data[b], - " exceeds rotary cache dimension 0 (", rotary_cache_max_seq, ")"); + " is out of range for rotary cache dimension 0 (", rotary_cache_max_seq, ")"); } } } diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index e750c84679798..eb025683e1813 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -359,7 +359,7 @@ TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_OOB) { std::vector> execution_providers; execution_providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectFailure, "exceeds rotary cache dimension 0", + tester.Run(OpTester::ExpectResult::kExpectFailure, "is out of range for rotary cache dimension 0", {}, nullptr, &execution_providers); } @@ -417,5 +417,64 @@ TEST(GroupQueryAttentionTest, SeqlensKWithinCosCache_Rotary) { {}, nullptr, &execution_providers); } +// Multi-batch test: one valid and one OOB seqlens_k value. +// Verifies the validation loop correctly identifies the offending batch index. +TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_MultiBatch) { + constexpr int num_heads = 1; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int rotary_half_dim = head_size / 2; + + constexpr int cos_cache_max_seq = 4; + constexpr int past_seq_len = 16; + constexpr int total_seq_len = 4; + constexpr int batch_size = 2; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + tester.AddAttribute("do_rotary", static_cast(1)); + + tester.AddInput("query", {batch_size, 1, hidden_size}, + std::vector(batch_size * hidden_size, 1.0f)); + tester.AddInput("key", {batch_size, 1, kv_hidden_size}, + std::vector(batch_size * kv_hidden_size, 1.0f)); + tester.AddInput("value", {batch_size, 1, kv_hidden_size}, + std::vector(batch_size * kv_hidden_size, 1.0f)); + + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.5f)); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.5f)); + + // seqlens_k: batch 0 is valid (3 < 4), batch 1 is OOB (10 >= 4) + tester.AddInput("seqlens_k", {batch_size}, {3, 10}); + tester.AddInput("total_sequence_length", {1}, {total_seq_len}); + + tester.AddInput("cos_cache", {cos_cache_max_seq, rotary_half_dim}, + std::vector(cos_cache_max_seq * rotary_half_dim, 1.0f)); + tester.AddInput("sin_cache", {cos_cache_max_seq, rotary_half_dim}, + std::vector(cos_cache_max_seq * rotary_half_dim, 0.0f)); + + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + tester.AddOutput("output", {batch_size, 1, hidden_size}, + std::vector(batch_size * hidden_size, 0.0f)); + tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.0f)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + // Error should reference batch index 1: seqlens_k[1] = 10 + tester.Run(OpTester::ExpectResult::kExpectFailure, "seqlens_k[1] = 10", + {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime From 7fc197349fa0a1000ef369128a0dd228377e8889 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 08:43:49 -0700 Subject: [PATCH 8/9] Address comments --- .../cpu/bert/group_query_attention.cc | 18 +++--------------- .../cpu/bert/group_query_attention_helper.h | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 8aa6e87f1d0c3..5a4f6795865c0 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -144,21 +144,9 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); if (do_rotary_) { ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); - // Validate seqlens_k values against cos/sin cache size to prevent OOB in rotary embedding lookup. - // Use the minimum of cos_cache and sin_cache dim-0 since CheckRotaryCaches does not enforce equality. - { - const int rotary_cache_max_seq = static_cast(std::min(cos_cache->Shape().GetDims()[0], - sin_cache->Shape().GetDims()[0])); - const int32_t* seqlens_k_data = seqlens_k->Data(); - for (int b = 0; b < batch_size; b++) { - // position_id = seqlens_k[b] (in token generation), must be < cache rows - if (seqlens_k_data[b] < 0 || seqlens_k_data[b] >= rotary_cache_max_seq) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "seqlens_k[", b, "] = ", seqlens_k_data[b], - " is out of range for rotary cache dimension 0 (", rotary_cache_max_seq, ")"); - } - } - } + // Validation of seqlens_k against rotary cache size is now performed in CheckInputs() + // to ensure all execution providers (CPU, CUDA, etc.) get the protection in one place. + // Initialize rotary parameters rotary_embedding_helper::RotaryParameters rotary_params = {}; rotary_params.batch_size = batch_size; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index f5399e307fbca..bfb454dfc05cf 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -284,6 +284,22 @@ Status CheckInputs(const T* query, int rotary_dim = 0; if (cos_cache != nullptr && sin_cache != nullptr) { ORT_RETURN_IF_ERROR(CheckRotaryCaches(cos_cache, sin_cache, head_size, total_sequence_length, rotary_dim)); + + // Validate seqlens_k against rotary cache size when rotary embeddings are enabled. + // This prevents OOB access when deriving position IDs from seqlens_k during rotary embedding. + const bool is_seqlens_k_on_cpu = (seqlens_k->Location().device.Type() == OrtDevice::CPU); + if (is_seqlens_k_on_cpu) { + const int rotary_cache_max_seq = static_cast(std::min(cos_cache->Shape().GetDims()[0], + sin_cache->Shape().GetDims()[0])); + const int32_t* seqlens_k_data = seqlens_k->template Data(); + for (int b = 0; b < batch_size; b++) { + if (seqlens_k_data[b] < 0 || seqlens_k_data[b] >= rotary_cache_max_seq) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "seqlens_k[", b, "] = ", seqlens_k_data[b], + " is out of range for rotary cache dimension 0 (", rotary_cache_max_seq, ")"); + } + } + } } else if (cos_cache != nullptr || sin_cache != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'cos_cache' and 'sin_cache' shall be both present or both absent."); From 097e16ff8b48f40e855d4678b28f634b9ee19a6d Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 21 May 2026 17:33:57 -0700 Subject: [PATCH 9/9] Address comments --- onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc | 5 +++-- .../contrib_ops/cpu/bert/group_query_attention_helper.h | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 5a4f6795865c0..bed927db43202 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -144,8 +144,9 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); if (do_rotary_) { ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); - // Validation of seqlens_k against rotary cache size is now performed in CheckInputs() - // to ensure all execution providers (CPU, CUDA, etc.) get the protection in one place. + // Validation of seqlens_k against rotary cache size is performed in CheckInputs() + // when seqlens_k is on CPU. GPU EPs where seqlens_k resides on device rely on + // RunRotaryEmbedding's position_ids validation for OOB protection. // Initialize rotary parameters rotary_embedding_helper::RotaryParameters rotary_params = {}; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index bfb454dfc05cf..a68a3cdd9efbf 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -289,11 +289,11 @@ Status CheckInputs(const T* query, // This prevents OOB access when deriving position IDs from seqlens_k during rotary embedding. const bool is_seqlens_k_on_cpu = (seqlens_k->Location().device.Type() == OrtDevice::CPU); if (is_seqlens_k_on_cpu) { - const int rotary_cache_max_seq = static_cast(std::min(cos_cache->Shape().GetDims()[0], - sin_cache->Shape().GetDims()[0])); + const int64_t rotary_cache_max_seq = std::min(cos_cache->Shape().GetDims()[0], + sin_cache->Shape().GetDims()[0]); const int32_t* seqlens_k_data = seqlens_k->template Data(); for (int b = 0; b < batch_size; b++) { - if (seqlens_k_data[b] < 0 || seqlens_k_data[b] >= rotary_cache_max_seq) { + if (seqlens_k_data[b] < 0 || static_cast(seqlens_k_data[b]) >= rotary_cache_max_seq) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "seqlens_k[", b, "] = ", seqlens_k_data[b], " is out of range for rotary cache dimension 0 (", rotary_cache_max_seq, ")");