Skip to content

fix(kv): key completion checkpoints on round-trip-stable prompt prefix - #45

Merged
marzukia merged 3 commits into
mainfrom
fix/disk-kv-prompt-key
Aug 5, 2026
Merged

marzukia merged 3 commits into
mainfrom
fix/disk-kv-prompt-key

Conversation

@marzukia

Copy link
Copy Markdown
Owner

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_boundary already receives prompt_token_ids separately, so it now keys on tokens[: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 same target_len, so the existing len(key) == token_offset guard and the _cache_offset_matches restore 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_trimmable keeps full-length keying for them. No-op for those models.

Also adds [kv_restore_divergence] logging in DiskCheckpointIndex.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.

HIT cached=24393  prompt=24685  to_prefill=292
HIT cached=24545  prompt=24907  to_prefill=362
HIT cached=24768  prompt=25149  to_prefill=381

Tests

  • 75 checkpoint tests pass (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.
  • New test_completion_checkpoint_keyed_on_stable_prompt_prefix drives the real _cleanup_finished store 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.
  • Prompts under 256 tokens keep the old full-key path, so short-prompt behaviour is unchanged.

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.

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.
@marzukia marzukia closed this Jul 20, 2026
@marzukia marzukia reopened this Aug 5, 2026
@marzukia
marzukia merged commit b4bab93 into main Aug 5, 2026
7 checks passed
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