Skip to content

Fix AppendNextTokensToSequences heap overflow - #2111

Merged
baijumeswani merged 10 commits into
mainfrom
asonawane/heap
Jul 16, 2026
Merged

Fix AppendNextTokensToSequences heap overflow#2111
baijumeswani merged 10 commits into
mainfrom
asonawane/heap

Conversation

@apsonawane

@apsonawane apsonawane commented Apr 30, 2026

Copy link
Copy Markdown
Contributor

Fix heap-buffer-overflow in GenerateNextToken when sequence is at max_length

Summary

Adds a fast-fail guard at the entry of Generator::GenerateNextToken() that throws a clear std::runtime_error when the sequence has already reached max_length. This prevents a heap-buffer-overflow write into the preallocated sequences buffer that is otherwise reachable via the C API.

Reported bug (AddressSanitizer)

A libFuzzer harness driving the OgaGenerator_* C API triggered the following crash:

==ERROR: AddressSanitizer: heap-buffer-overflow ... WRITE of size 4
  #0 Generators::GreedySearch_Cpu::AppendNextTokensToSequences()  src/search.cpp:323
  #1 Generators::GreedySearch_Cpu::SampleTopK(int, float)         src/search.cpp:192
  #2 Generators::Generator::GenerateNextToken()                   src/generators.cpp:602
  #3 OgaGenerator_GenerateNextToken                               src/ort_genai_c.cpp:473
  #4 <fuzzer harness>

0 bytes after 64-byte region  ← writes at index max_length
allocated by Generators::Sequences::Sequences(...) at construction

Root cause

Inside GreedySearch_Cpu::AppendTokens, ResetDone() is called unconditionally at the end. If the previous call filled the sequences buffer up to exactly max_length, done_ is cleared. A subsequent GenerateNextToken() then:

  1. Enters SampleTopK().
  2. The !done_ check falsely passes (buffer is actually full).
  3. AppendNextTokensToSequences() writes at index max_length → OOB write into the redzone immediately following the 64-byte sequences buffer.

Fix

Add an early runtime error at the top of Generator::GenerateNextToken():

if (search_->GetSequenceLength() >= state_->params_->search.max_length)
  throw std::runtime_error(
      "GenerateNextToken called with sequence length already at max_length (" +
      std::to_string(state_->params_->search.max_length) + ")");

This sits directly above frame #2 of the crash stack, so the OOB path in frames #0 and #1 is unreachable. The exception propagates through OgaGenerator_GenerateNextToken and is surfaced to the caller as a normal API error instead of memory corruption.

Copilot AI review requested due to automatic review settings April 30, 2026 20:54

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR aims to prevent out-of-bounds writes into the preallocated sequence buffers during generation by adding explicit max_length bounds checks across CPU and CUDA search implementations, and by failing fast when GenerateNextToken() is called after reaching max_length.

Changes:

  • Added early-return bounds checks in CPU greedy/beam append paths to prevent writing past the sequences buffer.
  • Added CUDA-side guards to avoid launching append kernels once the sequence length reaches max_length, and updated done/max-length checks.
  • Added an early GenerateNextToken() runtime error when sequence length is already at max_length.

Reviewed changes

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

File Description
src/search.cpp Adds CPU-side bounds checks and adjusts done/reset behavior around appending tokens.
src/generators.cpp Adds a fast-fail guard in GenerateNextToken() when already at max_length.
src/cuda/search_cuda.cpp Adds CUDA-side conditional guards to avoid appending once at/over max_length, plus logging/done handling.

Comment thread src/generators.cpp Outdated
Comment thread src/search.cpp Outdated
@apsonawane
apsonawane requested a review from a team as a code owner May 5, 2026 19:05
@kunal-vaishnavi

Copy link
Copy Markdown
Contributor

The changes in this PR seem to indicate a larger issue inside ORT GenAI. The done state should reflect the true state, and we should not have to insert max length checks in so many locations. Do we know why this issue is happening in the first place?

Comment thread src/cuda/search_cuda.cpp Outdated
@apsonawane
apsonawane enabled auto-merge (squash) May 20, 2026 21:45
@baijumeswani
baijumeswani disabled auto-merge July 16, 2026 16:23
@baijumeswani
baijumeswani merged commit 50bb2ab into main Jul 16, 2026
62 of 67 checks passed
@baijumeswani
baijumeswani deleted the asonawane/heap branch July 16, 2026 16:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants