LFM2: opt-in prefix-KV reuse; KVCacheSimple fork isolation + tests - #130
Merged
Conversation
MERERUN_LFM2_PREFIX_KV_CACHE=1 ports the Qwen-family in-memory prompt-prefix cache to the LFM2 chat runtime: forked layer caches (attention KV plus short-conv states, both already forkable) are stored at prefill chunk boundaries, the longest matching token prefix seeds later requests so only the prompt tail re-prefills, entries are bounded to 4 with the shared retention planner, and prefixKVCacheStats() exposes hits/misses/reused tokens. Chunk-boundary checkpoints only for now — the Gemma4-style semantic chat-template checkpoints are not yet derived for LFM2. Measured against a live api serve: a ~2.9k-token prompt repeats at 0.3s end-to-end versus 11.7s cold; a normal system+question prompt repeats coherently at 1.6s versus 3.8s. Note: LFM2.5 is an A1B MoE and its serve responses are nondeterministic run-to-run at temperature 0 even with the cache disabled (gather-scatter accumulation order, the same class as the Ornith 35B finding), so identity across requests is not a valid gate for this model; coherence plus the timing profile is. KVCacheSimple.fork now binds fresh MLXArray wrappers instead of sharing the parent's. The shared-wrapper fork was safe in practice for append-only usage (writes land at positions the fork never reads, and capacity growth rebinds the parent's property), but the aliasing is a loaded gun for any future consumer that interleaves reads and writes across fork generations; KVCacheForkTests pins the isolation and divergence contracts either way. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
SwiftLint (force_cast) failed CI on the two `as!` casts in the new fork tests. Merging main also combines the prefix-cache prefill restructure with #124's forwardPrefill last-position slicing; the per-chunk stored logits remain the last-position slice a seed consumer reads. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Member
Author
|
Post-merge re-measurement of the determinism caveat (now retracted in the description): The "nondeterministic at temp 0, gather/scatter atomics" baseline was taken against the stale AOT metallib that #131 root-caused and replaced (pre-0.30 kernels driven by 0.31.x host dispatch — corrupts any decode past ~1024 tokens of context, nondeterministically). Re-measured on current main (
LFM2.5 A1B is fully deterministic at temperature 0. Exact response identity is a valid regression gate for this model; the Ornith 35B "MoE atomics" theory is likewise retracted (same metallib root cause, fixed in #131). 🤖 Generated with Claude Code |
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
1 participant
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.
Lane I (part 2, final) of the platform perf sweep — the LFM2 prefix-KV gap from the verified map.
Feature
MERERUN_LFM2_PREFIX_KV_CACHE=1ports the Qwen-family prompt-prefix cache to LFM2: forked layer caches (attention KV + short-conv states) stored at prefill chunk boundaries, longest-prefix seeding, 4-entry bound via the shared retention planner,prefixKVCacheStats()for observability. Chunk-boundary checkpoints only (semantic chat-template checkpoints not yet derived for LFM2 — noted in docs).Measured (live
api serve)Determinism caveat — RETRACTED (2026-07-03, post-merge)
LFM2.5 is an A1B MoE: serve responses are nondeterministic at temperature 0 even with the cache disabled (gather/scatter accumulation order — same class as the Ornith 35B finding).This was wrong. The baseline was measured against a stale AOT metallib whose 2-pass SDPA kernels mismatched the mlx 0.31.x host dispatch (#131) — the same root cause as the Ornith 35B gibberish, which was likewise not MoE atomics. Re-measured on main after #131 with the stamped metallib:
text-chat-lfm25-a1b-8bitat temperature 0 is byte-identical 3/3 runs at 147 / 908 / 3509 prompt tokens via the CLI (including visible thinking content), and identical 3/3 through/v1/chat/completions(long prompt at two token budgets, plus a content-bearing short prompt). Exact-match IS a valid gate for this model.Fork hygiene
Chasing the (ultimately MoE-caused) divergence surfaced that
KVCacheSimple.fork()sharesMLXArraywrapper objects with the parent, whose subscript writes rebind those wrappers in place. Analysis shows the append-only usage kept it correct in practice (disjoint read/write regions + lazy views + growth rebinds) — but the aliasing is a loaded gun for future consumers.fork()now binds fresh wrappers, andKVCacheForkTestspins the isolation + divergence contracts.🤖 Generated with Claude Code