From b647e8daed2e7aff210ba67d4530e16461a15cb1 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 4 Aug 2026 15:51:38 -0700 Subject: [PATCH 1/7] Harden contrib CPU int narrowing for attention attrs Validate int64 attributes and shape-derived values fit in int before narrowing in LinearAttention, LongformerAttentionBase, and GptSubgraph validation. Add LinearAttention regressions for oversized q_num_heads and kv_num_heads attributes. --- .../contrib_ops/cpu/bert/linear_attention.cc | 13 +++-- .../cpu/bert/longformer_attention_base.h | 11 +++- .../cpu/transformers/subgraph_gpt.cc | 8 +++ .../contrib_ops/linear_attention_op_test.cc | 52 +++++++++++++++++++ 4 files changed, 78 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc index 052e7df8bda14..2d5c47c256a7b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc @@ -9,6 +9,7 @@ #include "core/platform/threadpool.h" #include +#include #include using onnxruntime::concurrency::ThreadPool; @@ -41,13 +42,17 @@ REGISTER_KERNEL_TYPED(float) template LinearAttention::LinearAttention(const OpKernelInfo& info) : OpKernel(info) { int64_t q_num_heads = 0; - ORT_ENFORCE(info.GetAttr("q_num_heads", &q_num_heads).IsOK() && q_num_heads > 0, - "q_num_heads must be a positive integer"); + ORT_ENFORCE(info.GetAttr("q_num_heads", &q_num_heads).IsOK() && + q_num_heads > 0 && + q_num_heads <= std::numeric_limits::max(), + "q_num_heads must be an integer in [1, INT_MAX]"); q_num_heads_ = static_cast(q_num_heads); int64_t kv_num_heads = 0; - ORT_ENFORCE(info.GetAttr("kv_num_heads", &kv_num_heads).IsOK() && kv_num_heads > 0, - "kv_num_heads must be a positive integer"); + ORT_ENFORCE(info.GetAttr("kv_num_heads", &kv_num_heads).IsOK() && + kv_num_heads > 0 && + kv_num_heads <= std::numeric_limits::max(), + "kv_num_heads must be an integer in [1, INT_MAX]"); kv_num_heads_ = static_cast(kv_num_heads); update_rule_ = info.GetAttrOrDefault("update_rule", "gated_delta"); diff --git a/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h index bb1dfea38ae80..e6f3bff4c73b0 100644 --- a/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h @@ -4,6 +4,7 @@ #pragma once #include "core/common/common.h" +#include #ifndef SHARED_PROVIDER #include "core/framework/op_kernel.h" #endif @@ -25,11 +26,17 @@ class LongformerAttentionBase { template LongformerAttentionBase(const KernelInfoType& info) { int64_t num_heads = 0; - ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); + ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && + num_heads > 0 && + num_heads <= std::numeric_limits::max(), + "num_heads must be an integer in [1, INT_MAX]"); num_heads_ = static_cast(num_heads); int64_t window = 0; - ORT_ENFORCE(info.GetAttr("window", &window).IsOK() && window > 0); + ORT_ENFORCE(info.GetAttr("window", &window).IsOK() && + window > 0 && + window <= std::numeric_limits::max(), + "window must be an integer in [1, INT_MAX]"); window_ = static_cast(window); } diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc index 030cdb1e1b17f..f256e34ebb472 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc @@ -7,6 +7,7 @@ #include "core/framework/utils.h" #include "core/providers/cpu/tensor/utils.h" #include +#include #include "contrib_ops/cpu/transformers/subgraph_gpt.h" #include "contrib_ops/cpu/utils/dump_tensor.h" @@ -172,6 +173,13 @@ Status GptSubgraph::Validate(const std::vector& subgraph_inputs, ORT_RETURN_IF(!logits_shape->dim(2).has_dim_value() || logits_shape->dim(2).dim_value() <= 0, "subgraph past state dimension 2 shall have a positive value for vocabulary size"); + ORT_RETURN_IF(past_shape->dim(2).dim_value() > std::numeric_limits::max(), + "subgraph past state dimension 2 is too large for int"); + ORT_RETURN_IF(past_shape->dim(4).dim_value() > std::numeric_limits::max(), + "subgraph past state dimension 4 is too large for int"); + ORT_RETURN_IF(logits_shape->dim(2).dim_value() > std::numeric_limits::max(), + "subgraph logits dimension 2 is too large for int"); + // Save parameters related to the subgraph. num_heads = static_cast(past_shape->dim(2).dim_value()); head_size = static_cast(past_shape->dim(4).dim_value()); diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index c8d2f6d24d71a..934fac13bae5c 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1382,5 +1382,57 @@ TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StandardGQA_N16) { RunStandard TEST(ContribOpLinearAttentionTest, LinearRule_StandardGQA_N4) { RunStandardGQA("linear", 8, 2, 32, 64); } TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StandardGQA_N4_Dim128) { RunStandardGQA("gated_delta", 8, 2, 128, 128); } +TEST(ContribOpLinearAttentionTest, RejectsQNumHeadsOverflow) { + auto ep = TryGetEpWithLinearAttention(); + if (!ep) { + GTEST_SKIP() << "LinearAttention kernel not registered"; + return; + } + + OpTester tester("LinearAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("update_rule", "linear"); + tester.AddAttribute("scale", 1.0f); + tester.AddAttribute("q_num_heads", 4294967296LL); + tester.AddAttribute("kv_num_heads", 1); + + tester.AddInput("query", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddInput("key", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddInput("value", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + + std::vector> execution_providers; + execution_providers.push_back(std::move(ep)); + tester.Run(OpTester::ExpectResult::kExpectFailure, "q_num_heads must be an integer in [1, INT_MAX]", + {}, nullptr, &execution_providers); +} + +TEST(ContribOpLinearAttentionTest, RejectsKvNumHeadsOverflow) { + auto ep = TryGetEpWithLinearAttention(); + if (!ep) { + GTEST_SKIP() << "LinearAttention kernel not registered"; + return; + } + + OpTester tester("LinearAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("update_rule", "linear"); + tester.AddAttribute("scale", 1.0f); + tester.AddAttribute("q_num_heads", 1); + tester.AddAttribute("kv_num_heads", 4294967296LL); + + tester.AddInput("query", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddInput("key", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddInput("value", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + + std::vector> execution_providers; + execution_providers.push_back(std::move(ep)); + tester.Run(OpTester::ExpectResult::kExpectFailure, "kv_num_heads must be an integer in [1, INT_MAX]", + {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime From 0e68739283bfb00263b498b6393441f0a998b5cf Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 5 Aug 2026 09:57:42 -0700 Subject: [PATCH 2/7] Fix misleading error message in GPT subgraph validation The error message on line 174 incorrectly referenced 'past state dimension 2' when it was actually validating the logits output dimension 2 (vocabulary size). This fix clarifies the message to correctly reference 'logits dimension 2' for better debugging experience when vocabulary size validation fails. Addresses review comment feedback from PR #31648. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc index f256e34ebb472..097ac27551c34 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc @@ -171,7 +171,7 @@ Status GptSubgraph::Validate(const std::vector& subgraph_inputs, "subgraph logits output is expected to have 3 dimension, got ", logits_shape->dim_size()); ORT_RETURN_IF(!logits_shape->dim(2).has_dim_value() || logits_shape->dim(2).dim_value() <= 0, - "subgraph past state dimension 2 shall have a positive value for vocabulary size"); + "subgraph logits dimension 2 shall have a positive value for vocabulary size"); ORT_RETURN_IF(past_shape->dim(2).dim_value() > std::numeric_limits::max(), "subgraph past state dimension 2 is too large for int"); From 24307864430d8b0e213870537b22211f1dc8f580 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 5 Aug 2026 10:29:41 -0700 Subject: [PATCH 3/7] Fix LinearAttention overflow tests to include required outputs The overflow tests for q_num_heads and kv_num_heads were failing schema validation before reaching the attribute range checks because the LinearAttention node requires two outputs. Add dummy output and present_state outputs so the invalid attribute validation is exercised as intended. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/contrib_ops/linear_attention_op_test.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index 934fac13bae5c..c6471777c49f7 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1401,6 +1401,8 @@ TEST(ContribOpLinearAttentionTest, RejectsQNumHeadsOverflow) { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); + tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddOutput("present_state", {1, 1, 4, 4}, std::vector(16, 0.0f)); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); @@ -1427,6 +1429,8 @@ TEST(ContribOpLinearAttentionTest, RejectsKvNumHeadsOverflow) { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); + tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddOutput("present_state", {1, 1, 4, 4}, std::vector(16, 0.0f)); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); From 0a7aa69056c2c4eb7675dde9abec0e1ff5523410 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 5 Aug 2026 11:08:59 -0700 Subject: [PATCH 4/7] Avoid shape inference mismatch in LinearAttention overflow tests The overflow tests were failing during schema shape inference before reaching the intended q_num_heads / kv_num_heads constructor validation. Use unknown output dimensions in the test harness so the failure path exercises the attribute range checks directly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/contrib_ops/linear_attention_op_test.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index c6471777c49f7..f3dab58671e67 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1401,8 +1401,8 @@ TEST(ContribOpLinearAttentionTest, RejectsQNumHeadsOverflow) { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); - tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddOutput("present_state", {1, 1, 4, 4}, std::vector(16, 0.0f)); + tester.AddOutput("output", {-1, -1, -1}, {}); + tester.AddOutput("present_state", {-1, -1, -1, -1}, {}); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); @@ -1429,8 +1429,8 @@ TEST(ContribOpLinearAttentionTest, RejectsKvNumHeadsOverflow) { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); - tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddOutput("present_state", {1, 1, 4, 4}, std::vector(16, 0.0f)); + tester.AddOutput("output", {-1, -1, -1}, {}); + tester.AddOutput("present_state", {-1, -1, -1, -1}, {}); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); From c6e7905b11667ea05c12a58a8726854445ea298a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 6 Aug 2026 10:42:25 -0700 Subject: [PATCH 5/7] Fix PR31648 overflow tests failing during setup Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test/contrib_ops/linear_attention_op_test.cc | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index f3dab58671e67..c48bb22759d3c 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1401,8 +1401,9 @@ TEST(ContribOpLinearAttentionTest, RejectsQNumHeadsOverflow) { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); - tester.AddOutput("output", {-1, -1, -1}, {}); - tester.AddOutput("present_state", {-1, -1, -1, -1}, {}); + tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddOutput("present_state", {1, 1, 4, 4}, + std::vector(16, 0.0f)); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); @@ -1429,8 +1430,9 @@ TEST(ContribOpLinearAttentionTest, RejectsKvNumHeadsOverflow) { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); - tester.AddOutput("output", {-1, -1, -1}, {}); - tester.AddOutput("present_state", {-1, -1, -1, -1}, {}); + tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddOutput("present_state", {1, 1, 4, 4}, + std::vector(16, 0.0f)); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); From c076922782e3c302bdf7a31b10edd561ca682925 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 6 Aug 2026 13:01:31 -0700 Subject: [PATCH 6/7] Fix PR31648 overflow tests to avoid shape inference conflicts Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/linear_attention_op_test.cc | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index c48bb22759d3c..0933183a37997 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1395,15 +1395,14 @@ TEST(ContribOpLinearAttentionTest, RejectsQNumHeadsOverflow) { tester.AddAttribute("q_num_heads", 4294967296LL); tester.AddAttribute("kv_num_heads", 1); - tester.AddInput("query", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddInput("key", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddInput("value", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddInput("query", {1, 1, 0}, {}); + tester.AddInput("key", {1, 1, 0}, {}); + tester.AddInput("value", {1, 1, 0}, {}); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); - tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddOutput("present_state", {1, 1, 4, 4}, - std::vector(16, 0.0f)); + tester.AddOutput("output", {1, 1, 0}, {}); + tester.AddOutput("present_state", {1, 1, 0, 0}, {}); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); @@ -1424,15 +1423,14 @@ TEST(ContribOpLinearAttentionTest, RejectsKvNumHeadsOverflow) { tester.AddAttribute("q_num_heads", 1); tester.AddAttribute("kv_num_heads", 4294967296LL); - tester.AddInput("query", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddInput("key", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddInput("value", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); + tester.AddInput("query", {1, 1, 0}, {}); + tester.AddInput("key", {1, 1, 0}, {}); + tester.AddInput("value", {1, 1, 0}, {}); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); - tester.AddOutput("output", {1, 1, 4}, {0.0f, 0.0f, 0.0f, 0.0f}); - tester.AddOutput("present_state", {1, 1, 4, 4}, - std::vector(16, 0.0f)); + tester.AddOutput("output", {1, 1, 0}, {}); + tester.AddOutput("present_state", {1, 1, 0, 0}, {}); std::vector> execution_providers; execution_providers.push_back(std::move(ep)); From 1e235df4b4df9f90d9802c78861073c103d30f05 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 11 Aug 2026 14:59:55 -0700 Subject: [PATCH 7/7] Align declared present_state shape in kv_num_heads test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/contrib_ops/linear_attention_op_test.cc | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index 9e2f9ad7f31f9..83c60cb3e5131 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1591,7 +1591,7 @@ static void RunLinearAttentionStateWindowTest(int B, int q_H, int kv_H, int n_k, false, 0.005f, 0.005f); tester.AddOutput("present_state", {W, B, kv_H, dk, dv}, expected_state_window, false, 0.005f, 0.005f); - + std::vector> execution_providers; execution_providers.push_back(std::move(ep)); tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); @@ -1645,7 +1645,10 @@ TEST(ContribOpLinearAttentionTest, RejectsKvNumHeadsOverflow) { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); tester.AddOutput("output", {1, 1, 0}, {}); - tester.AddOutput("present_state", {1, 1, 0, 0}, {}); + // present_state carries H_kv as its second dimension, so it must be declared with the + // oversized attribute value for shape inference to agree. The tensor is still empty + // because d_k and d_v are both 0, and the kernel rejects the attribute before any use. + tester.AddOutput("present_state", {1, 4294967296LL, 0, 0}, {}); std::vector> execution_providers; execution_providers.push_back(std::move(ep));