feat(core): shared snapshot serialization for KVCache / RotatingKVCache / ChunkedKVCache so Gemma 3, AFMoE and Llama 4 join exact-prefix prompt-cache reuse #1335

Summary

Prompt-cache reuse in the server has two producers: detached per-layer KVCache sets for families whose KV lives in the scheduler's CachePool (radix prefix matching, partial-block adoption), and exact-prefix ModelStateSnapshot copies for families that own their caches and opt in through LanguageModel::supports_snapshot_reuse(). The opt-in currently requires each family to hand-write per-cache-type serialization, and only Gemma 4 and Muse Glimmer have done so for the attention cache types. Gemma 3, AFMoE (which reuses Gemma 3's Cache enum) and Llama 4 hold ordinary KVCache / RotatingKVCache / ChunkedKVCache state in ModelOwnedSequenceState and are silently excluded from every prompt-cache path, so every multi-turn request on those families re-prefills the whole conversation. This change moves the KV-type serialization into one shared module and opts those three families in, with the same truncation rule Gemma 4 already uses.

Current behavior

Donation is skipped for any model-owned family that does not declare snapshot reuse (src/server/batch/scheduler.rs:2982-2996):

if self.model.supports_snapshot_reuse() {
    self.insert_model_state_snapshot(seq_id, &ctx, tokens, SnapshotOrigin::Completion);
    return;
}
let backend = self.cache_pool.get_mut(seq_id).map(|s| s.backend)
    .unwrap_or(SequenceStateBackend::ModelOwned);
// Other `ModelOwned` families carry no detachable cross-request KV.
if backend == SequenceStateBackend::ModelOwned {
    return;
}

Lookup is gated the same way (src/server/batch/scheduler.rs:2026-2036): supports_snapshot_reuse() == false yields SnapshotLookupOutcome::NoCandidate before the store is consulted.

Families affected, with the cache types they own:

Gemma 4 already carries a complete, tested serializer for Standard and Rotating caches (src/models/gemma4.rs:1196-1350, truncation rule at :1370-1440), and Muse Glimmer carries a near-identical copy (src/models/muse_glimmer_cache.rs:58-206). The scalar tags kv_cache_mode_to_i32 / kv_cache_mode_from_i32 are private to gemma4.rs:1139-1160. No serializer exists for ChunkedKVCache anywhere (grep -rn "chunked\." src/models/ is empty).

Cache types that are therefore not prefix-cacheable today, for the record:

• KVCache / RotatingKVCache inside a model-owned enum of gemma3, afmoe, gpt_oss, exaone4, ministral3, step3p5 (the last four hold them in RefCell<Vec<Cache>>, see src/models/gpt_oss.rs:1107, src/models/exaone4.rs:685, src/models/ministral3.rs:645, src/models/step3p5.rs:1026, and are single-stream).
• ChunkedKVCache (llama4 only).
• KimiDeltaCache (src/models/kimi_linear.rs:395), RGLRUCache (src/models/recurrent_gemma.rs:127), Rwkv7Cache (src/models/rwkv7.rs:78; has snapshot_into but the family never opts in). These three are blocked by fix(models): RWKV7, RecurrentGemma, and KimiLinear rebuild their recurrent state every decode step, so generation runs with no history #1220 (their state is rebuilt every forward) and stay out of scope here.
• RingSlidingKVCache (src/lib/mlxcel-core/src/cache/ring.rs:54, Unlimited-OCR only), by design.

DetachedRotatingKVCache (src/lib/mlxcel-core/src/cache/detach.rs:621) exists but has no caller outside its own module; the CachePool detach path only handles Vec<KVCache>.

Expected behavior

