Skip to content

Reduce CPU-side per-token overhead in GenerateNextToken and SampleTopP #2085

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
hanbitmyths merged 6 commits into from
Apr 20, 2026
Merged

Reduce CPU-side per-token overhead in GenerateNextToken and SampleTopP #2085

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
eeee379
Performance improvement
hanbitmyths Apr 14, 2026
fb0a270
Change the initial candidate count of top-p
hanbitmyths Apr 14, 2026
5b48c58
Resolve comments
hanbitmyths Apr 15, 2026
70788ca
Resolve lint errors
hanbitmyths Apr 15, 2026
8fdbe6a
coding convention
hanbitmyths Apr 16, 2026
054f0a6
Refactor InitPerTokenCache into InitializeSamplingMethod and Initiali…
hanbitmyths Apr 16, 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
111 changes: 68 additions & 43 deletions src/generators.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,7 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
// so skip the standard validations and just create the state.
if (ModelType::IsRNNT(model.config_->model.type)) {
state_ = model.CreateState({}, params);
is_nemotron_speech_model_ = dynamic_cast<NemotronSpeechState*>(state_.get()) != nullptr;
return;
}

Expand All @@ -367,6 +368,46 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
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

InitializePhi3RopeThreshold(params);
InitializeSamplingMethod(params);
}

void Generator::InitializePhi3RopeThreshold(const GeneratorParams& params) {
// TRT-RTX and DML EPs use a single rope factor for all tokens, so no ROPE rewind is needed.
const bool ep_uses_single_rope_factor = model_->p_device_->GetType() == DeviceType::NvTensorRtRtx ||
model_->p_device_->GetType() == DeviceType::DML;

// Phi3 ROPE factor rewind threshold: 4097 for phi3/phimoe, 8193 for phi3small, 0 otherwise
// TODO: Extend to support batch size > 1, num beams > 1, and multimodal models
const auto& model_type = model_->config_->model.type;
if (params.BatchBeamSize() == 1 && !ep_uses_single_rope_factor) {
if (model_type == "phi3" || model_type == "phimoe")
phi3_rope_threshold_ = 4097;
else if (model_type == "phi3small")
phi3_rope_threshold_ = 8193;
}
}

void Generator::InitializeSamplingMethod(const GeneratorParams& params) {
const auto& search = params.search;
if (!search.do_sample || search.top_k == 1 || search.temperature == 0) {
sampling_method_ = SamplingMethod::kGreedy;
} else {
if (search.num_beams != 1)
throw std::runtime_error("TopK and TopP cannot be used with a beam search");
if (search.top_p < 0.0f || search.top_p > 1.0f)
throw std::runtime_error("top_p must be between 0.0 and 1.0");
if (search.top_k < 0)
throw std::runtime_error("top_k must be 0 or greater");
if (search.top_p > 0.0f && search.top_p < 1.0f && search.top_k > 1) {
sampling_method_ = SamplingMethod::kTopKTopP;
} else if (search.top_k > 1) {
sampling_method_ = SamplingMethod::kTopK;
} else {
sampling_method_ = SamplingMethod::kTopP;
}
}
}

