diff --git a/onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h b/onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h index 984e8f2490a73..a8e7e0de6ed36 100644 --- a/onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h @@ -115,8 +115,25 @@ 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)); - if (head_size == 0) { - ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(cache_width, 2, "head_size", head_size)); + + int effective_rotary_dim = 0; + ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(cache_width, 2, "effective_rotary_dim", effective_rotary_dim)); + + const bool head_size_inferred = (head_size == 0); + if (head_size_inferred) { + head_size = effective_rotary_dim; + } + + // Only needed when head_size is inferred from the cache; the exact-width check below + // cannot catch a mismatch there because head_size == effective_rotary_dim by construction. + // When num_heads > 0 / rank-4, head_size is known and the exact-width check rejects an + // oversized cache with a more actionable message. + if (head_size_inferred && 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) { diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc index 880c10137f3fe..658c8695643e6 100644 --- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc +++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc @@ -1172,5 +1172,36 @@ 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 +// and head_size is inferred from the cache (rank-3 input without num_heads). This exercises the +// `effective_rotary_dim > hidden_size` guard, which covers the inference path that the +// exact-width check cannot catch (head_size == effective_rotary_dim by construction there). +TEST(RotaryEmbeddingTest, ContribRotaryEmbedding_RejectsCosCacheExceedsHiddenSize_NoNumHeads) { + int batch_size = 1; + int sequence_length = 1; + int hidden_size = 64; + int half_rotary_dim = 64; // cos_cache_dims[1]*2 = 128 > hidden_size; head_size inferred to 128 + int max_sequence_length = 2; + + OpTester test("RotaryEmbedding", 1, onnxruntime::kMSDomain); + test.AddAttribute("interleaved", static_cast(0)); + // num_heads intentionally NOT set so head_size stays 0 on entry and is inferred from cos_cache. + + test.AddInput("input", {batch_size, sequence_length, hidden_size}, + std::vector(hidden_size, 42.0f)); + test.AddInput("position_ids", {1}, {0}); + test.AddInput("cos_cache", {max_sequence_length, half_rotary_dim}, + std::vector(max_sequence_length * half_rotary_dim, 0.0f)); + test.AddInput("sin_cache", {max_sequence_length, half_rotary_dim}, + std::vector(max_sequence_length * half_rotary_dim, 1.0f)); + test.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(hidden_size, 0.0f)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "exceeds input hidden_size", {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc b/onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc index 2f51b8a7a5690..24b152131cc0d 100644 --- a/onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc +++ b/onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc @@ -1412,5 +1412,40 @@ TEST(RotaryEmbeddingTest, RotaryEmbedding_PositionIds_OOB_InBatch_WebGPU_Passthr test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +// Test that a cos_cache whose width does not match head_size / 2 (or rotary_embedding_dim / 2) +// is rejected by the mainline op. The mainline op needs no additional hidden_size/OOB guard +// because it requires num_heads > 0 for a rank-3 input, so head_size is always derived as +// hidden_size / num_heads (never inferred from the cache) and the exact-width check below +// already rejects an over-sized cos_cache. +TEST(RotaryEmbeddingTest, RotaryEmbedding_RejectsCosCacheWidthMismatch) { + // hidden_size = 64, num_heads = 1 => head_size = 64, expected cache width = 32. + // cos_cache dim1 = 64 mismatches the expected 32 and is rejected by the existing width check. + int batch_size = 1; + int sequence_length = 1; + int hidden_size = 64; + int half_rotary_dim = 64; // mismatches expected head_size / 2 = 32 + int max_sequence_length = 2; + + OpTester test("RotaryEmbedding", 23, onnxruntime::kOnnxDomain); + test.AddAttribute("interleaved", static_cast(0)); + test.AddAttribute("num_heads", static_cast(1)); + + test.AddInput("input", {batch_size, sequence_length, hidden_size}, + std::vector(hidden_size, 42.0f)); + test.AddInput("cos_cache", {max_sequence_length, half_rotary_dim}, + std::vector(max_sequence_length * half_rotary_dim, 0.0f)); + test.AddInput("sin_cache", {max_sequence_length, half_rotary_dim}, + std::vector(max_sequence_length * half_rotary_dim, 1.0f)); + test.AddInput("position_ids", {batch_size, sequence_length}, {0}); + test.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(hidden_size, 0.0f)); + + std::vector> 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", + {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime