Skip to content

prefix: state-slot entries are exact-length-only, so any divergence costs the whole prefill #11

Description

@xiaguan

Summary

For a manifest with a recurrent state slot (tray.has_seq_state()), prefix-cache entries are exact-length-only. Any divergence between two consecutive requests in a conversation — even a single token, even one that appears near the very end of a long prompt — yields a hit of zero, not "everything up to the divergence".

For long multi-turn prompts this is the difference between re-prefilling a handful of tokens and re-prefilling the entire context.

docs/serve.md already documents the mild form of this ("同一 prompt 重发不命中"). The sharp form is that it also discards every usable page of a prefix that does match, which I think is not intended.

Where

Two places combine.

1. Entries are only registered at one length. crates/kern-serve/src/scheduler.rs:

let every_page = !tray.has_seq_state();

With a state slot, checkpoint() returns immediately and the only prefix.insert is in finish() — one entry, at the sequence's full retired length.

2. Lookup refuses to fall back for those entries. crates/kern-runtime/src/prefix.rs, in lookup:

match (tail_ok, e.held.has_slot()) {
    (true, _)      => full + k.tail_len,
    (false, false) => full,      // pages-only: fall back to whole pages
    (false, true)  => continue,  // state entry: exact length or nothing
}

Entry::buckets() reinforces it — a slot entry has heads empty, so it is registered in exactly one (depth, chain) bucket and is invisible at any shallower depth.

The pages-only path is fine and does what you'd want (divergence inside page k still returns k * unit tokens; covered by crates/kern-runtime/tests/prefix.rs).

Why this bites in practice

The rendered prompt of turn N+1 is generally not a byte-identical extension of turn N's, because chat templates rewrite history. A concrete, fully reproducible example from the pinned renderer (this part is in the vllm-chat dep, not in this repo — rust/src/chat/src/renderer/deepseek.rs):

let emit_thinking_block = thinking_mode == ThinkingMode::Thinking
    && (!drop_thinking || current_render_index > last_user_render_index);

last_user_render_index is computed once over the whole message list. Appending a new user turn moves it, so the most recent assistant turn flips from "thinking kept" to "thinking dropped". The rendered prefix therefore diverges at the start of that assistant turn.

Note the shape of this: the divergence point is the last assistant turn. Everything before it is stable. Under page-granular matching almost the entire prompt is still reusable; under exact-length matching the hit is 0.

Minimal repro, two turns, client echoing the assistant's reasoning_content back verbatim:

turn 2 settings rendered prompt prefix hit
defaults 73 tokens none
chat_template_kwargs: {"thinking": true, "drop_thinking": false} 141 tokens 126 tokens

The workaround changes the render so it happens to stay stable; it does not address the matching behaviour, and it only works if the client echoes reasoning back.

Suggestion

Decouple the state slot from the KV pages, so the two can be reused independently:

  • register slot-bearing entries at every whole page as well, so lookup can return full for them like any other entry;
  • on a hit shorter than the entry, reuse the KV pages and treat the recurrent state as absent rather than disqualifying — rebuild it, or cold-start it.

For a speculative-decoding draft the state is the draft's own cache, and re-running the draft over a short window is cheap relative to re-prefilling the whole context. For a GDN-style state the trade is a real one and may need a flag, but "exact length or nothing" seems like the wrong default either way.

A narrower alternative, if the above is too invasive: keep the disqualification but let lookup fall through to a shallower pages-only entry on the same chain instead of continue-ing out of the whole depth. That recovers partial hits whenever a pages-only entry happens to exist, without changing the state semantics.

Related, smaller: cached_tokens is never reported

KernScheduler knows the hit length (start = row.prefix(), logged as prefix_hit at debug level and summed into stats.prefix_hit_tokens) but never calls ledger.set_cached_tokens(...). The frontend's usage.prompt_tokens_details.cached_tokens is therefore always absent, so over HTTP a large hit and a total miss are indistinguishable. This makes the issue above hard to diagnose from the client side — you have to turn on debug logging to see it at all.

Looks like a one-line addition next to the existing debug!(... prefix_hit = start ...).

Environment

  • kern at e6c1487
  • manifest with a recurrent state slot, so every_page = false

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions