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

src/generators.cpp

@@ -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;
   }
 
@@ -367,6 +368,44 @@ 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
+
+  InitPerTokenCache(params);
+}
+
+void Generator::InitPerTokenCache(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;
+  }
+
+  // Pre-determine sampling method and validate parameters at construction
+  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) {
@@ -511,18 +550,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_) {
@@ -556,31 +595,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_) {
@@ -609,28 +639,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");
   }
 }
 

src/generators.h

@@ -130,6 +130,16 @@ 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 InitPerTokenCache(const GeneratorParams& params);
 };
 
 struct OrtGlobals {

src/search.cpp

@@ -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) {
@@ -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;
 
@@ -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;