lablup / lablup/mlxcel

fix(prompt-cache): snapshot bucket sizing is not model-aware for the attention-cache families, so large-context reuse silently never happens

Open
#1,761 0 comments 0 reactions 0 assignees View on GitHub
area:inference priority:high status:ready type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.