Reject vocab_size < 2 for beam search to fix OOB in BeamSearch_Cpu::SelectTop

src/generators.cpp

@@ -398,6 +398,12 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
     throw std::runtime_error("num_beams (" + std::to_string(params.search.num_beams) + ") must be in [1, " + std::to_string(max_num_beams) + "]");
   if (params.config.model.vocab_size < 1)
     throw std::runtime_error("vocab_size must be 1 or greater, is " + std::to_string(params.config.model.vocab_size));
+  // Beam search selects the top 2*num_beams (beam, token) candidates out of
+  // num_beams*vocab_size entries in BeamSearch_Cpu::SelectTop, which requires
+  // num_beams*vocab_size >= 2*num_beams, i.e. vocab_size >= 2. A smaller
+  // vocabulary would drive an out-of-bounds partial_sort.
+  if (params.search.num_beams > 1 && params.config.model.vocab_size < 2)
+    throw std::runtime_error("vocab_size (" + std::to_string(params.config.model.vocab_size) + ") must be 2 or greater when using beam search (num_beams=" + std::to_string(params.search.num_beams) + ")");
 
   search_ = CreateSearch(params);
   state_ = model.CreateState(search_->GetSequenceLengths(), params);    // Search sequence lengths set when creating state

src/search.cpp

@@ -117,7 +117,12 @@ void BeamSearch_Cpu::SelectTop() {
   // Use partial_sort to find only the top 2*num_beams elements per batch,
   // instead of heapifying the entire vocab*beams array via priority_queue.
   const size_t total_elements = static_cast<size_t>(params_->search.num_beams) * params_->config.model.vocab_size;
-  assert(total_elements >= top_k);
+  // Defense-in-depth: a plain assert() is compiled out under NDEBUG (release
+  // builds), so enforce the invariant that partial_sort relies on at runtime.
+  // This is normally guaranteed by the vocab_size >= 2 validation for beam
+  // search in Generator::Generator.
+  if (total_elements < top_k)
+    throw std::runtime_error("Beam search requires num_beams * vocab_size (" + std::to_string(total_elements) + ") to be at least 2 * num_beams (" + std::to_string(top_k) + "); vocab_size is too small");
 
   // Reuse class member to avoid re-allocating on every call (size is constant).
   select_top_idx_.resize(total_elements);

test/sampling_tests.cpp

@@ -81,6 +81,27 @@ TEST(SamplingTests, BatchedSamplingTopKCpu) {
   }
 }
 
+// Regression test: beam search with a vocab_size too small to supply
+// 2*num_beams candidates must be rejected at generator creation instead of
+// driving an out-of-bounds partial_sort in BeamSearch_Cpu::SelectTop.
+TEST(SamplingTests, BeamSearchVocabSizeTooSmallThrowsCpu) {
+  auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  config->Overlay(R"({ "model": { "vocab_size" : 1 } })");
+
+  auto model = OgaModel::Create(*config);
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", 10);
+  params->SetSearchOption("num_beams", 2);
+
+  try {
+    OgaGenerator::Create(*model, *params);
+    FAIL() << "Expected std::runtime_error for beam search with vocab_size < 2";
+  } catch (const std::runtime_error& e) {
+    EXPECT_NE(std::string(e.what()).find("vocab_size"), std::string::npos)
+        << "Unexpected error message: " << e.what();
+  }
+}
+
 TEST(SamplingTests, BatchedSamplingTopPAndKCpu) {
   std::vector<int32_t> input_ids{0, 1, 2, 3};
   std::vector<float> logits_cpu{2.0f, 1.5f, 1.25f, 0.25f, 0.25f,