Skip to content

Reuse cached chat transcript prefixes#370

Open
aleroot wants to merge 4 commits into
ml-explore:mainfrom
aleroot:auto-chat-prefix-cache-reuse
Open

Reuse cached chat transcript prefixes#370
aleroot wants to merge 4 commits into
ml-explore:mainfrom
aleroot:auto-chat-prefix-cache-reuse

Conversation

@aleroot

@aleroot aleroot commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

Proposed changes

Changes ChatSession to reuse the KV cache for chat prefixes automatically.
Instead of repeatedly prefilling prompt-rendered chat-prefix tokens on later turns, the session keeps an exact token ledger for the prompt and generated assistant output. On the next turn it re-renders the full transcript, checks that the cached tokens are an exact prefix and that KV cache offsets match, then prefills only the new suffix. If anything does not line up, it falls back to full prefill or throws when the prompt configuration changed.
Reuse is intentionally disabled for unsafe paths like logit processors, speculative decoding, bounded/non-reusable caches, and multimodal or masked inputs.

This should improve performance on coding agents using the library.

Notes

This is the first step to make a Codex-like app based on the library. At the moment prefix reuse is enabled just for KVCacheSimple with no logit processors, no speculative decoding, and no dynamic KV quantization. The safest subset. Each opt-out is a known performance floor with a clear unlock path:

  • RotatingKVCache — The transcript prefix eventually gets overwritten once offset > maxSize. A future two-tier cache (static prefix + rotating suffix) would enable reuse for the dominant production config (anyone setting maxKVSize).

  • Logit processorsPenaltyProcessor's TokenRing state is stale when skipping the prefix. On a prefix-reused turn, calling processor.prompt() with the prefix tokens before generation would sync the rings. Small mechanical change once plumbing is ready.

  • Speculative decodingtrimPromptCache after rejected draft tokens can desync KV offsets from the ledger. The speculative iterator already tracks accepted/rejected per round — wiring GenerationTrace.committedTokens to the net effect would unlock this.

The supportsStaticPrefixReuse / allowsPrefixReuse gates are single-point switches designed for independent enablement.

Important

Once #360 lands, PromptLedger.Message should carry the structured tool metadata, not just role and content, so cached replay preserves correlated tool calls exactly.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@aleroot
aleroot force-pushed the auto-chat-prefix-cache-reuse branch 3 times, most recently from f4c9879 to c180cb5 Compare June 25, 2026 17:57
@aleroot

aleroot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor Author

I have reworked completely from my previous #367 since that was too simplistic

@aleroot
aleroot force-pushed the auto-chat-prefix-cache-reuse branch 3 times, most recently from b2a9f7b to 4553688 Compare June 27, 2026 18:45
aleroot added 4 commits June 29, 2026 22:50
ChatSession now keeps an exact token ledger next to the KV cache and reuses a cached prefix only when the re-rendered transcript starts with those tokens and all KV offsets match. Otherwise it falls back to full prefill or throws on prompt configuration changes.

Reuse is disabled for stateful or unsafe paths, including logit processors, speculative decoding, bounded/non-reusable caches, and multimodal or masked inputs.

Tests cover suffix-only prefill, fallback behavior, prompt/tool/context mutations, clear behavior, and cache capability reporting.
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