Repository navigation
fix(kv): key completion checkpoints on round-trip-stable prompt prefix - #45
Merged
Merged
Conversation
The disk-KV completion store keyed each checkpoint on prompt + generated output tokens. The model output (empty <think></think> scaffold, tool-call formatting) does not re-tokenize to the same IDs when the client echoes the conversation back next turn, so the stored key was never a prefix of the next prompt. Every multi-turn request fell back to the system-prompt checkpoint and cold-prefilled the whole conversation. Persist the completion checkpoint keyed on the prompt prefix only (trimmable caches slice the KV to the prompt length so key length, token_offset and the stored KV all stay consistent). Recurrent/hybrid caches keep the full key. Add a lookup divergence diagnostic that logs where a near-match checkpoint left the incoming prompt.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The completion-boundary checkpoint write keyed the content index on the full token sequence (prompt + model output). Model output tokens (the empty
<think>scaffold Qwen3.5 emits in non-thinking mode, tool-call formatting) do not re-tokenise to the same IDs when the client rebuilds the conversation from text. So a growing per-turn checkpoint never matched the next turn as a prefix, and every multi-turn request fell back to the system-prompt checkpoint and cold-prefilled everything after it. TTFT was ~121s on a 70k conversation.The lookup layer was already correct (there is a test documenting that keying on generated output does not help). The bug was purely on the write side: the reliably-firing completion store keyed on the volatile output tail, so no prompt-keyed growing checkpoint was ever deposited.
Fix
Key the completion checkpoint on the round-trip-stable prompt prefix only, not prompt + output.
_disk_persist_boundaryalready receivesprompt_token_idsseparately, so it now keys ontokens[:stable_prefix_len]and slices the KV blob to that length via_slice_kv_to_length. Key length,token_offset, and the physical KV rows are all set from the sametarget_len, so the existinglen(key) == token_offsetguard and the_cache_offset_matchesrestore gate both hold. A restore can never over-read into output-token KV that is not in the query. The few dropped output tokens re-prefill cheaply next turn.Recurrent/hybrid caches (GatedDeltaNet) cannot be prefix-sliced (their state has already consumed the output tokens), so
_cache_has_non_trimmablekeeps full-length keying for them. No-op for those models.Also adds
[kv_restore_divergence]logging inDiskCheckpointIndex.lookup()that reports the first divergence position and the two differing token IDs when a checkpoint shared more tokens than were restored but was not a clean prefix.Result
Live on helium (Qwen3.5-122B-A10B-oQ4-mtp). Within-thread chaining is restored: consecutive turns restore a climbing checkpoint and prefill only the delta instead of thousands of tokens.
Tests
test_disk_kv_checkpoint.py,test_scheduler_disk_kv_hook.py,test_delta_kv_checkpoints.py), including all existing round-trip, quarantine, and delta tests.test_completion_checkpoint_keyed_on_stable_prompt_prefixdrives the real_cleanup_finishedstore site and asserts the checkpoint is keyed on the prompt, not prompt + output, and that a next-turn prompt with different re-tokenised output still matches. It is non-vacuous: removing the fix makes it fail.Out of scope
The divergence logging surfaced a separate cross-thread break: OpenCode injects a per-thread session ID and Discord channel ID near the top of the system prompt, so different threads diverge at ~token 693 and cannot share the static prefix. That fix lives in the OpenCode/Kimaki prompt assembly and is tracked separately.