diff --git a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc index a385eb1f95b96..adb325ceca80e 100644 --- a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc @@ -12,6 +12,8 @@ #include #include +#include +#include #include using onnxruntime::concurrency::ThreadPool; @@ -46,13 +48,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..097ac27551c34 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" @@ -170,7 +171,14 @@ 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"); + 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()); diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index e0eaa28c4012e..5f20ff14e007b 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1608,6 +1608,65 @@ static void RunLinearAttentionStateWindowTest(int B, int q_H, int kv_H, int n_k, tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +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, 0}, {}); + tester.AddInput("key", {1, 1, 0}, {}); + tester.AddInput("value", {1, 1, 0}, {}); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + 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)); + 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, 0}, {}); + tester.AddInput("key", {1, 1, 0}, {}); + tester.AddInput("value", {1, 1, 0}, {}); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOutput("output", {1, 1, 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)); + tester.Run(OpTester::ExpectResult::kExpectFailure, "kv_num_heads must be an integer in [1, INT_MAX]", + {}, nullptr, &execution_providers); +} + // d_k = 4 is not a decode fast-path shape, so this lands on the generic recurrent kernel. TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StateWindow) { RunLinearAttentionStateWindowTest(/*B=*/1, /*q_H=*/2, /*kv_H=*/2, /*n_k=*/2, /*T=*/5,