1. A shared module src/models/kv_snapshot.rs exposes, for the three standard attention cache types, snapshot_into(cache, snapshot, prefix) -> Result<(), String>, restore_from(cache, snapshot, prefix) -> Result<(), String> and snapshot_truncatable_to(snapshot, prefix, target_len) -> bool, plus kv_cache_mode_to_i32 / kv_cache_mode_from_i32 moved out of gemma4.rs. Gemma 4 and Muse Glimmer call the shared functions; their behavior and tensor names do not change.
2. Serialization contract (all scalars stored through recurrent_snapshot::push_i32, all tensors copied through ModelStateSnapshot::push_tensor, which already deep-copies; only KVCacheMode::Fp16 is accepted, any other mode returns Err so the donation is skipped, exactly as Gemma 4 does today):
   • KVCache: {prefix}.standard.keys [B, H_kv, T_alloc, D], {prefix}.standard.values [B, H_kv, T_alloc, D], {prefix}.standard.offset (i32, logical length; T_alloc >= offset because buffers grow by step), {prefix}.standard.mode (i32 tag). Restore sets keys, values, offset, mode = Fp16. Truncation to target_len <= offset: always allowed (restore then KVCache::trim(offset - target_len)).
   • RotatingKVCache: {prefix}.rotating.keys, {prefix}.rotating.values (physical ring buffers, [B, H_kv, min(offset, max_size) or the over-window prefill length, D]), {prefix}.rotating.max_size, {prefix}.rotating.buffer_size, {prefix}.rotating.offset, {prefix}.rotating.start_position, {prefix}.rotating.idx, {prefix}.rotating.step, {prefix}.rotating.mode, {prefix}.rotating.turbo_seed. Restore goes through RotatingKVCache::restore_fp16_snapshot_state(RotatingKVCacheSnapshotState, keys, values) (src/lib/mlxcel-core/src/cache.rs:5260). Truncation allowed only when RotatingKVCacheSnapshotState::can_truncate_to(target_len) (cache.rs:4185) holds, i.e. buffer_size == 0 && idx == offset && offset <= max_size && target_len <= offset: the ring has not wrapped so logical token t sits at slot t.
   • ChunkedKVCache (new): {prefix}.chunked.keys, {prefix}.chunked.values ([B, H_kv, offset - start_position, D] after maybe_trim_front), {prefix}.chunked.chunk_size, {prefix}.chunked.offset, {prefix}.chunked.start_position. Restore assigns the three scalars and the two tensors (the step field stays at its constructor default of 256; it only affects future growth). Truncation to target_len allowed only when start_position == 0 && target_len <= offset: once the front has been trimmed the window a fresh prefill would hold, [target_len - chunk_size, target_len), is wider than [start_position, target_len), so a truncated restore would attend over fewer tokens than a cold prefill and is not exact. When allowed, restore then slice keys/values to [.., .., :target_len, ..] and set offset = target_len.
3. Gemma 3 and AFMoE: gemma3::Cache gains snapshot_into / restore_from / snapshot_truncatable_to / truncate_to delegating to the shared module with the same layer{idx} prefix scheme Gemma 4 uses, and Gemma3Model / AfmoeModel implement supports_snapshot_reuse() -> true, snapshot_sequence_state, restore_sequence_state, snapshot_truncatable_to, restore_sequence_state_truncated following the Gemma 4 template at src/models/gemma4.rs:5968-6070. Family tags: "gemma3", "afmoe".
4. Llama 4: Llama4Cache gains the same four methods (Chunked arm uses the chunked.* contract, Regular arm the standard.* contract) and Llama4Model opts in with family tag "llama4".
5. A second request whose token vector extends a finished request's prompt + generated vector on any of these families restores the snapshot and prefills only the appended tokens; the generated text is token-identical to a cold run at temperature 0.

Implementation plan

1. Create src/models/kv_snapshot.rs (register in src/models/mod.rs). Move kv_cache_mode_to_i32 / kv_cache_mode_from_i32 from src/models/gemma4.rs:1139-1160 into it (make the error message family-neutral). Add:
   • pub(crate) fn snapshot_standard(cache: &KVCache, snapshot: &mut ModelStateSnapshot, prefix: &str) -> Result<(), String> and restore_standard(cache: &mut KVCache, snapshot, prefix) -> Result<(), String>: lift the Self::Standard arms of gemma4.rs:1204-1229 and :1283-1308 verbatim.
   • snapshot_rotating / restore_rotating: lift the Self::Rotating arms of gemma4.rs:1230-1273 and :1309-1349.
   • snapshot_chunked(cache: &ChunkedKVCache, ...) / restore_chunked(cache: &mut ChunkedKVCache, ...) per the contract above. ChunkedKVCache.step is private (src/lib/mlxcel-core/src/cache.rs:5497); leave it at the constructor value.
   • standard_truncatable_to, rotating_truncatable_to (lift gemma4.rs:1370-1440), chunked_truncatable_to(snapshot, prefix, target_len) -> bool returning start_position == 0 && 0 <= target_len <= offset.
   • truncate_standard(cache, target_len) (calls KVCache::trim), truncate_rotating (reuse the Gemma 4 Cache::truncate_to logic at the rotating arm), truncate_chunked (slice keys/values on axis 2 to target_len, set offset = target_len).
   • Use recurrent_snapshot::{push_i32, push_optional, restore_i32, restore_optional} (src/models/recurrent_snapshot.rs:24-51) for every field.
