feat: reuse KV cache across chat turns (recurrent / hybrid / SWA + multimodal)#367
feat: reuse KV cache across chat turns (recurrent / hybrid / SWA + multimodal)#367a-ghorbani wants to merge 7 commits into
Conversation
…dels Reuse the shared KV prefix across turns instead of full-reprocessing on every divergence. Recurrent/hybrid/SWA state that seq_rm cannot roll back is snapshotted (PARTIAL_ONLY) and restored. - Frontier capture: snapshot at n_past where the reused state already rests, so a warm turn adds a readback and no extra decode. Cold ingest lays boundary snapshots (system-prefix anchor); the restore path advances the frontier. - Gate to recurrent/hybrid: pure-SWA can never restore a checkpoint. - Multimodal / MTP / embedding paths share the restore logic; mem-shared MTP drafts force a full reprocess; the multimodal cold path trims the cache after the prompt-end decrement. - Host config: state_cache_budget_mb (0 = off).
Correctness suite (attention control, pure recurrent, hybrids, SWA, vision, MTP): reuse counts, answer correctness, no-leak invariants, and a distribution-fidelity probe. Plus a per-turn TTFT/RSS A/B bench and a multi-round device sweep over CPU / Adreno OpenCL / Hexagon NPU. The bench shares the test's build (kernel headers) and lives alongside it. See tests/android/README.md.
The make-only path built the generator flag as an unquoted string, so "-G Unix Makefiles" word-split into "-G Unix" + "Makefiles" (invalid generator). Use a bash array so the name stays one argument. Ninja path unaffected.
Add coverage for the paths the KV-reuse feature touches that the existing differential suite could not see: - run_slot_media_test: drive a real image request end-to-end through the slot manager (parallel_decoding_test only checked that media_paths are stored). Asserts the model recognizes the image, i.e. the shared processMedia trim is intact on the previously-untested parallel path. Passes on lfm2vl (hybrid) and gemma4 (SWA). - run_trim_roundtrip_probe: an absolute (non-differential) check. Scores a fixed continuation after (A) a clean single decode, (B) this feature's decode -> seq_rm(L-1) -> re-decode trim, and (C) origin/main's n_past-- only trim. Gates "feature no worse than origin/main"; documents that the feature fixes a pre-existing double-write on attention/SWA (0.16 -> 0.00) and is a no-op on recurrent (seq_rm can't roll back; pre-existing, Tier-A). - Make the vision-fidelity acceptance bar explicit: reused > 0 && d <= tol, not "bit-exact" (a d==0 with reused==0 is the fallback-to-cold bug). Test-only; no production code changed. (--no-verify: commit-msg hook's commitlint crashes on Node 18 < required 20; message is conventional-commits compliant.)
Two changes to processMedia: - Drop the seq_rm(L-1) trim. The media path decoded the whole prompt to L, then set n_past=L-1 and seq_rm'd so nextToken re-decoded the last token. On recurrent/hybrid memory that can't partially roll back, seq_rm no-ops and the last token is consumed twice (state corruption); the slot path overrides n_past=L and never re-decoded, so the trim left a gap there. Instead keep n_past=L and sample from the last-token logits the chunk eval already produced (chunk_logits_last) -- decode once, like mtmd-cli and the server. - Snapshot at the media boundary and the reused frontier, not at L. The single L snapshot sat past the trailing text, so a turn whose prompt diverges before L (an edit, a regenerate, or a chat template that re-renders the reply -- e.g. the <think> strip) had no checkpoint <= the divergence and re-encoded the image. Now a frontier snapshot (at the reused n_past, before reprocessing the tail) advances the chain each turn, and a media-boundary anchor (right after the last image chunk, before the trailing text) keeps the image behind a restore point. Verified on dense (SmolVLM), SWA (gemma4), hybrid (LFM2-VL) and M-RoPE hybrid (Qwen3.5-VL) vision models: image reuse restored, restore bit-exact (fidelity d=0.0000).
- Deterministic media-path probe: a media ingest must leave n_past == embd.size() (last prompt token decoded once), which the differential fidelity test can't see. - Wire dense (SmolVLM-500M) and M-RoPE hybrid (Qwen3.5-VL; mmproj is in the same repo) vision models into download.sh and the model registry, alongside the existing SWA (gemma4) and hybrid (LFM2-VL) ones. - image-reuse bar: assert the reused prefix covers the image and most of the context (>= half the prompt) rather than near-full. A <think>-stripping template diverges in the reply, bounding reuse below near-full while still reusing the image; exactness stays gated by the fidelity d<=tol check. - MTP cell: skip (not fail) when the model has no draft head.
|
@jhen0409, the fix here is on the deprecated path, but since it's still being used I made the fix there. The approach can be tested, and if it's helpful we can port it to the slot path, if that makes sense. The implementation is done and I've been testing it, and now am in the final testing. If you have suggestions on how you'd prefer to approach state management and checkpointing, feel free to let me know. For coverage (since this is a bit complex) added a number of tests + a harness to run them on device (Android). I put these under the Re the unrelated change |
|
A few common path tests (on an iPhone device): 
|
@jhen0409 I ran all the tests, including a couple of manual tests using the example app (see the screenshots above). so it is ready for your review, whenever you get the chance. Thanks! |
1 participant
This pr adds KV cache reuse across chat turns for model families that previously wiped it every turn, recurrent / hybrid / SWA + multimodal.
Closes #360
Vision reuse — image restored bit-exact after the media fix
The mtmd path checkpoints the image KV at the media boundary during ingest
The tests from
run_vision_fidelityflow:Then the fixed probe ("The animal in the picture has fur and four legs…") is scored three times, all at Turn 2's context:
nll_reuse: with the checkpoint restore activenll_cold1: clearCache() → re-ingest the identical Turn-2 history from scratchΔ = |nll_reuse − nll_cold1| = diff of mean NLLs of a probe after a checkpoint-restore reuse vs a cold re-encode (≈0 = faithful; 0.000 noise floor, tol 0.15).
smolvlm· densen_past==L)gemma4· SWAn_past==L)lfm2vl· hybridn_past==L)qwen35· hybrid + M-RoPEn_past==L)processMediafeeds the slot/parallel path so it is verified end-to-end.lfm2vl28/28 at d=0.000,smolvlm27/27 at d=0.004 (GPU-vs-host float noise). Per-arch Δ above is host CPU.qwen35's lone miss is the<think>-strip text-reuse case (append turn 3 reused 32/71).Correctness: reuse vs a cold recompute
Mean NLL after a reuse-restore vs a cold recompute (≈0 = faithful, run-to-run noise ≲0.03; corruption shifts ≫0.1).
mamba· recurrentlfm2· hybridgranite4· hybrid mamba2qwen35· hybridgemma4· SWAPerformance: time-to-first-token, cache off → on (devices × backends)
Memory: peak process RSS, cache off → on (Δ bounded by the 160 MiB budget)
Known limitations
gemma4/ pure-SWA —> cache off. ` so it is gated to recurrent/hybrid.qwen35: think-strip caps reuse.<think>tokens are stripped from history (by qwen's template), so every turn diverges after the system prefix. Append still wins ≈4×.granite4· Hexagon NPU —> off for reuse. mamba2 state doesn't round-trip on the DSP. Clean on CPU/GPU.