From 78acea648016bdc67dd60a60b43dc6fbcbbc9418 Mon Sep 17 00:00:00 2001 From: David Fan Date: Mon, 6 Jul 2026 18:03:33 +0000 Subject: [PATCH] Reject vocab_size < 2 for beam search to prevent OOB in SelectTop BeamSearch_Cpu::SelectTop partial_sorts an index array of num_beams*vocab_size entries and asks for the top 2*num_beams (top_k). When vocab_size == 1 (and num_beams >= 2), total_elements < top_k, so the partial_sort middle iterator points past the end of select_top_idx_, causing an out-of-bounds read/write (heap-buffer-overflow, CWE-787/CWE-125). The only guard was an assert() that is compiled out under NDEBUG in release builds. Fix: - Reject beam search (num_beams > 1) with vocab_size < 2 in the Generator::Generator validator. - Replace the compiled-out assert in SelectTop with a runtime check that throws, so release builds are protected as defense-in-depth. - Add a CPU regression test overlaying vocab_size = 1 with num_beams = 2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/generators.cpp | 6 ++++++ src/search.cpp | 7 ++++++- test/sampling_tests.cpp | 21 +++++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/src/generators.cpp b/src/generators.cpp index 30d4e9844e..c6798b54f0 100644 --- a/src/generators.cpp +++ b/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 diff --git a/src/search.cpp b/src/search.cpp index b8ca196e60..d401ace919 100644 --- a/src/search.cpp +++ b/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(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); diff --git a/test/sampling_tests.cpp b/test/sampling_tests.cpp index afca86345a..ff172857f2 100644 --- a/test/sampling_tests.cpp +++ b/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 input_ids{0, 1, 2, 3}; std::vector logits_cpu{2.0f, 1.5f, 1.25f, 0.25f, 0.25f,