[Feature]: KV-cache reuse for shared-prefix one-shot requests (prefix caching / partial truncate)

Suggestion Description

FLM's KV reuse only serves conversations that grow. A very common serving pattern gets zero reuse today: many independent one-shot requests sharing a long stable prefix (instructions, reference documents, few-shot rubrics, RAG context) and differing only in a small trailing payload. Every such request re-prefills the entire shared prefix. On NPU hardware this dominates latency and power: in our workload (~1,900-token prompts sharing ~1,880 tokens, ~35-token payloads), every request pays ~7 s of NPU prefill for ~35 tokens of new content.

The blocker is architectural, and it sits in the engine's cache API. We traced the full request path and are asking for two small primitives that would unlock this class of workloads — and make FLM's caching equivalent to what llama-server has shipped for years.

The workload

Typical shapes: batch raters/classifiers, evaluation harnesses scoring many samples against one long rubric, RAG pipelines with a fixed retrieved corpus, templated extraction over fixed documents.

Client-side workarounds exist — e.g. the client maintains a growing fake conversation so the round-based cache matches — but they:

• multiply upload volume O(N²),
• make the model see prior requests' payloads (contaminating outputs),
• push complexity into every client.

Why the current architecture can't serve it

We traced the full request path; these are the specific blockers:

1. PromptCache::can_use_cache() is conversation-shaped, not prefix-shaped. It requires messages.size() > 2 and that all previously cached messages reappear as a strict prefix of the new message list. One-shot requests fail the first gate immediately, no matter how much text they share with previous requests.

2. _shared_insert() is append-only; divergence clears everything. The token-level prefix walk against token_history skips matched leading tokens, but any mid-stream divergence calls clear_context() — the whole KV cache is dropped and re-prefilled. Two requests sharing their first 90% and diverging in the last 10% reuse nothing.

3. The engine exposes no partial-truncate primitive. The causal_lm interface offers prefill, forward, clear_context(), and a single-slot checkpoint()/restore(). There is no truncate(len)/seq_rm(pos) equivalent, and get_k_cache()/get_v_cache() are read-only — so even a wrapper-level implementation is impossible against the prebuilt engine libraries (libqwen3_npu.so etc.).

4. Checkpoint placement is the deeper bottleneck. A checkpoint can only be captured at the current end of cache. Prefix reuse would be nearly free if a checkpoint — or a raw rewind — could be placed inside an already-prefilled prompt (i.e. at the prefix/payload boundary). Today the entire prefix must be re-prefilled just to place a cache state at that boundary.

How llama.cpp solves the same problem (existence proof)

llama.cpp's KV cache is cell-addressable with per-cell sequence tags and a logical n_past cursor. Its core primitives — llama_kv_cache_seq_rm(seq, pos0, pos1) (remove a token range), plus seq_keep/seq_cp/seq_add (keep/copy/branch sequences) — mean "rewind to position N" is simply dropping cells ≥ N; attention only reads below the cursor, so it's O(range) bookkeeping, not recompute.

On top of that, llama-server keeps a multi-entry prompt cache: it stores the tokenized prompts of the last N conversations, and per request computes the longest common token prefix (LCP) against a cached entry, truncates the cached KV to the LCP via seq_rm, and prefills only the remainder. Consequences worth highlighting:

• Mid-prompt divergence is fine: two requests sharing 90% and diverging at token 1,700 still reuse 1,700 tokens. There is no "clear on mismatch" cliff.
• Many prefixes coexist with LRU-style eviction, so interleaved clients with different prefixes don't destroy each other.
• Shape-agnostic: one-shot requests, multi-turn chats, tool loops — anything — reuse via the same raw token-LCP check, no message-structure gates.
• Branching is a copy: seq_cp lets forked conversations (regeneration, speculative decoding) share parent KV without re-prefill.
• Correctness model: the LCP token match at request time is the only source of truth, with full-prefill fallback on any mismatch — reuse is never load-bearing.

FLM's wrapper layer is already well-aligned with this model: _shared_insert()'s token-prefix walk and the full-prefill fallback are the same safety-net philosophy. What's missing is entirely in the engine API.

Requested changes (in order of impact)

1. A token-range truncate/rewind primitive on the engine — e.g. truncate(len) (a logical cursor move + cell invalidation is enough; attention already reads only below the current length). This single primitive enables: partial-prefix reuse on mid-prompt divergence, checkpoint placement inside prefills, and forked-conversation support at the wrapper level.
2. Multiple checkpoint slots (save/restore by handle, or several named pins) — enables several distinct stable prefixes and pinned system prompts to coexist across interleaved clients.
3. Longer term: an LCP-based cache in _shared_insert — on divergence, keep KV for the longest common token prefix instead of clearing, with a small LRU over cached prompts (the llama-server design). This is the general solution and would make caching shape-agnostic for all clients.
4. Smaller hardening items observed along the way:
   • Centralize the per-wrapper checkpoint bookkeeping (restore_allowed → restore(); token_history = checkpoint_his + checkpoint_his = token_history; engine->checkpoint() is duplicated in every model wrapper, with a documented silent-desync trap). One shared hook pair in AutoModel would make engine cache features far cheaper to add across models.
   • Make the restore_allowed path verify consistency (e.g. compare restore()'s returned length against checkpoint_his.size()) so out-of-band checkpointing degrades to a cache miss instead of silent corruption.

Expected impact

For prefix-heavy serving, per-request prefill drops from the full shared prefix to just the new payload — in our measurements from ~1,900 tokens to ~35 (wall time ~7.7 s → ~4.0 s on qwen3:8b, generation-bound; prefill-bound workloads gain proportionally more), with outputs identical to cold inference.

Batch/RAG/eval workloads would see the largest wins: effectively full-prefill cost only for the first request, near-free reuse after.

Operating System

No response

GPU

No response

ROCm Component

No response

Start by tracing PromptCache::can_use_cache(), _shared_insert(), and the causal_lm checkpoint/restore interface to confirm the current cache and divergence behavior. Done means the engine exposes the requested rewind capability and the wrapper can reuse the longest shared token prefix for one-shot requests without changing cold-inference outputs; identify the relevant tests while tracing these entry points.