Fix security vulnerabilities - #2224
Merged
Merged
Conversation
Copilot started reviewing on behalf of
Akshay Sonawane (apsonawane)
June 12, 2026 04:01
View session
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens sampling parameter handling by preventing top_k from exceeding the model’s vocab_size, avoiding potential out-of-bounds behavior during Top-K / Top-K+Top-P sampling.
Changes:
- Added generator initialization validation to reject
top_k > vocab_sizewith a detailedstd::runtime_error. - Clamped
ktovocab_sizeinside CPU Top-K and Top-K+Top-P sampling to preventpartial_sortfrom indexing past the vocabulary. - Added a unit test ensuring generator creation rejects invalid
top_kfor a known test model.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
src/generators.cpp |
Rejects invalid top_k during sampling method initialization to prevent unsafe configurations. |
src/search.cpp |
Caps CPU sampling k at vocab_size to prevent out-of-bounds in partial_sort-based Top-K selection. |
test/sampling_tests.cpp |
Adds regression coverage asserting top_k > vocab_size is rejected at generator creation. |
Akshay Sonawane (apsonawane)
enabled auto-merge (squash)
June 12, 2026 04:09
kunal-vaishnavi
approved these changes
Jun 12, 2026
kunal-vaishnavi
disabled auto-merge
June 18, 2026 21:26
David Fan (jiafatom)
added a commit
that referenced
this pull request
Jul 12, 2026
…yMinLength (#2266) ## Summary Fixes an out-of-bounds write (CWE-787 / CWE-129) reachable from a malicious `genai_config.json`. `Search_Cpu::ApplyMinLength` suppresses EOS tokens while the sequence is shorter than `min_length` by writing `std::numeric_limits<float>::lowest()` into the per-beam score row at each configured `eos_token_id`: ```cpp std::span<float> const beam_token_scores = GetScores(i); // vocab_size elements for (auto token_id : params_->config.model.eos_token_id) beam_token_scores[token_id] = std::numeric_limits<float>::lowest(); // unchecked index ``` `eos_token_id` comes straight from config (`src/config.cpp`) with only a `static_cast<int>` and is never bounded against `vocab_size`. A model whose `eos_token_id >= vocab_size` (or `< 0`) drives the store past the end of the `vocab_size`-sized `GetScores()` subspan — a heap-buffer-overflow write on the first decode step (reached via `Generator::GenerateNextToken` -> `ApplyMinLength`). The same pattern exists in the CUDA backend (`Search_Cuda::ApplyMinLength`). ## Changes - **`src/generators.cpp`**: reject any `eos_token_id` outside `[0, vocab_size)` in the `Generator::Generator` validator, mirroring the existing `top_k <= vocab_size` check (added in #2224). - **`src/search.cpp`** and **`src/cuda/search_cuda.cpp`**: defense-in-depth — skip out-of-range ids at the sink in both backends. - **`test/sampling_tests.cpp`**: add `EosTokenIdExceedsVocabSizeThrowsCpu`, which overlays an out-of-range `eos_token_id` and asserts `OgaGenerator::Create` throws instead of proceeding to the OOB write. ## Testing New regression test follows the existing `SamplingTests` overlay pattern (`tiny-random-gpt2-fp32` + `vocab_size`/`eos_token_id` overlay). Local build isn't available in my environment; relying on CI to build and run the C++ unit tests. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request adds stricter validation and handling for the
top_ksampling parameter to prevent it from exceeding the model's vocabulary size, which could previously lead to out-of-bounds errors. The changes ensure that both the generator initialization and sampling methods enforce this constraint, and a new test verifies the behavior.Sampling parameter validation and safety:
src/generators.cpp: Added a check inGenerator::InitializeSamplingMethodto throw a runtime error iftop_kis set higher thanvocab_size.src/search.cpp: UpdatedGreedySearch_Cpu::SampleTopKandGreedySearch_Cpu::SampleTopKTopPto clampktovocab_size, preventing out-of-bounds access during sampling. [1] [2]Testing:
test/sampling_tests.cpp: Added a test (TopKExceedingVocabSizeIsRejected) to ensure that settingtop_kgreater thanvocab_sizeis properly rejected with a runtime error.