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

[David Fan (jiafatom)]

Summary

Fixes a heap out-of-bounds read/write (CWE-787 / CWE-125) reachable from a malicious genai_config.json.

BeamSearch_Cpu::SelectTop builds an index array of total_elements = num_beams * vocab_size entries and partial_sorts it to surface the top top_k = 2 * num_beams candidates:

const size_t top_k = 2 * params_->search.num_beams;
const size_t total_elements = num_beams * vocab_size;
assert(total_elements >= top_k);              // compiled out under NDEBUG
select_top_idx_.resize(total_elements);
std::partial_sort(begin, begin + top_k, end, ...);  // middle past end when top_k > total_elements

When vocab_size == 1 (and num_beams >= 2), total_elements = num_beams < top_k = 2*num_beams, so begin + top_k points past the end of select_top_idx_. std::partial_sort heapifies [begin, begin+top_k), reading and move-writing slots beyond the allocation. The only guard was an assert, which is compiled out in the shipped release builds (NDEBUG). The path is reached from the normal generation loop (GenerateNextToken -> SelectTop).

Fix

• src/generators.cpp: reject beam search (num_beams > 1) with vocab_size < 2 in the Generator::Generator validator (primary fix, at creation). total_elements >= top_k reduces exactly to vocab_size >= 2.
• src/search.cpp: replace the compiled-out assert(total_elements >= top_k) with a runtime check that throws, so release builds are protected as defense-in-depth.
• test/sampling_tests.cpp: add BeamSearchVocabSizeTooSmallThrowsCpu, overlaying vocab_size = 1 with num_beams = 2 and asserting OgaGenerator::Create throws.

Note: the greedy path (GreedySearch_Cpu::SelectTop, num_beams == 1) uses std::max_element and is unaffected.

Testing

New regression test follows the existing SamplingTests overlay pattern. Local build isn't available in my environment (root-owned build dir / jiafa-dev container needed); relying on CI to build and run the C++ unit tests.

[Copilot]

Pull request overview

This PR hardens beam search configuration validation to prevent an out-of-bounds access in BeamSearch_Cpu::SelectTop when a malicious (or invalid) genai_config.json sets vocab_size too small for the top_k = 2 * num_beams selection logic.

Changes:

• Add generator-creation validation to reject num_beams > 1 when vocab_size < 2.
• Add a runtime invariant check in BeamSearch_Cpu::SelectTop (replacing a release-stripped assert) to prevent std::partial_sort from operating with an invalid middle iterator.
• Add a regression test asserting OgaGenerator::Create throws for num_beams = 2 with vocab_size = 1.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
src/generators.cpp	Adds upfront validation to reject beam search configs with vocab_size < 2.
src/search.cpp	Adds defense-in-depth runtime check guarding partial_sort against top_k > total_elements.
test/sampling_tests.cpp	Adds regression coverage for the invalid vocab_size + beam search combination.