2. Replace the bodies of gemma4::Cache::{snapshot_into, restore_from, snapshot_truncatable_to, truncate_to} and muse_glimmer_cache::MuseCache::{snapshot_into, restore_from, ...} with calls into kv_snapshot. Keep the existing tensor names and error strings so src/models/gemma4_tests.rs and src/models/muse_glimmer_tests.rs snapshot tests pass unchanged.
3. src/models/gemma3.rs: add to impl Cache (:644) the four methods delegating to kv_snapshot. In impl LanguageModel for Gemma3Model (:1210-1240), add supports_snapshot_reuse, snapshot_sequence_state (iterate self.sequence_state.with_sequence_state_ref(seq_id, ...), prefix layer{idx}, return None on any Err after a tracing::warn!), restore_sequence_state (check snapshot.family() == "gemma3", build self.model.make_caches(), restore_from per layer, self.sequence_state.replace_sequence_state(seq_id, state)), snapshot_truncatable_to (all layers agree), restore_sequence_state_truncated (re-check, restore, then truncate_to per layer). Mirror src/models/gemma4.rs:5968-6070.
4. src/models/afmoe.rs: same five trait methods on AfmoeModel (:1658-1690 region), family tag "afmoe", using the gemma3::Cache methods from step 3.
5. src/models/llama4.rs: add snapshot_into / restore_from / snapshot_truncatable_to / truncate_to on Llama4Cache (:38), then the five trait methods on the model wrapper (:1713-1740 region), family tag "llama4". The VLM wrapper around Llama 4 (if it forwards LanguageModel calls) must forward the five new methods too; check src/loaded_model.rs dispatch macro at :370 picks them up automatically (it forwards every trait method).
6. src/server/prompt_cache/hybrid_ssm.rs: no change; gemma3 / afmoe / llama4 are not in HYBRID_SSM_MODEL_TYPES, so APC block hashing stays enabled for them. The snapshot path does not use block hashes, so nothing else in src/server/prompt_cache/ changes.
7. docs/turbo-kv-cache.md, section "Exact-prefix snapshots for recurrent state": list the snapshot-capable families and the per-type truncation rule (standard: always; rotating: unwrapped only; chunked: untrimmed only).

Validation

(a) Unit tests, src/models/kv_snapshot_tests.rs:

• standard_round_trip_preserves_keys_values_offset: build a KVCache, push 37 tokens of random K/V, snapshot, restore into a fresh cache, assert offset == 37 and update_and_fetch on one more token returns K/V whose first 37 positions are bitwise equal.
• rotating_round_trip_preserves_ring_state_after_wrap: RotatingKVCache::new(16), push 40 tokens one at a time, snapshot, restore, assert snapshot_state() equals the original (offset, idx, max_size) and that the next update_and_fetch returns the same 16-token window as the original cache.
• rotating_truncatable_only_while_unwrapped: at 10 of 16 tokens truncatable_to(7) == true; at 40 tokens == false.
• chunked_round_trip_after_front_trim: ChunkedKVCache::new(8), push 20 tokens, maybe_trim_front, snapshot, restore, assert start_position == 12, offset == 20, keys shape [1, H, 8, D].
• chunked_truncatable_only_before_front_trim: at 6 of 8 tokens truncatable_to(4) == true; at 20 tokens == false.
• non_fp16_mode_is_rejected: a KVCache::new_with_mode(KVCacheMode::Int8) snapshot returns Err.
• src/models/gemma3_tests.rs: snapshot_restore_matches_cold_decode: run a 24-token prefill plus 8 decode steps through Gemma3Model with a synthetic config (sliding_window = 8, sliding_window_pattern = 2, 2 layers), snapshot at 32, restore, decode 4 more tokens, compare logits to a cold 36-token run at the same positions (allclose, atol 1e-3 in f16).
• src/models/llama4_tests.rs: same shape with attention_chunk_size = 8.

