fix(prompt-cache): snapshot bucket sizing is not model-aware for the attention-cache families, so large-context reuse silently never happens
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem
`snapshot_family_is_model_aware` in `src/server/prompt_cache/snapshot_sizing.rs:145-171` does not list `gemma3`, `llama4` or `afmoe`, and none of the three classify as `KvArchKind::Hybrid` or `KvArchKind::PureSsm` at `snapshot_sizing.rs:154`, so `recommend_model_snapshot_capacity_from_config` returns `None` for all of them and they fall back to the fixed `DEFAULT_SNAPSHOT_CAPACITY_BYTES` of 512 MiB (`src/server/prompt_cache/policy.rs:178`). A fixed bucket is the right policy for a recurrent family, whose state is O(1) in context, and the wrong policy for an attention-cache family, whose snapshot is O(context).
Consequence: once a conversation's snapshot exceeds the bucket it is rejected `Oversized` (`src/server/prompt_cache/store.rs:832-839`) and the prompt cache silently does nothing for exactly the long conversations it exists to accelerate. Llama 4 Scout stores roughly 192 KiB per token (48 layers, 8 KV heads, head dim 128, fp16, K and V), so a conversation past about 2731 tokens cannot be stored at all.
## Evidence
Measured on the #1752 real-checkpoint validation run (`mlxcel-server` on port 19335, prompt cache on, entry sizes read from the `prompt-cache snapshot inserted ... bytes=` log line).
- `models/llama-4-scout-17b-16e-instruct-4bit`: a 1817-token completion entry was 402,653,712 bytes (384 MiB) and the next turn's 2352-token entry was 503,317,008 bytes (480 MiB), which is within 6 percent of the 512 MiB bucket after only two turns of a roughly 1300-token conversation. A third turn cannot fit. Both figures are step-aligned buffer sizes (2048 and 2560 tokens at 192 KiB per token), so the stored bytes run ahead of the logical token count.
- `models/gemma-3-4b-it-4bit`: entries grew 188,253,128 bytes at 1320 tokens, 341,346,192 at 1412, and 306,661,264 to 318,572,432 across the later turns. One turn stores two entries (a `Boundary` one and a `Completion` one), so a single three-turn conversation already pushes the bucket into continuous eviction.
## Second defect in the same function
`model_type()` at `snapshot_sizing.rs:136-143` prefers `text_config.model_type` (`src/execution/config_fields.rs:82-84`), and the checkpoints under `models/` report `gemma4_text`, `gemma4_unified_text` and `muse_glimmer_text`, none of which equal the `gemma4` or `muse_glimmer` entries in the list. The model-aware path is currently dead code for its existing entries too.
## Suggested direction
Size the bucket from the model's real per-token KV footprint (layers, KV heads, head dim, dtype) times a target context rather than from a family name, and match `model_type()` against the text-config spellings the checkpoints actually report, or normalize the `_text` suffix before the comparison. Whichever is chosen, it changes a default memory policy, so it needs measurement on real checkpoints rather than an assertion.
## Acceptance criteria
- [ ] A Llama 4 Scout or Gemma 3 conversation whose snapshot exceeds the current 512 MiB bucket inserts and is reused on the next turn, with no `Oversized` snapshot rejection recorded.
- [ ] `gemma4`, `gemma4_unified` and `muse_glimmer` checkpoints take the model-aware path, covered by a test that uses the `text_config.model_type` spellings the checkpoints actually carry.
- [ ] The change to default memory use is measured and reported on at least two real checkpoints, not asserted.
## Verification
```bash
cargo test --profile test-fast --features metal,accelerate prompt_cache::snapshot_sizing
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
```
Then a real multi-turn run against `models/llama-4-scout-17b-16e-instruct-4bit` and `models/gemma-3-4b-it-4bit`, comparing the `prompt-cache snapshot inserted` sizes and the oversized-rejection counter before and after.
Found by the security and performance review of PR #1752 and deliberately not fixed inside it. Refs #1335 and #1752. Sibling follow-ups: #1754, #1755, #1756, #1760.
Contributor guide
Research direction
Start with snapshot_family_is_model_aware and model_type() in src/server/prompt_cache/snapshot_sizing.rs, then inspect recommend_model_snapshot_capacity_from_config and the Oversized path in src/server/prompt_cache/store.rs. Run the prompt-cache snapshot-sizing tests first, then the real Llama 4 Scout and Gemma 3 multi-turn runs described in the issue. Done means model-aware sizing is exercised for the reported text-config spellings, large snapshots are reused without Oversized rejection, and memory effects are measured on at least two checkpoints.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100