Fix rotary embedding oob issue

onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h

@@ -115,8 +115,30 @@ Status CheckInputs(const T* input,
   if (rotary_embedding_dim == 0) {
     int cache_width = 0;
     ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(cos_cache_dims[1], "cache_width", cache_width));
+
+    int effective_rotary_dim = 0;
+    ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(cache_width, 2, "effective_rotary_dim", effective_rotary_dim));
+
     if (head_size == 0) {
-      ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(cache_width, 2, "head_size", head_size));
+      head_size = effective_rotary_dim;
+    }
+
+    // Rotary embedding is applied per head, so the inferred dimension must not exceed head_size.
+    if (head_size > 0 && effective_rotary_dim > head_size) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "RotaryEmbedding: cos_cache dimension (", cache_width,
+                             " * 2 = ", effective_rotary_dim,
+                             ") exceeds head_size (", head_size,
+                             ") when rotary_embedding_dim is 0");
+    }
+
+    // Also guard against exceeding the full hidden_size (covers num_heads==0 / rank-3 without num_heads).
+    if (hidden_size > 0 && effective_rotary_dim > hidden_size) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "RotaryEmbedding: cos_cache dimension (", cache_width,
+                             " * 2 = ", effective_rotary_dim,
+                             ") exceeds input hidden_size (", hidden_size,
+                             ") when rotary_embedding_dim is 0");
     }
   } else {
     if (!transposed) {
@@ -158,8 +180,8 @@ Status CheckInputs(const T* input,
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'cos_cache' dimension 0 should be same as ",
                            "max_sequence_length, got ", cos_cache_dims[0]);
   }
-  if ((head_size / 2) != static_cast<int>(cos_cache_dims[1]) &&
-      (rotary_embedding_dim <= 0 || (rotary_embedding_dim / 2) != static_cast<int>(cos_cache_dims[1]))) {
+  const int expected_cache_width = rotary_embedding_dim > 0 ? (rotary_embedding_dim / 2) : (head_size / 2);
+  if (expected_cache_width != static_cast<int>(cos_cache_dims[1])) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'cos_cache' dimension 1 should be same as ",
                            "head_size / 2 or rotary_embedding_dim / 2, got ", cos_cache_dims[1]);
   }

onnxruntime/core/providers/cpu/llm/rotary_embedding_helper.h

@@ -131,6 +131,20 @@ Status CheckInputs(const T* input,
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "rotary_embedding_dim must be less than or equal to ",
                              "head_size");
     }
+
+    // Validate that the rotary embedding dimension derived from cos_cache does not exceed head_size.
+    if (rotary_embedding_dim == 0 && head_size > 0) {
+      int cache_width = static_cast<int>(cos_cache_dims[2]);
+      int effective_rotary_dim = cache_width * 2;
+      if (effective_rotary_dim > head_size) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "RotaryEmbedding: cos_cache dimension (", cache_width,
+                               " * 2 = ", effective_rotary_dim,
+                               ") exceeds head_size (", head_size,
+                               ") when rotary_embedding_dim is 0");
+      }
+    }
+
     // Check cos_cache input shapes
     if (cos_cache_dims[2] != (rotary_embedding_dim > 0 ? rotary_embedding_dim : head_size) / 2) {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'cos_cache' dimension 2 should be same as ",
@@ -166,6 +180,20 @@ Status CheckInputs(const T* input,
                              "head_size");
     }
 
+    // Validate that the rotary embedding dimension derived from cos_cache does not exceed head_size,
+    // which would cause an out-of-bounds read on the input tensor.
+    if (rotary_embedding_dim == 0 && head_size > 0) {
+      int cache_width = static_cast<int>(cos_cache_dims[1]);
+      int effective_rotary_dim = cache_width * 2;
+      if (effective_rotary_dim > head_size) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "RotaryEmbedding: cos_cache dimension (", cache_width,
+                               " * 2 = ", effective_rotary_dim,
+                               ") exceeds head_size (", head_size,
+                               ") when rotary_embedding_dim is 0");
+      }
+    }
+
     int position_ids_batch = 0;
     int position_ids_sequence = 0;
     ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(position_ids_dims[0], "position_ids_dim0", position_ids_batch));

onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc

@@ -1034,7 +1034,7 @@ TEST(RotaryEmbeddingTest, ContribRotaryEmbedding_RejectsRank3MalformedCacheWidth
   std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
   execution_providers.push_back(DefaultCpuExecutionProvider());
   test.Run(OpTester::ExpectResult::kExpectFailure,
-           "Input 'cos_cache' dimension 1 should be same as head_size / 2 or rotary_embedding_dim / 2, got 8",
+           "cos_cache dimension",
            {}, nullptr, &execution_providers);
 }
 
