perf(qwen-flash): prefetch n-gram rows across cache checkpoints

src/cli/serve/generation_metrics.hpp

@@ -188,6 +188,11 @@ inline std::string GenerationLogDetails(
   if (result.lookup_tokens > 0)
     out << " lookup_accepted=" << result.lookup_accepted_tokens
         << " lookup_proposed=" << result.lookup_tokens;
+  if (result.ngram_rows > 0)
+    out << " ngram_rows=" << result.ngram_rows
+        << " ngram_dedup_rows=" << result.ngram_dedup_rows
+        << " ngram_cache_hits=" << result.ngram_cache_hits
+        << " ngram_disk_reads=" << result.ngram_disk_reads;
   if (result.cache_snapshot_bytes > 0)
     out << " cache_snapshot_bytes=" << result.cache_snapshot_bytes;
   if (result.cache_disk_queued_bytes > 0)

src/cli/serve/inference_backend.cpp

@@ -2883,14 +2883,17 @@ class QwenFlashNextTextRunner final : public HipTextModelRunner {
     const std::size_t next_position = offset + consumed;
     const auto prefix = QwenFlashNextEngineTokens(prompt.first(next_position));
     // The next step most likely takes the same budget; its n-gram rows are
-    // read during this one and used only if it does.
-    const std::size_t after = prompt.size() - next_position;
-    const auto next = QwenFlashNextEngineTokens(prompt.subspan(
-        next_position,
-        std::min<std::size_t>({max_input_tokens, after,
-                               after <= model_->PrefillThroughCapacity()
-                                   ? model_->PrefillThroughCapacity()
-                                   : model_->PrefillCapacity()})));
+    // read during this one and used for the tokens it shares. A step ending
+    // at a cache frontier continues into the pool's lookahead.
+    const auto rest = next_position == prompt.size()
+                          ? state.prefill_lookahead()
+                          : prompt.subspan(next_position);
+    const std::size_t after = rest.size();
+    const auto next = QwenFlashNextEngineTokens(rest.subspan(
+        0, std::min<std::size_t>({max_input_tokens, after,
+                                  after <= model_->PrefillThroughCapacity()
+                                      ? model_->PrefillThroughCapacity()
+                                      : model_->PrefillCapacity()})));
     std::string error;
     if (!qfn.session().Sync(prefix, &error, next)) {
       qfn.set_position(0);
@@ -3155,6 +3158,14 @@ class QwenFlashNextTextRunner final : public HipTextModelRunner {
     return RequireQwenFlashNextState(state).position();
   }
 
+  [[nodiscard]] TextRunnerNgramRows NgramRows(
+      const TextRunnerState& state) const override {
+    const auto totals =
+        RequireQwenFlashNextState(state).session().NgramRowTotals();
+    return {totals.rows, totals.dedup_rows, totals.cache_hits,
+            totals.disk_reads};
+  }
+
 private:
   mutable models::qwen::vision::ImageCache image_cache_;
   std::shared_ptr<QwenFlashNextModel> model_;

src/cli/serve/text_generation_backend.hpp

@@ -209,6 +209,12 @@ class TextGenerationBackend {
     std::size_t draft_accepted_tokens{0};
     std::size_t lookup_tokens{0};
     std::size_t lookup_accepted_tokens{0};
+    /// Hashed n-gram embedding rows gathered for this request, split by how
+    /// each was served.
+    std::size_t ngram_rows{0};
+    std::size_t ngram_dedup_rows{0};
+    std::size_t ngram_cache_hits{0};
+    std::size_t ngram_disk_reads{0};
     std::size_t prefill_tokens{0};
     std::size_t prefill_chunks{0};
     std::size_t active_decode_prefill_chunks{0};

src/cli/serve/text_generation_scheduler.cpp

@@ -953,6 +953,12 @@ struct TextGenerationScheduler::Impl {
     // A stop may land on a preview, an unadvanced AR token, or inside an
     // accepted speculative block. Retain only correctly labelled executed
     // state, using the same cache reconciliation as an interrupted request.
+    // Read before the commit returns the leased state.
+    const auto ngram = request->runner_request.ngram_rows();
+    request->result.ngram_rows = ngram.rows;
+    request->result.ngram_dedup_rows = ngram.dedup_rows;
+    request->result.ngram_cache_hits = ngram.cache_hits;
+    request->result.ngram_disk_reads = ngram.disk_reads;
     const auto cache_commit =
         finish_reason == TextGenerationBackend::FinishReason::kStopSequence
             ? request->runner_request.Cancel()

src/cli/serve/text_model_runner.cpp

@@ -873,6 +873,7 @@ struct TextRunnerPool::Request::Impl {
           state,
           std::span<const TextRunnerToken>(prompt).first(prefill_offset));
     }
+    ngram_baseline = runner->NgramRows(state);
     const auto capabilities = runner->Descriptor().capabilities;
     // Component adapters supply their own pass boundaries below. Building the
     // generic plan first only repeats prefix lookups before discarding it.
@@ -1232,6 +1233,8 @@ struct TextRunnerPool::Request::Impl {
   std::vector<TextRunnerToken> generated;
   std::size_t prefill_offset{0};
   bool decode_ready{false};
+  /// The leased state's cumulative n-gram rows before this request ran.
+  TextRunnerNgramRows ngram_baseline;
   // A throwing model call may have partially mutated device state. Never
   // publish that state as a live continuation, even if its position is stale.
   bool state_reusable{true};
@@ -1328,6 +1331,16 @@ double TextRunnerPool::Request::cache_restore_ms() const noexcept {
   return impl_ != nullptr ? impl_->lease.restore_ms() : 0.0;
 }
 
+TextRunnerNgramRows TextRunnerPool::Request::ngram_rows() const {
+  if (impl_ == nullptr)
+    return {};
+  const auto now = impl_->runner->NgramRows(
+      dynamic_cast<const TextRunnerState&>(impl_->lease.state()));
+  const auto& base = impl_->ngram_baseline;
+  return {now.rows - base.rows, now.dedup_rows - base.dedup_rows,
+          now.cache_hits - base.cache_hits, now.disk_reads - base.disk_reads};
+}
+
 bool TextRunnerPool::Request::cache_disk_hit() const noexcept {
   return impl_ != nullptr && impl_->lease.restored_from_disk();
 }
@@ -1518,6 +1531,23 @@ TextPrefillStep TextRunnerPool::Request::Prefill(std::size_t max_input_tokens) {
   // Fallback captures and admission must finish before pinning starts. The
   // chunk's GPU work can overlap preparation; cached decode never calls it.
   impl_->lease.PreparePrefill(std::min(remaining, max_input_tokens));
+  // The following step starts at this frontier and stops at the next one, so
+  // a runner may read its inputs during this step.
+  auto lookahead_end = impl_->prompt.size();
+  for (const auto point : impl_->checkpoints)
+    if (point > frontier)
+      lookahead_end = std::min(lookahead_end, point);
+  for (const auto point : impl_->boundaries)
+    if (point > frontier)
+      lookahead_end = std::min(lookahead_end, point);
+  if (snapshot_position > frontier)
+    lookahead_end = std::min(lookahead_end, snapshot_position);
+  struct ClearLookahead {
+    TextRunnerState& state;
+    ~ClearLookahead() { state.SetPrefillLookahead({}); }
+  } clear_lookahead{state};
+  state.SetPrefillLookahead(std::span<const TextRunnerToken>(impl_->prompt)
+                                .subspan(frontier, lookahead_end - frontier));
   auto step =
       in_pass ? impl_->runner->PrefillThrough(
                     state, model_prompt, impl_->prefill_offset,

src/cli/serve/text_model_runner.hpp

@@ -180,6 +180,15 @@ struct TextPrefillStep {
   double checkpoint_ms{0};
 };
 
+/// Hashed n-gram embedding rows a request gathered. Each row was copied from
+/// a duplicate in the same gather, served by the row cache or read from disk.
+struct TextRunnerNgramRows {
+  std::uint64_t rows{0};
+  std::uint64_t dedup_rows{0};
+  std::uint64_t cache_hits{0};
+  std::uint64_t disk_reads{0};
+};
+
 struct TextDecodeSelection {
   bool stop{false};
   TextRunnerToken token{0};
@@ -210,6 +219,17 @@ class TextRunnerState : public ContinuationState {
   void SetStopAtEos(bool value) noexcept { stop_at_eos_ = value; }
   [[nodiscard]] bool stop_at_eos() const noexcept { return stop_at_eos_; }
 
+  /// Prompt tokens past the next Prefill's frontier, up to the frontier
+  /// planned after it. A runner may read their inputs ahead during that call
+  /// but never consumes them. Empty outside a prefill step.
+  void SetPrefillLookahead(std::span<const TextRunnerToken> tokens) noexcept {
+    prefill_lookahead_ = tokens;
+  }
+  [[nodiscard]] std::span<const TextRunnerToken> prefill_lookahead()
+      const noexcept {
+    return prefill_lookahead_;
+  }
+
   /// Installs a request-scoped cancellation check for model calls that can
   /// yield internally. Implementations that only yield between work units may
   /// keep the default no-op behavior.
@@ -223,6 +243,7 @@ class TextRunnerState : public ContinuationState {
 
 private:
   bool stop_at_eos_{true};
+  std::span<const TextRunnerToken> prefill_lookahead_;
 };
 
 /// Immutable model-owned continuation payload.
@@ -382,6 +403,13 @@ class TextModelRunner {
     (void)state;
     (void)prefix;
   }
+  /// Cumulative n-gram rows gathered through this state; models without an
+  /// n-gram table report none.
+  [[nodiscard]] virtual TextRunnerNgramRows NgramRows(
+      const TextRunnerState& state) const {
+    (void)state;
+    return {};
+  }
   /// Tells a request state that its next execution will use a multi-request
   /// batch. Runners with request-local speculative state may discard work that
   /// the batched path cannot consume.
@@ -524,6 +552,8 @@ class TextRunnerPool {
     [[nodiscard]] std::size_t cache_restore_bytes() const noexcept;
     [[nodiscard]] double cache_restore_ms() const noexcept;
     [[nodiscard]] bool cache_disk_hit() const noexcept;
+    /// N-gram rows gathered since this request leased its state.
+    [[nodiscard]] TextRunnerNgramRows ngram_rows() const;
     [[nodiscard]] std::optional<ComponentCacheMetrics> component_cache_metrics()
         const;
     [[nodiscard]] std::size_t prompt_tokens() const noexcept;

src/models/qwen38_flash_next/engine.cpp

@@ -273,6 +273,10 @@ std::size_t Session::AllocatedBytes() const noexcept {
   return session_->AllocatedBytes();
 }
 
+NgramTotals Session::NgramRowTotals() const noexcept {
+  return session_->NgramRowTotals();
+}
+
 bool Session::MtpEnabled() const noexcept {
   return session_->mtp_enabled();
 }
@@ -1185,7 +1189,8 @@ bool Session::FinishDecode(const DecodeRequest& request,
 
 bool Session::DecodeStep(std::size_t max_tokens,
                          sampling::SamplerState& sampler, DecodeResult* result,
-                         std::string* error_msg, bool stop_at_eos) {
+                         std::string* error_msg, bool stop_at_eos,
+                         std::uint32_t occupancy) {
   if (continuation_hooks_)
     continuation_hooks_->BeforeExecution();
   if (!valid_) {
@@ -1196,7 +1201,8 @@ bool Session::DecodeStep(std::size_t max_tokens,
   valid_ = false;
   const DecodeRequest request{this, max_tokens, &sampler, result, stop_at_eos};
   PendingDecode pending;
-  if (!PrepareDecode(request, &pending, error_msg)) {
+  if (!PrepareDecode(request, &pending, error_msg, false, std::nullopt,
+                     occupancy)) {
     return false;
   }
   if (pending.chain.empty()) {

src/models/qwen38_flash_next/engine.hpp

@@ -18,6 +18,7 @@
 #include "src/models/qwen38_flash_next/config.hpp"
 #include "src/models/qwen38_flash_next/draft_vocab.hpp"
 #include "src/models/qwen38_flash_next/mtp_policy.hpp"
+#include "src/models/qwen38_flash_next/ngram.hpp"
 #include "src/models/qwen38_flash_next/prompt_lookup.hpp"
 
 namespace gufo::core {
@@ -28,7 +29,6 @@ namespace gufo::models::qwen38_flash_next {
 
 struct ModelWeights;
 struct MtpWeights;
-class NgramTable;
 struct MtpCandidateLogits;
 class ContinuationAdapter;
 class ContinuationHooks;
@@ -163,11 +163,14 @@ class Session final {
   /// RNG consumption includes proposal and verification draws. A
   /// one-token budget or a model without MTP uses ordinary decoding.
   /// Benchmarks may continue past EOS by setting stop_at_eos to false.
+  /// A greedy step prices its drafts as one of `occupancy` decoding requests,
+  /// as DecodeBatch does for an unplanned cohort of that width.
   [[nodiscard]] bool DecodeStep(std::size_t max_tokens,
                                 sampling::SamplerState& sampler,
                                 DecodeResult* result,
                                 std::string* error_msg = nullptr,
-                                bool stop_at_eos = true);
+                                bool stop_at_eos = true,
+                                std::uint32_t occupancy = 1);
   struct BatchOutcome {
     bool completed{false};
     std::string error{};
@@ -223,6 +226,9 @@ class Session final {
   [[nodiscard]] const SpeculativeStats& Statistics() const noexcept {
     return stats_;
   }
+  /// PLE rows gathered by this execution slot since creation; zero without
+  /// a PLE table.
+  [[nodiscard]] NgramTotals NgramRowTotals() const noexcept;
 
   /// Compatibility version; bump on payload or inference arithmetic changes.
   static constexpr std::uint32_t kSnapshotPayloadVersion = 18;

src/models/qwen38_flash_next/kernels/rocm/batch.cpp

@@ -839,6 +839,7 @@ bool Executor::ForwardBatch(std::span<const BatchItem> items,
   }
   FinishPrefetch();
   batch_rows_ = 0;
+  ngram_owners_.clear();
   for (std::size_t i = 0; i < items.size(); ++i) {
     const auto& item = items[i];
     auto& session = *item.session;
@@ -850,6 +851,8 @@ bool Executor::ForwardBatch(std::span<const BatchItem> items,
     std::copy(item.tokens.begin(), item.tokens.end(),
               tokens_host_ + offsets[i]);
     if (c.ple_layer >= 0) {
+      ngram_owners_.push_back(
+          {item.tokens.size() * c.ple_heads, &session.ngram_counters_});
       for (std::size_t t = 0; t < item.tokens.size(); ++t) {
         HashNgramRows(
             c, session.ngram_, item.tokens.subspan(t, 1),
@@ -865,7 +868,8 @@ bool Executor::ForwardBatch(std::span<const BatchItem> items,
     if (ngram_ == nullptr ||
         !ngram_->StartRead(std::span(host_rows_).first(rows * c.ple_heads),
                            std::span(host_emb_, static_cast<std::size_t>(rows) *
-                                                    c.PleEmbeddingDim()))) {
+                                                    c.PleEmbeddingDim()),
+                           ngram_owners_)) {
       return Fail(error, "batched n-gram read could not start");
     }
     ple_pending_ = true;