Harden contrib CPU int narrowing for attention attrs

onnxruntime/contrib_ops/cpu/bert/linear_attention.cc

@@ -12,6 +12,8 @@
 
 #include <algorithm>
 #include <cmath>
+#include <limits>
+#include <vector>
 #include <cstring>
 
 using onnxruntime::concurrency::ThreadPool;
@@ -46,13 +48,17 @@
 template <typename T>
 LinearAttention<T>::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<int>::max(),
+              "q_num_heads must be an integer in [1, INT_MAX]");
   q_num_heads_ = static_cast<int>(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<int>::max(),
+              "kv_num_heads must be an integer in [1, INT_MAX]");
   kv_num_heads_ = static_cast<int>(kv_num_heads);
 
   update_rule_ = info.GetAttrOrDefault<std::string>("update_rule", "gated_delta");

onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #include "core/common/common.h"
+#include <limits>
 #ifndef SHARED_PROVIDER
 #include "core/framework/op_kernel.h"
 #endif
@@ -25,11 +26,17 @@
   template <typename KernelInfoType>
   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<int>::max(),
+                "num_heads must be an integer in [1, INT_MAX]");
     num_heads_ = static_cast<int>(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<int>::max(),
+                "window must be an integer in [1, INT_MAX]");
     window_ = static_cast<int>(window);
   }
 

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 <gsl/gsl>
+#include <limits>
 #include "contrib_ops/cpu/transformers/subgraph_gpt.h"
 #include "contrib_ops/cpu/utils/dump_tensor.h"
 
@@ -170,7 +171,14 @@
                 "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<int>::max(),
+                "subgraph past state dimension 2 is too large for int");
+  ORT_RETURN_IF(past_shape->dim(4).dim_value() > std::numeric_limits<int>::max(),
+                "subgraph past state dimension 4 is too large for int");
+  ORT_RETURN_IF(logits_shape->dim(2).dim_value() > std::numeric_limits<int>::max(),
+                "subgraph logits dimension 2 is too large for int");
 
   // Save parameters related to the subgraph.
   num_heads = static_cast<int>(past_shape->dim(2).dim_value());

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<std::string>("update_rule", "linear");
+  tester.AddAttribute<float>("scale", 1.0f);
+  tester.AddAttribute<int64_t>("q_num_heads", 4294967296LL);
+  tester.AddAttribute<int64_t>("kv_num_heads", 1);
+
+  tester.AddInput<float>("query", {1, 1, 0}, {});
+  tester.AddInput<float>("key", {1, 1, 0}, {});
+  tester.AddInput<float>("value", {1, 1, 0}, {});
+  tester.AddOptionalInputEdge<float>();
+  tester.AddOptionalInputEdge<float>();
+  tester.AddOptionalInputEdge<float>();
+  tester.AddOutput<float>("output", {1, 1, 0}, {});
+  tester.AddOutput<float>("present_state", {1, 1, 0, 0}, {});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<std::string>("update_rule", "linear");
+  tester.AddAttribute<float>("scale", 1.0f);
+  tester.AddAttribute<int64_t>("q_num_heads", 1);
+  tester.AddAttribute<int64_t>("kv_num_heads", 4294967296LL);
+
+  tester.AddInput<float>("query", {1, 1, 0}, {});
+  tester.AddInput<float>("key", {1, 1, 0}, {});
+  tester.AddInput<float>("value", {1, 1, 0}, {});
+  tester.AddOptionalInputEdge<float>();
+  tester.AddOptionalInputEdge<float>();
+  tester.AddOptionalInputEdge<float>();
+  tester.AddOutput<float>("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<float>("present_state", {1, 4294967296LL, 0, 0}, {});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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,