feat(ext): keep indivisible prompt spans whole across prefill stops and checkpoints (#1011)

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

docs/KV-CACHE.md

@@ -310,6 +310,14 @@ The cache reserves the full payload before enabling this capture. Disk,
 image, intermediate and learned-prefix checkpoints still stop prefill at
 their boundary.
 
+**A checkpoint never splits an indivisible input.** A model whose prompt
+context declares spans through `TextPromptContext::PrefillStop`, such as image
+rows that must share one forward, never has a prefill step, planned checkpoint,
+stable snapshot or shared-prefix boundary end inside one, and never restores a
+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.
+
 **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/indivisible-prompt-spans.md

@@ -0,0 +1,17 @@
+# Indivisible prompt spans
+
+Add `TextPromptContext::PrefillStop(position, up)` (`src/ext/text.hpp`), a
+virtual hook through which a model-owned prompt context keeps inputs whose
+tokens must be evaluated in one forward, such as image rows that attend each
+other, out of split prefill steps. It moves a candidate stop to the first row
+of the span containing it, or past the span's last token when `up`. The text
+runner pool applies it to every position it chooses: step ends, planned cache
+captures, the stable snapshot boundary and shared-prefix checkpoints. The cache
+never restores a checkpoint inside a span. A step starting at a span consumes it
+whole, beyond its token budget, so a runner's scratch must fit its largest span.
+A runner stopping inside a span, or a hook moving a stop the wrong way, fails
+the request. With `prefill_alignment`, every stop is both aligned and outside
+the spans.
+
+The default returns `position`, so existing contexts and models behave
+identically. No other declaration changes.

src/cli/serve/component_text_cache.cpp

@@ -1140,7 +1140,13 @@ std::unique_ptr<ComponentTextCache::Lease> ComponentTextCache::Acquire(
       return boundary < prompt.size() ||
              impl_->runner->HasDecodeFrontier(*request.state);
     };
-    if (live && !decodable(found.selected->boundary)) {
+    // Prefill must not resume inside an indivisible input of this prompt,
+    // which a checkpoint from another prompt with the same tokens can do.
+    const auto whole_inputs = [&](std::size_t boundary) {
+      return !context || context->PrefillStop(boundary, false) == boundary;
+    };
+    if (live && (!decodable(found.selected->boundary) ||
+                 !whole_inputs(found.selected->boundary))) {
       live.reset();
       const std::lock_guard lock(impl_->mutex);
       entry->branch = std::make_shared<const std::uint8_t>(0);
@@ -1158,8 +1164,9 @@ std::unique_ptr<ComponentTextCache::Lease> ComponentTextCache::Acquire(
     for (const auto& candidate : found.candidates) {
       if (restored || (cancelled && cancelled()))
         break;
-      if (candidate.live || (candidate.boundary % alignment != 0 &&
-                             candidate.boundary < prompt.size()))
+      if (candidate.live || !whole_inputs(candidate.boundary) ||
+          (candidate.boundary % alignment != 0 &&
+           candidate.boundary < prompt.size()))
         continue;
       try {
         if (candidate.checkpoint) {

src/cli/serve/text_model_runner.cpp

@@ -281,6 +281,12 @@ struct TextRunnerPool::Request::Impl {
     // Freeze it before prefill, then retain this turn's own stable boundary
     // before mutable assistant framing.
     auto count = cache_prefix_tokens == 0 ? prompt.size() : cache_prefix_tokens;
+    // Neither snapshots nor any other planned stop may split an indivisible
+    // input. Without declared spans every position is unchanged.
+    // A stable boundary inside the first input leaves nothing to retain.
+    count = LegalStop(count, false);
+    if (count == 0)
+      count = prompt.size();
     // A newly consumed image needs its own stable checkpoint. Otherwise the
     // next assistant turn can invalidate the complete prompt's framing and
     // force the image through prefill again from the older text frontier.
@@ -351,7 +357,7 @@ struct TextRunnerPool::Request::Impl {
     }
     ngram_baseline = runner->NgramRows(state);
     checkpoints = lease->PlannedCaptures(cache_prefix_tokens);
-    if (alignment > 1) {
+    if (alignment > 1 || context) {
       for (auto& position : checkpoints)
         position = Aligned(position);
       std::erase_if(checkpoints, [&](std::size_t position) {
@@ -363,8 +369,43 @@ struct TextRunnerPool::Request::Impl {
     }
   }
 
-  [[nodiscard]] std::size_t Aligned(std::size_t position) const noexcept {
-    return position - position % alignment;
+  /// The prompt context's stop for `position`. A hook moving the wrong way
+  /// would keep the loops below from ending.
+  [[nodiscard]] std::size_t LegalStop(std::size_t position, bool up) const {
+    if (!context)
+      return position;
+    const auto stop = context->PrefillStop(position, up);
+    if (up ? stop < position : stop > position)
+      throw std::logic_error(
+          "prompt context moved a prefill stop the wrong way");
+    return stop;
+  }
+
+  /// The nearest position at or before `position` where a pass may end: on
+  /// the prefill alignment and outside every indivisible input.
+  [[nodiscard]] std::size_t Aligned(std::size_t position) const {
+    for (;;) {
+      const auto next = LegalStop(position - position % alignment, false);
+      if (next == position)
+        return position;
+      position = next;
+    }
+  }
+
+  /// The nearest such position after `position`, or the prompt end.
+  [[nodiscard]] std::size_t NextStop(std::size_t position) const {
+    for (;;) {
+      if (position >= prompt.size())
+        return prompt.size();
+      auto next = position;
+      if (const auto over = next % alignment; over != 0)
+        next += alignment - over;
+      if (next < prompt.size())
+        next = LegalStop(next, true);
+      if (next == position)
+        return position;
+      position = next;
+    }
   }
 
   ~Impl() {
@@ -901,6 +942,16 @@ TextPrefillStep TextRunnerPool::Request::Prefill(std::size_t max_input_tokens) {
             : std::min(max_input_tokens + impl_->alignment - over,
                        impl_->prompt.size() - impl_->prefill_offset);
   }
+  // Nor does it end inside an indivisible input: it ends before one, or a
+  // step starting at one consumes it whole, beyond the budget. Every planned
+  // stop is legal already, so later clamps keep the step legal.
+  if (const auto end = impl_->prefill_offset + max_input_tokens;
+      impl_->context && end < impl_->prompt.size()) {
+    auto stop = impl_->Aligned(end);
+    if (stop <= impl_->prefill_offset)
+      stop = impl_->NextStop(impl_->prefill_offset + 1);
+    max_input_tokens = stop - impl_->prefill_offset;
+  }
 
   auto& state = dynamic_cast<TextRunnerState&>(impl_->lease->state());
   // Captures this chunk takes in place, ascending, without ending a pass
@@ -1100,6 +1151,10 @@ TextPrefillStep TextRunnerPool::Request::Prefill(std::size_t max_input_tokens) {
   }
 
   impl_->prefill_offset += step.consumed_tokens;
+  if (impl_->LegalStop(impl_->prefill_offset, false) != impl_->prefill_offset) {
+    throw std::runtime_error(
+        "text runner stopped prefill inside an indivisible input");
+  }
   const bool reached_frontier = impl_->prefill_offset == impl_->prompt.size();
   if (step.decode_ready != (impl_->prefill_offset == model_prompt.size())) {
     throw std::runtime_error(

src/ext/text.hpp

@@ -45,6 +45,23 @@ struct TextPromptContext {
       std::size_t token_count) const {
     return PrefixInputIdentity(cache_identity, cache_prefixes, token_count);
   }
+  /// Moves a candidate prefill stop, a token count, to a position where
+  /// prefill may end: the nearest legal one at or before `position`, or at or
+  /// after it when `up`. A model whose input rows must be evaluated in one
+  /// forward, such as image rows attending each other, returns the first row
+  /// of the span containing `position` (down) or the position after its last
+  /// token (up). Positions at span edges and outside spans are legal and must
+  /// be returned unchanged, and the result must be a fixed point of this
+  /// function. The pool applies it to cache checkpoints, snapshot and shared
+  /// prefix boundaries, and never restores a checkpoint inside a span. A step
+  /// that starts at a span consumes the whole span even beyond its token
+  /// budget, so the runner's scratch must fit its largest span, and a runner
+  /// consuming less than a step must not stop inside one. The default declares
+  /// no spans.
+  [[nodiscard]] virtual std::size_t PrefillStop(std::size_t position,
+                                                bool /*up*/) const {
+    return position;
+  }
 };
 
 struct TextPreparedPrompt {

tests/cli/component_text_cache_test.cpp

@@ -2653,7 +2653,188 @@ void RamCacheOffRefusesCapturesWithinTheDiskStep() {
     }
   }
 }
+// Inputs whose tokens must be evaluated together, such as images. A span is
+// [first, end): positions strictly inside it are not legal stops.
+struct SpanContext final : TextPromptContext {
+  std::vector<std::pair<std::size_t, std::size_t>> spans;
+  std::size_t PrefillStop(std::size_t position, bool up) const override {
+    for (const auto& [first, end] : spans)
+      if (position > first && position < end)
+        return up ? end : first;
+    return position;
+  }
+};
+std::shared_ptr<SpanContext> Spans(
+    std::vector<std::pair<std::size_t, std::size_t>> spans) {
+  auto context = std::make_shared<SpanContext>();
+  context->spans = std::move(spans);
+  return context;
+}
+std::vector<TextRunnerToken> SpanPrompt(std::size_t tokens) {
+  std::vector<TextRunnerToken> prompt(tokens);
+  std::iota(prompt.begin(), prompt.end(), 10);
+  return prompt;
+}
+void IndivisibleInputsAreNotSplit() {
+  using Stops = std::vector<std::size_t>;
+  using Context = std::shared_ptr<const TextPromptContext>;
+  // Prefill in steps of `budget` and report where the steps ended.
+  const auto steps = [](std::size_t tokens, std::size_t budget, Context context,
+                        std::size_t alignment = 1, std::size_t stable = 0) {
+    auto runner = std::make_shared<Runner>();
+    runner->alignment = alignment;
+    TextRunnerPool pool(runner, 1);
+    auto request = pool.Acquire(SpanPrompt(tokens), sampling::SamplingConfig{},
+                                {}, std::move(context), true, stable);
+    while (!request.prefill_complete())
+      (void)request.Prefill(budget);
+    request.Invalidate();
+    return runner->pass_ends;
+  };
+  // A budget stops before an input, and a step starting at one takes it whole.
+  assert(steps(20, 3, Spans({{4, 14}})) == Stops({3, 4, 14, 17, 20}));
+  assert(steps(20, 3, Spans({{4, 10}, {10, 16}})) ==
+         Stops({3, 4, 10, 16, 19, 20}));
+  assert(steps(12, 3, Spans({{0, 6}})) == Stops({6, 9, 12}));
+  assert(steps(12, 3, Spans({{8, 12}})) == Stops({3, 6, 8, 12}));
+  // A budget ending exactly at an input's edges is not moved.
+  assert(steps(20, 3, Spans({{5, 8}})) == Stops({3, 5, 8, 11, 14, 17, 20}));
+  // No declared spans change no step.
+  assert(steps(20, 3, {}) == Stops({3, 6, 9, 12, 15, 18, 20}));
+  assert(steps(20, 3, Spans({})) == steps(20, 3, {}));
+  // With an alignment a stop is both aligned and outside every input; only
+  // the prompt end may be off the alignment.
+  assert(steps(42, 6, Spans({{10, 19}}), 4) ==
+         Stops({4, 8, 20, 24, 28, 32, 36, 40, 42}));
+  assert(steps(42, 64, Spans({{22, 29}}), 4, 30) == Stops({20, 42}));
+  assert(steps(42, 64, Spans({}), 4, 30) == Stops({28, 42}));
+  // A planned capture inside an input moves before it.
+  const auto planned = steps(3000, 4096, {});
+  assert(planned.size() > 1);
+  const auto capture = planned.front();
+  const auto moved = steps(3000, 4096, Spans({{capture - 8, capture + 8}}));
+  assert(moved.front() == capture - 8);
+  assert(steps(3000, 4096, Spans({{capture, capture + 8}})).front() == capture);
+  // The stable boundary lies inside the input: its checkpoint moves before it.
+  for (const bool in_pass : {false, true}) {
+    auto runner = std::make_shared<Runner>();
+    runner->in_pass = in_pass;
+    runner->max_in_pass_tail = 16;
+    TextRunnerPool pool(runner, 1);
+    const auto prompt = SpanPrompt(20);
+    const auto spans = Spans({{4, 14}});
+    auto first =
+        pool.Acquire(prompt, sampling::SamplingConfig{}, {}, spans, true, 10);
+    while (!first.prefill_complete())
+      (void)first.Prefill(64);
+    assert(runner->pass_ends == (in_pass ? Stops({20}) : Stops({4, 20})));
+    assert(runner->in_pass_points == (in_pass ? Stops({4}) : Stops{}));
+    (void)first.Commit();
+    auto edited = prompt;
+    edited[12] += 1000;
+    auto second =
+        pool.Acquire(edited, sampling::SamplingConfig{}, {}, spans, true, 10);
+    assert(second.cached_prompt_tokens() == 4);
+    second.Invalidate();
+  }
+  // A peer's shared prefix ending inside an input stops before the input.
+  {
+    auto runner = std::make_shared<Runner>();
+    TextRunnerPool pool(runner, 1);
+    auto leader = pool.Acquire(SpanPrompt(1600), sampling::SamplingConfig{}, {},
+                               Spans({{900, 1100}}));
+    assert(leader.ShareCheckpoint(1000) == 900);
+    while (!leader.prefill_complete())
+      (void)leader.Prefill(4096);
+    assert(runner->pass_ends.front() == 900);
+    leader.Invalidate();
+  }
+  // A checkpoint or live frontier from a prompt without the input cannot
+  // resume inside it, even with the same tokens.
+  {
+    auto runner = std::make_shared<Runner>();
+    TextRunnerPool pool(runner, 1);
+    const auto prompt = SpanPrompt(20);
+    auto first =
+        pool.Acquire(prompt, sampling::SamplingConfig{}, {}, {}, true, 10);
+    Prefill(first);
+    (void)first.Commit();
+    auto edited = prompt;
+    edited[15] += 1000;
+    auto plain =
+        pool.Acquire(edited, sampling::SamplingConfig{}, {}, {}, true, 0);
+    assert(plain.cached_prompt_tokens() == 10);
+    plain.Invalidate();
+    auto spanned = pool.Acquire(edited, sampling::SamplingConfig{}, {},
+                                Spans({{4, 14}}), true, 0);
+    assert(spanned.cached_prompt_tokens() == 0);
+    spanned.Invalidate();
+  }
+  for (const bool spanned : {false, true}) {
+    auto runner = std::make_shared<Runner>();
+    TextRunnerPool pool(runner, 1);
+    auto history = SpanPrompt(10);
+    auto first = pool.Acquire(history);
+    Prefill(first);
+    history.push_back(Generate(first));
+    (void)first.Commit();
+    history.resize(30, 7);
+    Context context;
+    if (spanned)
+      context = Spans({{4, 14}});
+    auto next =
+        pool.Acquire(history, sampling::SamplingConfig{}, {}, context, true, 0);
+    assert(next.cached_prompt_tokens() == (spanned ? 0 : 11));
+    next.Invalidate();
+  }
+  // A runner stopping inside an input fails the request.
+  {
+    class ShortStepRunner final : public Runner {
+    public:
+      TextPrefillStep Prefill(TextRunnerState& state,
+                              std::span<const TextRunnerToken> prompt,
+                              std::size_t offset,
+                              std::size_t budget) const override {
+        return Runner::Prefill(state, prompt, offset,
+                               std::min<std::size_t>(budget, 6));
+      }
+    };
+    TextRunnerPool pool(std::make_shared<ShortStepRunner>(), 1);
+    auto request = pool.Acquire(SpanPrompt(20), sampling::SamplingConfig{}, {},
+                                Spans({{4, 14}}));
+    bool failed = false;
+    try {
+      (void)request.Prefill(64);
+    } catch (const std::runtime_error&) {
+      failed = true;
+    }
+    assert(failed);
+    request.Invalidate();
+  }
+  // A hook moving a stop the wrong way fails instead of hanging.
+  {
+    struct WrongWayContext final : TextPromptContext {
+      std::size_t PrefillStop(std::size_t position, bool) const override {
+        return position > 4 && position < 14 ? 4 : position;
+      }
+    };
+    auto runner = std::make_shared<Runner>();
+    TextRunnerPool pool(runner, 1);
+    auto request = pool.Acquire(SpanPrompt(20), sampling::SamplingConfig{}, {},
+                                std::make_shared<WrongWayContext>());
+    bool failed = false;
+    try {
+      while (!request.prefill_complete())
+        (void)request.Prefill(3);
+    } catch (const std::logic_error&) {
+      failed = true;
+    }
+    assert(failed);
+    request.Invalidate();
+  }
+}
 int main() {
+  IndivisibleInputsAreNotSplit();
   PromptReuseCanBeDisabledPerRequest();
   StableChatPrefixSurvivesInterruptedFraming();
   WarmTurnsAdvanceTheStableCheckpoint();