DeviceSpan<int32_t> Generator::AllocateInputIdsOnDevice(cpu_span<const int32_t> input_ids) {
Expand Down Expand Up @@ -511,18 +552,18 @@ void Generator::SetRuntimeOption(const char* key, const char* value) {
}

size_t Generator::TokenCount() const {
if (auto* speech_state = dynamic_cast<NemotronSpeechState*>(state_.get()))
return speech_state->TokenCount();
if (is_nemotron_speech_model_)
return static_cast<NemotronSpeechState*>(state_.get())->TokenCount();
return static_cast<size_t>(search_->GetSequenceLength());
}

bool Generator::IsDone() {
ThrowErrorIfSessionTerminated(state_->session_terminated_);

if (auto* speech_state = dynamic_cast<NemotronSpeechState*>(state_.get())) {
if (is_nemotron_speech_model_) {
// Pending mel input means we haven't started processing this chunk yet
if (!extra_inputs_.empty()) return false;
return speech_state->IsChunkDone();
return static_cast<NemotronSpeechState*>(state_.get())->IsChunkDone();
}

if (computed_logits_) {
Expand Down Expand Up @@ -556,31 +597,22 @@ void Generator::GenerateNextToken() {
ThrowErrorIfSessionTerminated(state_->session_terminated_);

// RNNT models: yield one token per call from the decoder state machine
if (auto* speech_state = dynamic_cast<NemotronSpeechState*>(state_.get())) {
if (is_nemotron_speech_model_) {
state_->SetExtraInputs(extra_inputs_);
extra_inputs_.clear();
speech_state->StepToken();
static_cast<NemotronSpeechState*>(state_.get())->StepToken();
return;
}

if (search_->GetSequenceLength() == 0 && !computed_logits_)
throw std::runtime_error("GenerateNextToken called with no prior state. Please call AppendTokens, SetLogits, or SetInputs before calling GenerateNextToken.");

// TRT-RTX and DML EPs use a single rope factor for all tokens: https://github.com/microsoft/onnxruntime-genai/blob/d5dc8cb02fd02b0dce99c6938449566371da0d28/src/python/py/models/builder.py#L1464-L1473
// TODO: change this when these EPs support multi rope factors
const bool epUsesSingleRopeFactor = model_->p_device_->GetType() == DeviceType::NvTensorRtRtx || model_->p_device_->GetType() == DeviceType::DML;

// TODO: Extend the solution to make it work for batch size > 1, num beams > 1, multimodal and DML
// Phi3 model switches from short factor to long factor at 4097 (original_max_position_embeddings+1) token, needs Recomputation of Position IDs and KV Cache
// at this stage which is achieved by rewinding to zero and appending the current sequence
// Scenarios where this solution works: Batch size = 1, Num beams = 1, decoder model, EP is either CPU or CUDA
// Scenarios where it doesn't work: Batch size > 1 OR Num beams > 1 OR Multimodal model (like phi3 vision) OR EP is DML
if (search_->params_->BatchBeamSize() == 1 && !epUsesSingleRopeFactor) {
if (((search_->GetSequenceLength() == 4097) && (model_->config_->model.type == "phi3" || model_->config_->model.type == "phimoe")) || ((search_->GetSequenceLength() == 8193) && (model_->config_->model.type == "phi3small"))) {
auto current_seq = cpu_span<int32_t>(GetSequence(0).CopyDeviceToCpu());
RewindToLength(0);
AppendTokens(current_seq);
}
// Phi3 model switches from short factor to long factor at the ROPE threshold token,
// needs recomputation of Position IDs and KV Cache via rewind + re-append.
if (phi3_rope_threshold_ != 0 && search_->GetSequenceLength() == phi3_rope_threshold_) {
auto current_seq = cpu_span<int32_t>(GetSequence(0).CopyDeviceToCpu());
RewindToLength(0);
AppendTokens(current_seq);
}

if (!computed_logits_) {
Expand Down Expand Up @@ -609,28 +641,21 @@ void Generator::GenerateNextToken() {
}

last_action_ = Action::generated;
if (!search.do_sample || search.top_k == 1 || search.temperature == 0) {
search_->SelectTop();
return;
}

// The user explicitly called TopK_TopP on a beam search
if (search.num_beams != 1)
throw std::runtime_error("TopK and TopP cannot be used with a beam search");

// Sanity checks
if (search.top_p < 0.0f || search.top_p > 1.0f)
throw std::runtime_error("top_p must be between 0.0 and 1.0");
if (search.top_k < 0)
throw std::runtime_error("top_k must be 0 or greater");

if (search.top_p > 0.0f && search.top_p < 1.0f && search.top_k > 1) {
search_->SampleTopKTopP(search.top_k, search.top_p, search.temperature);
} else if (search.top_k > 1) {
search_->SampleTopK(search.top_k, search.temperature);
} else {
assert(search.top_k == 0);
search_->SampleTopP(search.top_p, search.temperature);
switch (sampling_method_) {
case SamplingMethod::kGreedy:
search_->SelectTop();
return;
case SamplingMethod::kTopKTopP:
search_->SampleTopKTopP(search.top_k, search.top_p, search.temperature);
return;
case SamplingMethod::kTopK:
search_->SampleTopK(search.top_k, search.temperature);
return;
case SamplingMethod::kTopP:
search_->SampleTopP(search.top_p, search.temperature);
return;
default:
throw std::runtime_error("Unknown sampling method");
}
}

Expand Down
11 changes: 11 additions & 0 deletions src/generators.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,17 @@ struct Generator : LeakChecked<Generator> {
generated, // Set after GenerateNextToken
rewound }; // Set after RewindToLength
Action last_action_{standard};

// Pre-computed per-token decisions: avoid repeated checks each token
bool is_nemotron_speech_model_{};
int phi3_rope_threshold_{}; // 0 means no ROPE rewind needed
enum class SamplingMethod { kGreedy,
kTopK,
kTopP,
kTopKTopP };
SamplingMethod sampling_method_{SamplingMethod::kGreedy};
void InitializeSamplingMethod(const GeneratorParams& params);
void InitializePhi3RopeThreshold(const GeneratorParams& params);
};

struct OrtGlobals {
Expand Down
117 changes: 63 additions & 54 deletions src/search.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,34 +192,36 @@ void GreedySearch_Cpu::SampleTopK(int k, float temperature) {
AppendNextTokensToSequences();
}

// Adaptive partial-sort nucleus (top-p) search.
// Find the minimal nucleus of tokens whose cumulative probability >= p using
// adaptive partial sort with a log-space tail bound. This avoids computing the
// global softmax partition function (O(V) exp() calls) by bounding the unsorted
// tail probability at each step.
//
// Given a score array and precomputed global partition stats (max_score, inv_temp,
// inv_global_exp_sum), finds the minimal set of tokens whose cumulative true
// probability >= p using adaptive partial_sort with geometrically growing K.
// At sorted position i, all V-i-1 remaining elements have score <= scores[indices[i]]
// (by sort order within the prefix, and by partial_sort guarantee beyond it). So:
// tail_sum <= (V - i - 1) * exp((scores[indices[i]] - max_score) * inv_temp)
// The cumulative probability lower bound is:
// prefix_sum / (prefix_sum + tail_bound)
// When this exceeds p, the true cumulative (which is >= this bound) also exceeds p.
//
// Returns the number of tokens in the nucleus (cutoff_index).
// On return, indices[0..cutoff_index-1] are the nucleus token indices sorted by
// descending score.
//
// Complexity (V = vocab_size, typically 128K-200K):
// - Best case (common): O(V log K0) where K0=256 — covers p>=0.99 for most LLM distributions
// - Rare fallback: O(V log 1024) — very flat distributions
// - Worst case: O(V log V) — uniform distribution (essentially never)
// The nucleus found may include a few extra tokens compared to exact global softmax,
// but is always a valid nucleus (true cumulative prob >= p). For typical peaked LLM
// distributions, the bound is tight and the result is identical.
static int FindNucleus(std::span<const float> scores, std::span<int32_t> indices,
float p, float max_score, float inv_temp, float inv_global_exp_sum) {
float p, float max_score, float inv_temp) {
const int vocab_size = static_cast<int>(scores.size());

// 256 covers p>=0.99 for most LLM distributions. 4x geometric growth handles
// the rare case of unusually flat distributions without over-sorting.
constexpr int kInitialCandidateCount = 256;
// Start small (16) to minimize wasted work when the top few tokens dominate
// (common at low temperature). Grow 4× to avoid O(V log V) full-sort fallback
// that a 2× growth with a hard cap would require for flat distributions.
constexpr int kInitialCandidateCount = 16;
constexpr int kCandidateCountGrowthFactor = 4;

std::iota(indices.begin(), indices.end(), 0);

int sorted_count = 0;
int k = std::min(kInitialCandidateCount, vocab_size);
float cumulative_prob = 0.0f;
float prefix_exp_sum = 0.0f;
int cutoff_index = 0;

while (k <= vocab_size) {
Expand All @@ -228,53 +230,69 @@ static int FindNucleus(std::span<const float> scores, std::span<int32_t> indices
std::partial_sort(indices.begin() + sorted_count, indices.begin() + k, indices.end(),
[&scores](int32_t i, int32_t j) { return scores[i] > scores[j]; });

// Accumulate true probabilities for the newly sorted elements [sorted_count, k).
// Note: We recompute exp() here rather than caching a V-sized array because
// cutoff_index is typically << V (e.g., 10-50 tokens for p=0.95), so the few
// recomputed exp() calls are far cheaper than allocating and filling a ~800KB
// float array per batch item. The Softmax below (Step 3) only operates on the
// nucleus, not the full vocab.
// Accumulate exp() for newly sorted elements and check the tail bound.
cutoff_index = k;
for (int i = sorted_count; i < k; ++i) {
cumulative_prob += std::exp((scores[indices[i]] - max_score) * inv_temp) * inv_global_exp_sum;
if (cumulative_prob >= p) {
float exp_i = std::exp((scores[indices[i]] - max_score) * inv_temp);
prefix_exp_sum += exp_i;

// Upper bound on remaining elements: all V-i-1 unsorted entries have score <= current.
float tail_bound = static_cast<float>(vocab_size - i - 1) * exp_i;
if (prefix_exp_sum >= p * (prefix_exp_sum + tail_bound)) {
cutoff_index = i + 1;
break;
}
}
sorted_count = k;

if (cumulative_prob >= p || k == vocab_size)
if (cutoff_index < k || k == vocab_size)
break;

k = std::min(k * kCandidateCountGrowthFactor, vocab_size);
}

// When the vocabulary is small enough that the tail bound may be loose,
// compute the exact cutoff using the true partition function.
// For large V this path is never reached (tail bound is tight for peaked distributions).
if (vocab_size <= kInitialCandidateCount * kCandidateCountGrowthFactor) {
// Ensure all elements are sorted for exact computation
if (sorted_count < vocab_size) {
std::partial_sort(indices.begin() + sorted_count, indices.begin() + vocab_size, indices.end(),
[&scores](int32_t i, int32_t j) { return scores[i] > scores[j]; });
}
float global_exp_sum = 0.0f;
for (int i = 0; i < vocab_size; ++i) {
global_exp_sum += std::exp((scores[indices[i]] - max_score) * inv_temp);
}
float cum = 0.0f;
for (int i = 0; i < vocab_size; ++i) {
cum += std::exp((scores[indices[i]] - max_score) * inv_temp);
if (cum >= p * global_exp_sum) {
return i + 1;
}
}
}

return cutoff_index;
}

// Top-P (nucleus) sampling using adaptive partial sort.
// Top-P (nucleus) sampling using adaptive partial sort with tail-bound
// early termination.
//
// Instead of fully sorting the entire vocabulary O(V log V), this uses
// FindNucleus() which applies std::partial_sort with an adaptive K that starts
// small and grows geometrically only when needed.
// Instead of a separate O(V) pass to compute the global softmax partition
// function (max + sum-of-exp over the full vocabulary), FindNucleus uses
// a conservative upper bound on the unsorted tail to determine when the
// nucleus is complete. This eliminates ~V calls to std::exp() per token
// (e.g., ~128K for typical LLM vocabularies).
//
// Correctness: A single O(V) pass computes the global softmax partition function
// (max + sum-of-exp) before sorting. Cumulative probabilities are then measured
// against the true full-vocab distribution, ensuring the nucleus contains exactly
// the minimal set of tokens whose true probabilities sum to >= p.
//
// Additional optimizations over the previous implementation:
// - Buffers (indices, top_logits) are allocated once outside the batch loop,
// eliminating ~800KB of heap allocation per batch item per decode step.
// - Softmax for final sampling is computed only over the nucleus (not the full vocab).
// - discrete_distribution is constructed over the nucleus (cutoff_index elements) instead
// of the full vocabulary, avoiding an O(V) CDF construction on zeroed-out entries.
// Additional optimizations:
// - Buffers (indices, top_logits) are allocated once outside the batch loop.
// - Softmax for final sampling is computed only over the nucleus.
// - discrete_distribution is constructed over the nucleus only.
void GreedySearch_Cpu::SampleTopP(float p, float temperature) {
const int vocab_size = params_->config.model.vocab_size;

// Buffers allocated once outside the batch loop to avoid per-call heap allocations.
// For vocab_size=200K, this saves ~800KB of malloc/free per batch item per token step.
std::vector<int32_t> indices(vocab_size);
std::vector<float> top_logits;

Expand All @@ -285,22 +303,13 @@ void GreedySearch_Cpu::SampleTopP(float p, float temperature) {

std::span<float> scores = next_token_scores_.CpuSpan().subspan(batch_id * vocab_size, vocab_size);

// Step 1: Compute the global partition function (one O(V) pass) so that
// cumulative probabilities measured over the sorted prefix are relative to
// the true full-vocab distribution, not a renormalized subset.
float inv_temp = 1.0f / temperature;
float max_score = *std::max_element(scores.begin(), scores.end());
float global_exp_sum = 0.0f;
for (int i = 0; i < vocab_size; ++i) {
global_exp_sum += std::exp((scores[i] - max_score) * inv_temp);
}
float inv_global_exp_sum = 1.0f / global_exp_sum;

// Step 2: Find the nucleus via adaptive partial sort.
int cutoff_index = FindNucleus(scores, indices, p, max_score, inv_temp, inv_global_exp_sum);
// Find the nucleus via adaptive partial sort with tail-bound early termination.
int cutoff_index = FindNucleus(scores, indices, p, max_score, inv_temp);

// Step 3: Re-normalize the nucleus probabilities and sample a token.
// Build softmax over only the cutoff_index elements for proper sampling.
// Re-normalize the nucleus probabilities and sample a token.
top_logits.resize(cutoff_index);
for (int i = 0; i < cutoff_index; ++i) {
top_logits[i] = scores[indices[i]] * inv_temp;
Expand Down
Loading
Loading