@@ -1051,7 +1051,7 @@ TEST(RotaryEmbeddingTest, ContribRotaryEmbedding_RejectsRank4MalformedCacheWidth
   std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
   execution_providers.push_back(DefaultCpuExecutionProvider());
   test.Run(OpTester::ExpectResult::kExpectFailure,
-           "Input 'cos_cache' dimension 1 should be same as head_size / 2 or rotary_embedding_dim / 2, got 8",
+           "cos_cache dimension",
            {}, nullptr, &execution_providers);
 }
 
@@ -1172,5 +1172,33 @@ TEST(RotaryEmbeddingTest, ContribRotaryEmbedding_PositionIds_Negative_WebGPU_Pas
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
 }
 
+// Test that cos_cache dimension exceeding hidden_size is rejected when rotary_embedding_dim=0.
+TEST(RotaryEmbeddingTest, ContribRotaryEmbedding_RejectsCosCacheExceedsHiddenSize) {
+  int batch_size = 1;
+  int sequence_length = 1;
+  int hidden_size = 64;
+  int half_rotary_dim = 64;  // makes cos_cache_dims[1]*2 = 128 > hidden_size
+  int max_sequence_length = 2;
+
+  OpTester test("RotaryEmbedding", 1, onnxruntime::kMSDomain);
+  test.AddAttribute<int64_t>("interleaved", static_cast<int64_t>(0));
+  test.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(1));
+
+  test.AddInput<float>("input", {batch_size, sequence_length, hidden_size},
+                       std::vector<float>(hidden_size, 42.0f));
+  test.AddInput<int64_t>("position_ids", {1}, {0});
+  test.AddInput<float>("cos_cache", {max_sequence_length, half_rotary_dim},
+                       std::vector<float>(max_sequence_length * half_rotary_dim, 0.0f));
+  test.AddInput<float>("sin_cache", {max_sequence_length, half_rotary_dim},
+                       std::vector<float>(max_sequence_length * half_rotary_dim, 1.0f));
+  test.AddOutput<float>("output", {batch_size, sequence_length, hidden_size},
+                        std::vector<float>(hidden_size, 0.0f));
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectFailure,
+           "cos_cache dimension", {}, nullptr, &execution_providers);
+}
+
 }  // namespace test
 }  // namespace onnxruntime

onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc

@@ -1412,5 +1412,34 @@ TEST(RotaryEmbeddingTest, RotaryEmbedding_PositionIds_OOB_InBatch_WebGPU_Passthr
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
 }
 
+// Test that cos_cache dimension exceeding hidden_size is rejected when rotary_embedding_dim=0.
+TEST(RotaryEmbeddingTest, RotaryEmbedding_RejectsCosCacheExceedsHiddenSize) {
+  // hidden_size = 64, cos_cache dim1 = 64 => effective rotary dim = 128 > 64
+  int batch_size = 1;
+  int sequence_length = 1;
+  int hidden_size = 64;
+  int half_rotary_dim = 64;  // makes cos_cache_dims[1]*2 = 128 > hidden_size
+  int max_sequence_length = 2;
+
+  OpTester test("RotaryEmbedding", 23, onnxruntime::kOnnxDomain);
+  test.AddAttribute<int64_t>("interleaved", static_cast<int64_t>(0));
+  test.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(1));
+
+  test.AddInput<float>("input", {batch_size, sequence_length, hidden_size},
+                       std::vector<float>(hidden_size, 42.0f));
+  test.AddInput<float>("cos_cache", {max_sequence_length, half_rotary_dim},
+                       std::vector<float>(max_sequence_length * half_rotary_dim, 0.0f));
+  test.AddInput<float>("sin_cache", {max_sequence_length, half_rotary_dim},
+                       std::vector<float>(max_sequence_length * half_rotary_dim, 1.0f));
+  test.AddInput<int64_t>("position_ids", {batch_size, sequence_length}, {0});
+  test.AddOutput<float>("output", {batch_size, sequence_length, hidden_size},
+                        std::vector<float>(hidden_size, 0.0f));
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectFailure,
+           "cos_cache dimension", {}, nullptr, &execution_providers);
+}
+
 }  // namespace test
 }  // namespace onnxruntime