Skip to content

Validate eos_token_id against vocab_size to prevent OOB write in ApplyMinLength #2266

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
David Fan (jiafatom) merged 4 commits into from
Jul 12, 2026
Merged

Validate eos_token_id against vocab_size to prevent OOB write in ApplyMinLength #2266

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
46efa7b
Validate eos_token_id against vocab_size to prevent OOB write
jiafatom Jul 6, 2026
a00cfcd
Address review: assert runtime_error in test and clarify range message
jiafatom Jul 6, 2026
6f6725b
Address review: drop per-inference eos_token_id checks (validate once…
jiafatom Jul 7, 2026
d4ef5ea
Fix SamplingTests broken by eos_token_id vocab-bounds validation
jiafatom Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/generators.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,6 +405,14 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
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) + ")");

// eos_token_id values are used directly as indices into the per-token score
// row (of size vocab_size), e.g. in Search::ApplyMinLength. An out-of-range
// value would cause an out-of-bounds write, so reject it here.
for (auto eos_token_id : params.config.model.eos_token_id) {
if (eos_token_id < 0 || eos_token_id >= params.config.model.vocab_size)
throw std::runtime_error("eos_token_id (" + std::to_string(eos_token_id) + ") must be in range [0, " + std::to_string(params.config.model.vocab_size) + ") (vocab_size)");
}

search_ = CreateSearch(params);
state_ = model.CreateState(search_->GetSequenceLengths(), params); // Search sequence lengths set when creating state
guidance_logits_processor_ = CreateGuidanceLogitsProcessor(*state_); // Could be nullptr if use_guidance (constrained decoding) is not used
Expand Down
37 changes: 29 additions & 8 deletions test/sampling_tests.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ TEST(SamplingTests, BatchedSamplingTopPCpu) {
0.1f, 0.1f, 0.1f, 0.1f, 0.6f};

auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
config->Overlay(R"({ "model": { "vocab_size" : 5 } })");
config->Overlay(R"({ "model": { "vocab_size" : 5, "eos_token_id" : 0 } })");

auto model = OgaModel::Create(*config);
auto params = OgaGeneratorParams::Create(*model);
Expand Down Expand Up @@ -56,7 +56,7 @@ TEST(SamplingTests, BatchedSamplingTopKCpu) {
1.25f, 0.25f, 1.5f, 0.25f, 2.0f};

auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
config->Overlay(R"({ "model": { "vocab_size" : 5 } })");
config->Overlay(R"({ "model": { "vocab_size" : 5, "eos_token_id" : 0 } })");

int batch_size = 4;

Expand Down Expand Up @@ -102,6 +102,27 @@ TEST(SamplingTests, BeamSearchVocabSizeTooSmallThrowsCpu) {
}
}

// Regression test: an eos_token_id outside [0, vocab_size) must be rejected at
// generator creation time instead of causing an out-of-bounds write in
// Search_Cpu::ApplyMinLength (which uses eos_token_id as a score-row index).
TEST(SamplingTests, EosTokenIdExceedsVocabSizeThrowsCpu) {
auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
config->Overlay(R"({ "model": { "vocab_size" : 5, "eos_token_id" : 5 } })");

auto model = OgaModel::Create(*config);
auto params = OgaGeneratorParams::Create(*model);
params->SetSearchOption("max_length", 10);
params->SetSearchOption("min_length", 5);

try {
OgaGenerator::Create(*model, *params);
FAIL() << "Expected std::runtime_error for eos_token_id >= vocab_size";
} catch (const std::runtime_error& e) {
EXPECT_NE(std::string(e.what()).find("eos_token_id"), 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,
Expand All @@ -110,7 +131,7 @@ TEST(SamplingTests, BatchedSamplingTopPAndKCpu) {
1.25f, 0.25f, 1.5f, 0.25f, 2.0f};

auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
config->Overlay(R"({ "model": { "vocab_size" : 5 } })");
config->Overlay(R"({ "model": { "vocab_size" : 5, "eos_token_id" : 0 } })");

int batch_size = 4;

Expand Down Expand Up @@ -244,7 +265,7 @@ void Softmax(std::span<float> scores, float temperature) {
void RunSamplingTest(int batch_size, int k, float p, int vocab_size, int num_iter, float temperature, bool use_cuda) {
// --- 1. Setup Model and Generator Parameters ---
auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
std::string overlay_json = R"({ "model": { "vocab_size" : )" + std::to_string(vocab_size) + R"( } })";
std::string overlay_json = R"({ "model": { "vocab_size" : )" + std::to_string(vocab_size) + R"(, "eos_token_id" : 0 } })";
config->Overlay(overlay_json.c_str());

if (use_cuda) {
Expand Down Expand Up @@ -506,7 +527,7 @@ TEST(SamplingTests, BatchedSamplingTopPCuda) {
int vocab_size = 5;

auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
config->Overlay(R"({ "model": { "vocab_size" : 5 } })");
config->Overlay(R"({ "model": { "vocab_size" : 5, "eos_token_id" : 0 } })");
config->ClearProviders();
config->AppendProvider("cuda");
auto model = OgaModel::Create(*config);
Expand Down Expand Up @@ -536,7 +557,7 @@ TEST(SamplingTests, BatchedSamplingTopKCuda) {
int vocab_size = 5;

auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
config->Overlay(R"({ "model": { "vocab_size" : 5 } })");
config->Overlay(R"({ "model": { "vocab_size" : 5, "eos_token_id" : 0 } })");
config->ClearProviders();
config->AppendProvider("cuda");
auto model = OgaModel::Create(*config);
Expand Down Expand Up @@ -569,7 +590,7 @@ TEST(SamplingTests, BatchedSamplingTopPAndKCuda) {
int vocab_size = 5;

auto config = OgaConfig::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
config->Overlay(R"({ "model": { "vocab_size" : 5 } })");
config->Overlay(R"({ "model": { "vocab_size" : 5, "eos_token_id" : 0 } })");
config->ClearProviders();
config->AppendProvider("cuda");
auto model = OgaModel::Create(*config);
Expand Down Expand Up @@ -728,7 +749,7 @@ struct NvTensorRtRtxTestSetup {

// Create config with vocab_size overlay
auto config = OgaConfig::Create(resolved_path.c_str());
std::string overlay = R"({ "model": { "vocab_size" : )" + std::to_string(vocab_size) + R"( } })";
std::string overlay = R"({ "model": { "vocab_size" : )" + std::to_string(vocab_size) + R"(, "eos_token_id" : 0 } })";
config->Overlay(overlay.c_str());
config->ClearProviders();
config->AppendProvider("NvTensorRtRtx");
Expand Down
Loading