Validate BeamSearch vocab_size against logits width

onnxruntime/contrib_ops/cpu/transformers/beam_search_parameters.cc

@@ -144,9 +144,13 @@ void BeamSearchParameters::ParseFromInputs(OpKernelContext* context) {
 
 void BeamSearchParameters::SetSubgraphParameters(int vocabulary_size, int heads, int hidden_size_per_head, int layers) {
   // Override vocab_size using the inferred shape from the decoder subgraph ONLY IF
-  // the vocab_size hasn't been explicitly specified by the user (as an attribute of BeamSearch)
+  // the vocab_size hasn't been explicitly specified by the user (as an attribute of BeamSearch).
+  // Reject an explicitly supplied vocab_size that exceeds the decoder logits width.
   if (vocab_size == -1 || vocab_size == 0) {
     vocab_size = vocabulary_size;
+  } else if (vocab_size > vocabulary_size) {
+    ORT_THROW("vocab_size attribute (", vocab_size,
+              ") cannot exceed decoder subgraph logits width (", vocabulary_size, ")");
   }
   num_heads = heads;
   head_size = hidden_size_per_head;

onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc

@@ -305,6 +305,10 @@ Status ProcessLogits(const OrtValue& logits,                                 //
   // where input_length equals to parameters_->sequence_length for first subgraph call, and 1 for the remaining calls.
   const TensorShape& logits_shape = logits.Get<Tensor>().Shape();
   ORT_ENFORCE(logits_shape.NumDimensions() == 3);
+  const auto logits_vocab_size = static_cast<int>(logits_shape[2]);
+  ORT_RETURN_IF_NOT(vocab_size > 0 && vocab_size <= logits_vocab_size,
+                    "BeamSearch: vocab_size attribute (", vocab_size,
+                    ") must be positive and not exceed decoder logits width (", logits_vocab_size, ")");
   auto input_length = logits_shape[1];
   auto logits_batch_size = logits_shape[0];
 
@@ -472,6 +476,10 @@ Status GreedySearchProcessLogits(
   // where input_length equals to parameters_->sequence_length for first subgraph call, and 1 for the remaining calls.
   const TensorShape& logits_shape = logits.Get<Tensor>().Shape();
   ORT_ENFORCE(logits_shape.NumDimensions() == 3);
+  const auto logits_vocab_size = static_cast<int>(logits_shape[2]);
+  ORT_RETURN_IF_NOT(vocab_size > 0 && vocab_size <= logits_vocab_size,
+                    "GreedySearch: vocab_size attribute (", vocab_size,
+                    ") must be positive and not exceed decoder logits width (", logits_vocab_size, ")");
   auto input_length = logits_shape[1];
 
   // Get logits for the last token:

onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc

@@ -362,6 +362,10 @@ Status ProcessLogits(const OrtValue& logits,                                 //
   // where input_length equals to parameters_->sequence_length for first subgraph call, and 1 for the remaining calls.
   const TensorShape& logits_shape = logits.Get<Tensor>().Shape();
   ORT_ENFORCE(logits_shape.NumDimensions() == 3);
+  const auto logits_vocab_size = static_cast<int>(logits_shape[2]);
+  ORT_RETURN_IF_NOT(vocab_size > 0 && vocab_size <= logits_vocab_size,
+                    "BeamSearch: vocab_size attribute (", vocab_size,
+                    ") must be positive and not exceed decoder logits width (", logits_vocab_size, ")");
   auto input_length = logits_shape[1];
   auto logits_batch_size = logits_shape[0];
 
@@ -849,6 +853,10 @@ Status GreedySearchProcessLogits(
   // where input_length equals to parameters_->sequence_length for first subgraph call, and 1 for the remaining calls.
   const TensorShape& logits_shape = logits.Get<Tensor>().Shape();
   ORT_ENFORCE(logits_shape.NumDimensions() == 3);
+  const auto logits_vocab_size = static_cast<int>(logits_shape[2]);
+  ORT_RETURN_IF_NOT(vocab_size > 0 && vocab_size <= logits_vocab_size,
+                    "GreedySearch: vocab_size attribute (", vocab_size,
+                    ") must be positive and not exceed decoder logits width (", logits_vocab_size, ")");
   auto input_length = logits_shape[1];
 
   // NOTE: `padded_vocab_size` MAY be different from `vocab_size`.

onnxruntime/test/contrib_ops/beam_search_test.cc

@@ -11,6 +11,7 @@
 #include "test/unittest_util/model_tester.h"
 #include "test/util/include/scoped_env_vars.h"
 #include "contrib_ops/cpu/transformers/generation_shared.h"
+#include "contrib_ops/cpu/transformers/beam_search_parameters.h"
 
 #ifdef USE_CUDA
 #include "core/providers/cuda/cuda_provider_options.h"
@@ -21,6 +22,33 @@ extern std::unique_ptr<Ort::Env> ort_env;
 namespace onnxruntime {
 namespace test {
 
+TEST(BeamSearchParametersTest, SetSubgraphParametersRejectsOversizedVocabSize) {
+  contrib::transformers::BeamSearchParameters parameters;
+  parameters.vocab_size = 150;
+
+  EXPECT_THROW(parameters.SetSubgraphParameters(128, 1, 1, 1), OnnxRuntimeException);
+}
+
+TEST(BeamSearchParametersTest, SetSubgraphParametersAllowsPaddedVocabSize) {
+  contrib::transformers::BeamSearchParameters parameters;
+  parameters.vocab_size = 64;
+
+  parameters.SetSubgraphParameters(128, 2, 4, 6);
+
+  EXPECT_EQ(parameters.vocab_size, 64);
+  EXPECT_EQ(parameters.num_heads, 2);
+}
+
+TEST(BeamSearchParametersTest, SetSubgraphParametersUsesSubgraphSizeWhenAttributeIsDefault) {
+  contrib::transformers::BeamSearchParameters parameters;
+  parameters.vocab_size = -1;
+
+  parameters.SetSubgraphParameters(128, 2, 4, 6);
+
+  EXPECT_EQ(parameters.vocab_size, 128);
+  EXPECT_EQ(parameters.num_heads, 2);
+}
+
 void RunGptBeamSearchFp32() {
   std::vector<int64_t> input_ids_shape{3, 12};
   std::vector<int32_t> input_ids{