Skip to content

feat(cache): keep recurrent states where prompts diverge - #2

Merged
mzinner merged 1 commit into
mainfrom
feat/prefix-state-boundaries
Oct 23, 2026
Merged

mzinner merged 1 commit into
mainfrom
feat/prefix-state-boundaries

Conversation

@mzinner

@mzinner mzinner commented Oct 23, 2026

Copy link
Copy Markdown
Owner

Summary

Qwen3.8-Flash-Next is a hybrid model: its recurrent (GDN) state can only be resumed where a state snapshot was saved. The engine saved them at the end of each prompt and at rolling 4096-token checkpoints, so any prompt that shared a long prefix with an earlier one but differed near its end re-prefilled from the start (or from the first 32 tokens). Agent harnesses do that all the time:

  • a new question in the same session: the same system prompt, tools and reminders, a different last message;
  • a side agent (e.g. a memory-extraction subagent) with the same system prompt and a different first user message.

This PR makes the engine save states where prompts actually diverge, and keep the most widely shared one:

  1. State hints (protocol v7). The server sends state_hints: prompt positions where the engine should keep a state.
    • It finds them from the rendered tokens: the end of the system prompt and tools, the start of the last message, and the starts of that message's final content parts.
    • The content parts come from the request as sent. The frontend joins them into one string, and harnesses often put reminders before the question.
    • The engine plans them as state boundaries and skips any within 4 blocks of another planned state.
  2. Retention.
    • The first hint (end of the shared setup) is a boundary state, kept ahead of all other states and of plain KV, at most 4 of them (SPLASH_BOUNDARY_STATES).
    • The other hints are hint states. They enter the cache as the oldest ordinary state and become ordinary once a later prompt resumes from them, so a speculative hint never pushes out an end-of-turn state that plain continuations resume from.
    • StateCache now has one StateRole per entry (ordinary, checkpoint, hint, boundary) instead of a checkpoint flag.
  3. Status reports hint_publications, hint_publication_failures, boundary_entries, boundary_capacity and boundary_demotions.

Measurements

Replay of 31 requests captured from a real agent harness session (CLI mode, ACP mode, and its memory-extraction subagent), sent one at a time in order, generating one token each, on an M5 Pro 64 GB (macOS 27.0.1):

main this PR
prompt tokens reused 63.5% 78.2%
total time, 31 requests 315 s 198 s (-37%)
requests slower than main (>1 s) — 2

Requests that changed by more than a second:

# kind prompt tokens main: reused / time this PR: reused / time
1 CLI session 11,822 0 / 32.7 s 0 / 29.7 s
3 CLI session 11,823 0 / 29.9 s 11,552 / 1.0 s
13 ACP session 17,654 0 / 46.5 s 0 / 48.3 s
14 ACP session 17,650 64 / 44.2 s 17,472 / 1.0 s
15 memory subagent 6,193 64 / 15.1 s 64 / 16.8 s
17 memory subagent 6,117 64 / 15.9 s 5,312 / 2.5 s
22 CLI session 11,848 64 / 29.9 s 8,992 / 8.2 s
28 CLI session 11,894 5,184 / 20.0 s 8,992 / 9.1 s
29 memory subagent 6,353 5,088 / 4.4 s 5,312 / 3.1 s

#1, npanj#13 and npanj#15 are cold full prefills with the same reuse in both runs. Their differences (−3.0 s, +1.8 s, +1.7 s) are run-to-run variation; in an earlier run of this branch npanj#13 took 43.7 s. Decode speed is unchanged: 28.6–32.7 tok/s across the runs, the same as main on this Mac.

Remaining misses:

Tests

  • New:
    • protocol v7 round trip and hint validation (C++ and Python, golden vectors updated);
    • boundary tier: eviction order, cap and demotion, upgrade on reuse;
    • hint states start cold and become ordinary on reuse;
    • server hint selection, including content parts with exact token positions.
  • make test-engine (CPU and Metal suites) passes. The two kv2_g12 verify tiles are still skipped under shader validation, as before.
  • dev/tests/test_server.py (175 tests) and the engine Python tests (417, 11 skipped) pass; ruff check and ruff format --check are clean.

Follow-ups, not in this PR

  • Record content-part offsets during message normalization (normalize_messages) and tokenize once with offsets, instead of finding parts by text in the rendered prompt. That would also cover Anthropic and Responses inputs uniformly.
  • Find message boundaries without the ChatML <|im_start|> convention (from the template itself), for other model families.
  • Let the server mark which hint ends the shared setup, instead of the engine inferring it from position. Alternatively, promote a state to boundary when a different request reuses it.
  • A memory split between the expert cache and the state cache that adapts to the workload (see the 30 GiB result above).

🤖 Generated with Claude Code

The model's recurrent state can only be resumed where a snapshot was
saved: at the end of each prompt and at rolling checkpoints. A prompt that
shared a long prefix with an earlier one but differed near its end - a new
question in the same session, a side agent with the same system prompt -
re-prefilled from the start.

The server now sends state hints (protocol v7): where the system prompt and
tools end, where the last message begins, and where that message's final
content parts begin. The engine keeps a state at each, skipping those within
four blocks of another. The first, shared by other conversations on the same
setup, is a boundary state kept ahead of all other states and of plain KV
(at most four, SPLASH_BOUNDARY_STATES). The others start as the oldest
ordinary state and become ordinary once a later prompt resumes from them,
so a guess never evicts an end-of-turn state. StateCache entries now carry
one StateRole instead of a checkpoint flag.

Replaying 31 requests captured from an agent harness: 63.5% -> 78.2% of
prompt tokens reused, 315 s -> 198 s in total; a new question in a 12K-18K
token session went from 30-44 s to 1 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@mzinner
mzinner merged commit bf079e9 into main Oct 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant