feat(ext): let text runners see the whole prompt end during prefill (#1069)

Co-authored-by: Claude Opus 5.5 <[email protected]>
Co-authored-by: Federico Izzo <[email protected]>

docs/KV-CACHE.md

@@ -318,6 +318,14 @@ checkpoint inside one. Such a position moves to the span's first row, and a step
 that starts at a span takes it whole even beyond its token budget. Contexts that
 declare no spans change no position.
 
+**A checkpoint may hold less than a cold request would.** Each prefill step's
+prompt ends at its own frontier, and `TextRunnerState::prefill_prompt_end()`
+gives the whole prompt's end. A model that runs some work only near that end
+can skip it in steps ending at earlier cache boundaries, so their checkpoints
+lack those rows. Its `TextModelRunner::CanContinue` then refuses prompts that
+would need them, and the cache serves those from a shallower boundary. By
+default every state continues every prompt.
+
 **The reused frontier must be frozen before prefill.** Prefill mutates the
 leased state in place, so a frontier another request might branch from has to
 be captured first.

include/gufo/ext/changes/0.2/prefill-prompt-end.md

@@ -0,0 +1,15 @@
+# Whole prompt end during prefill
+
+Add `TextRunnerState::prefill_prompt_end()` and `SetPrefillPromptEnd`
+(`src/ext/text.hpp`). During a prefill step, `prefill_prompt_end()` holds the
+whole prompt's token count; a step's own `prompt` span ends at its frontier,
+which may be a cache boundary short of it. It is 0 outside a step and when
+a caller drives the runner without the cache, meaning the step's own end.
+
+Add the virtual `TextModelRunner::CanContinue(state, prompt_tokens)`. A runner
+that keeps only the state longer prompts need returns false for a prompt its
+live or restored state cannot continue, and the cache serves that prompt from a
+shallower boundary.
+
+Additive: the default `CanContinue` returns true, so existing runners behave
+identically.

src/cli/serve/component_text_cache.cpp

@@ -1172,10 +1172,12 @@ std::unique_ptr<ComponentTextCache::Lease> ComponentTextCache::Acquire(
     // an unusable live frontier) can require another reset before cold use.
     bool clean = !live;
     // A frontier captured between prefill passes cannot select the next token
-    // of a prompt that ends there; serve that prompt from a shallower boundary.
+    // of a prompt that ends there, and a runner may refuse other prompts; serve
+    // those from a shallower boundary.
     const auto decodable = [&](std::size_t boundary) {
-      return boundary < prompt.size() ||
-             impl_->runner->HasDecodeFrontier(*request.state);
+      return (boundary < prompt.size() ||
+              impl_->runner->HasDecodeFrontier(*request.state)) &&
+             impl_->runner->CanContinue(*request.state, prompt.size());
     };
     // Prefill must not resume inside an indivisible input of this prompt,
     // which a checkpoint from another prompt with the same tokens can do.

src/cli/serve/text_model_runner.cpp

@@ -1197,10 +1197,14 @@ TextPrefillStep TextRunnerPool::Request::Prefill(std::size_t max_input_tokens) {
   }
   struct ClearLookahead {
     TextRunnerState& state;
-    ~ClearLookahead() { state.SetPrefillLookahead({}); }
+    ~ClearLookahead() {
+      state.SetPrefillLookahead({});
+      state.SetPrefillPromptEnd(0);
+    }
   } clear_lookahead{state};
   state.SetPrefillLookahead(std::span<const TextRunnerToken>(impl_->prompt)
                                 .subspan(frontier, lookahead_end - frontier));
+  state.SetPrefillPromptEnd(impl_->prompt.size());
   std::vector<std::size_t> in_pass_positions;
   for (const auto& capture : in_pass)
     in_pass_positions.push_back(capture.position);

src/ext/text.hpp

@@ -247,6 +247,16 @@ public:
       const noexcept {
     return prefill_lookahead_;
   }
+  /// The whole prompt's token count during a prefill step, 0 outside one. A
+  /// step's prompt ends at its own frontier, which may be a cache boundary
+  /// short of this. A caller driving the runner outside the serve cache
+  /// leaves it 0; treat 0 as the step's own prompt end.
+  void SetPrefillPromptEnd(std::size_t tokens) noexcept {
+    prefill_prompt_end_ = tokens;
+  }
+  [[nodiscard]] std::size_t prefill_prompt_end() const noexcept {
+    return prefill_prompt_end_;
+  }
 
   virtual void SetCancellationCheck(const CancellationCheck&) {}
 
@@ -258,6 +268,7 @@ public:
 private:
   bool stop_at_eos_{true};
   std::span<const TextRunnerToken> prefill_lookahead_;
+  std::size_t prefill_prompt_end_{};
 };
 
 struct TextRunnerAdvance {
@@ -510,6 +521,14 @@ public:
   [[nodiscard]] virtual bool HasDecodeFrontier(const TextRunnerState&) const {
     return true;
   }
+  /// Whether the state, live or restored, continues a prompt of
+  /// `prompt_tokens`. A model that skips work on rows far from the prompt's
+  /// end may keep only what longer prompts need; the cache serves a refused
+  /// prompt from a shallower boundary.
+  [[nodiscard]] virtual bool CanContinue(const TextRunnerState&,
+                                         std::size_t /*prompt_tokens*/) const {
+    return true;
+  }
 
   [[nodiscard]] virtual std::optional<TextDecodeSelection> PreviewFirstToken(
       TextRunnerState&, sampling::SamplerState&) const {

tests/cli/component_text_cache_test.cpp

@@ -839,6 +839,59 @@ void PrefillLookaheadReachesTheNextFrontier() {
     assert(end == (frontier == prompt.size() ? frontier : next));
   }
 }
+void PrefillStepsSeeTheWholePromptEnd() {
+  class PromptEndRunner final : public Runner {
+  public:
+    mutable std::vector<std::size_t> ends, frontiers;
+    mutable const TextRunnerState* state{};
+    void Record(const TextRunnerState& state,
+                std::span<const TextRunnerToken> prompt) const {
+      ends.push_back(state.prefill_prompt_end());
+      frontiers.push_back(prompt.size());
+      this->state = &state;
+    }
+    TextPrefillStep Prefill(TextRunnerState& state,
+                            std::span<const TextRunnerToken> prompt,
+                            std::size_t offset,
+                            std::size_t budget) const override {
+      Record(state, prompt);
+      return Runner::Prefill(state, prompt, offset, budget);
+    }
+    TextPrefillStep PrefillThrough(
+        TextRunnerState& state, std::span<const TextRunnerToken> prompt,
+        std::size_t offset, std::size_t budget,
+        std::span<const std::size_t> boundaries,
+        const InPassConsumer& consume) const override {
+      Record(state, prompt);
+      return Runner::PrefillThrough(state, prompt, offset, budget, boundaries,
+                                    consume);
+    }
+  };
+  auto runner = std::make_shared<PromptEndRunner>();
+  TextRunnerPool pool(runner, 1);
+  std::vector<TextRunnerToken> prompt(4000);
+  std::iota(prompt.begin(), prompt.end(), 10);
+  auto request = pool.Acquire(prompt);
+  while (!request.Prefill(1024).decode_ready) {
+  }
+  assert(runner->state && runner->state->prefill_prompt_end() == 0);
+  // A step ending at a cache boundary still sees the whole prompt's end.
+  assert(runner->frontiers.front() < prompt.size());
+  for (const auto end : runner->ends)
+    assert(end == prompt.size());
+
+  auto through_runner = std::make_shared<PromptEndRunner>();
+  through_runner->in_pass = true;
+  TextRunnerPool through_pool(through_runner, 1);
+  std::vector<TextRunnerToken> short_prompt(40);
+  std::iota(short_prompt.begin(), short_prompt.end(), 10);
+  auto through = through_pool.Acquire(short_prompt, sampling::SamplingConfig{},
+                                      {}, {}, true, 35);
+  assert(through.Prefill(64).decode_ready);
+  assert(through_runner->through_calls == 1);
+  assert(through_runner->ends == std::vector<std::size_t>{short_prompt.size()});
+  assert(through_runner->state->prefill_prompt_end() == 0);
+}
 void BackingPreparationFollowsCaptures() {
   class PrepareRunner final : public Runner {
   public:
@@ -942,6 +995,60 @@ void PartialFrontierDoesNotServeAnEqualPrompt() {
     resend.Invalidate();
   }
 }
+void RefusedFrontierServesFromShallowerBoundary() {
+  // A runner may keep only the state that longer prompts need. A prompt it
+  // refuses recomputes from a shallower boundary, live or restored; a longer
+  // one continues the frontier.
+  class RefusingRunner final : public Runner {
+  public:
+    bool CanContinue(const TextRunnerState& state,
+                     std::size_t prompt_tokens) const override {
+      return prompt_tokens >= CheckpointPosition(state) + 8;
+    }
+  };
+  // A live frontier left by a cancelled prefill.
+  for (const std::size_t next_size : {20U, 40U}) {
+    auto runner = std::make_shared<RefusingRunner>();
+    runner->partial = true;
+    TextRunnerPool pool(runner, 1);
+    std::vector<TextRunnerToken> prompt(40);
+    std::iota(prompt.begin(), prompt.end(), 10);
+    auto request = pool.Acquire(prompt, {}, [] { return false; }, {}, true, 0);
+    assert(request.Prefill(16).consumed_tokens == 16);
+    (void)request.Cancel();
+    const std::vector<TextRunnerToken> next_prompt(prompt.begin(),
+                                                   prompt.begin() + next_size);
+    auto next = pool.Acquire(next_prompt, {}, {}, {}, true, 0);
+    assert((next.cached_prompt_tokens() == 16) == (next_size == 40));
+    assert(next.cached_prompt_tokens() <= 16);
+    Prefill(next);
+    (void)Generate(next);
+    next.Invalidate();
+  }
+  // A checkpoint restored after another prompt replaced the slot.
+  for (const std::size_t next_size : {37U, 60U}) {
+    auto runner = std::make_shared<RefusingRunner>();
+    TextRunnerPool pool(runner, 1);
+    std::vector<TextRunnerToken> prompt(60);
+    std::iota(prompt.begin(), prompt.end(), 10);
+    auto seed = pool.Acquire(prompt, {}, {}, {}, true, 35);
+    Prefill(seed);
+    (void)seed.Commit();
+    auto unrelated = prompt;
+    unrelated.front() = 999;
+    auto other = pool.Acquire(unrelated, {}, {}, {}, true, 0);
+    Prefill(other);
+    (void)other.Commit();
+    auto next_prompt = prompt;
+    next_prompt.resize(next_size);
+    next_prompt.back() = 998;
+    auto next = pool.Acquire(next_prompt, {}, {}, {}, true, 0);
+    assert((next.cached_prompt_tokens() == 35) == (next_size == 60));
+    Prefill(next);
+    (void)Generate(next);
+    next.Invalidate();
+  }
+}
 void CompletedCancelledPromptGetsARetrySnapshot() {
   auto runner = std::make_shared<Runner>();
   runner->partial = true;
@@ -2980,7 +3087,9 @@ int main() {
   SharedPeerReusesTheCurrentStableCapture();
   BackingPreparationFollowsCaptures();
   PrefillLookaheadReachesTheNextFrontier();
+  PrefillStepsSeeTheWholePromptEnd();
   PartialFrontierDoesNotServeAnEqualPrompt();
+  RefusedFrontierServesFromShallowerBoundary();
   CompletedCancelledPromptGetsARetrySnapshot();
   BudgetUsesAllocatedSlotsInsteadOfLazyStateClaim();
   InPassCapturePreservesBoundaryAndLiveTail();