(b) Real checkpoint:

./target/release/mlxcel-server -m models/gemma-3-4b-it-4bit --prompt-cache --port 8080
# turn 1
curl -s localhost:8080/v1/chat/completions -d '{"model":"gemma-3-4b-it-4bit","messages":[{"role":"system","content":"<a 300-token system prompt>"},{"role":"user","content":"List three facts about Jupiter."}],"max_tokens":64,"temperature":0}'
# turn 2: same system prompt + turn-1 assistant text + a new user message
curl -s localhost:8080/v1/chat/completions -d '{... turn-1 messages + assistant reply + {"role":"user","content":"Now Saturn."} ...,"max_tokens":64,"temperature":0}'
curl -s localhost:8080/v1/cache/stats

Acceptance: /v1/cache/stats reports one snapshot hit for turn 2 with matched_len equal to the turn-1 prompt + generated length; turn-2 completion text is byte-identical to the same two requests served with --prompt-cache off; prompt_tokens time for turn 2 drops to the appended-tokens cost. Repeat with mlx-community/Llama-4-Scout-17B-16E-Instruct-4bit (chunked cache) and the AFMoE validation checkpoint mlx-community/Trinity-Nano-Preview-4bit. For Gemma 3 also run a third turn after the conversation exceeds sliding_window (1024) tokens and confirm the store still hits (exact restore) while snapshot_truncatable_to returns false for a shorter divergent prompt (a SnapshotDiverged reject, not a corrupt restore).

Acceptance criteria

• src/models/kv_snapshot.rs owns the standard / rotating / chunked serializers and truncation rules; gemma4.rs and muse_glimmer_cache.rs delegate to it with unchanged tensor names.
• Gemma 3, AFMoE and Llama 4 return supports_snapshot_reuse() == true and implement snapshot / restore / truncatable / truncated-restore.
• Turn-2 snapshot hit on mlx-community/gemma-3-4b-it-4bit and mlx-community/Llama-4-Scout-17B-16E-Instruct-4bit with byte-identical output versus prompt cache off. Partly met in PR feat(core): share KV snapshot serialization across attention cache families #1752: the turn-2 hit and exact restore are confirmed on both, including a wrapped-ring Gemma 3 turn over the sliding window, but comparing against --no-prompt-cache does not isolate the snapshot (the server defaults preserve_thinking differently and the cold prefill path itself differs), so the comparison was run with the cache on in both arms. Two gaps remain: trinity-nano-preview-4bit cannot be judged by token identity because it is a documented degenerate preview checkpoint, and Gemma 3 diverges reproducibly when the stored snapshot covers exactly one decoded token.
• Non-Fp16 KV modes are refused at snapshot time with a warning, never restored with a mismatched layout.
• docs/turbo-kv-cache.md lists the snapshot-capable families and per-type truncation rules.
• cargo test --workspace --profile test-fast --features metal,accelerate passes
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

Out of scope

• gpt_oss, exaone4, ministral3, step3p5: they keep caches in a RefCell<Vec<Cache>> that is not keyed by SequenceId; migrating them to ModelOwnedSequenceState is a separate refactor and a prerequisite for snapshot reuse.
• kimi_linear, recurrent_gemma, rwkv7: blocked by fix(models): RWKV7, RecurrentGemma, and KimiLinear rebuild their recurrent state every decode step, so generation runs with no history #1220.
• Block-level (radix / APC partial) reuse for rotating or chunked caches. Exact-prefix plus the unwrapped-only truncation rule is the correct contract for windowed caches.
• Snapshots of Int8 / Turbo KV modes (tracked separately in prompt-cache-snapshot-non-fp16-kv-mode-gate).