[inureyes]

Summary

The exact-prefix snapshot store had a fixed 512 MiB budget with no operator override, and no notion that a newer snapshot from the same conversation replaces the older one. A single 31B-class conversation snapshot runs 300-370 MB, so the second turn's insert LRU-evicted the first turn's snapshot rather than superseding it, and two concurrent conversations evicted each other. This adds an explicit session-chain supersede at insert time and wires the three existing PromptCacheConfig snapshot fields through CLI flags and MLXCEL_* env vars.

What changed

• src/server/prompt_cache/store.rs: insert_snapshot now removes every stored snapshot whose token vector is a strict prefix of the incoming one, within the same session key and the same model/lora/template/multimodal bucket. It runs before the new entry's bytes are accounted, so enforce_snapshot_caps sees the freed bytes and can admit an extension that would not have fit alongside its own ancestor. A None session key never triggers it (no conversation identity to chain on) and two different session keys never touch each other.
• src/server/prompt_cache/store.rs, policy.rs, routes/cache.rs: new snapshot_supersedes counter, held apart from snapshot_evictions_lru so capacity tuning can distinguish deterministic in-session replacement from real budget pressure. Surfaced on /v1/cache/stats.
• src/server/cli_input.rs: three new ServerStartupInput fields, three new env_fallback_prompt_cache_snapshot_* helpers, and build_prompt_cache_config now chains with_snapshot_limits(...) after with_apc(...). Same CLI-beats-env precedence as the sibling prompt-cache knobs.
• src/main.rs, src/bin/mlx_server.rs: --prompt-cache-snapshot-capacity-bytes, --prompt-cache-snapshot-max-entries, --prompt-cache-snapshot-ttl on both binaries, with matching help text; src/commands/serve.rs applies the env fallbacks and maps the fields through.
• src/server/prompt_cache/entry.rs: ModelSnapshotEntry::new_for_test(tokens, family, size_bytes), mirroring CacheEntry::new_for_test, so budget tests can control the accounted size without allocating MLX tensors.
• docs/environment-variables.md: the three new variables, plus how to size the budget from a measured snapshot_bytes.
• Tests: five new store tests (chain collapse, cross-session isolation, no supersede without a session key, strict-extension requirement, and byte accounting proving supersede runs before the capacity check) and five new cli_input tests (CLI override, defaults when unset, and the three env fallbacks including CLI precedence and an unparseable TTL).

Scope note on the issue's "model-aware default"

Proposed Solution (b) also asked to derive the default budget from the model's per-token state footprint. Not done here, and the reason is scope rather than architecture. There is a usable hook point: PromptCacheStore is constructed in start_server (src/server/startup.rs:2006) with startup.model_path in hand and before the model loads, and the twenty lines above it already read config.json through detect_hybrid_ssm_from_path and mutate the same config.prompt_cache struct to auto-disable APC for hybrid SSM families. A model-aware default could be computed at that same point without deferring store construction or introducing a two-phase config. What it needs is the part this issue does not carry: a per-family state-size formula (layer count, state dimensions, dtype) for each of the thirteen supports_snapshot_reuse() families, validated against measured snapshot_bytes. That is its own piece of work, and getting it wrong silently mis-sizes the store, so the default stays at 512 MiB and the knobs plus sizing guidance in --help and docs/environment-variables.md cover the operator-facing need: measure snapshot_bytes and multiply by the entries expected resident.

Corrected after review: an earlier version of this note claimed the deferral was forced by config resolution happening during CLI parsing. That is true of build_prompt_cache_config but does not establish the constraint, given the start_server hook above.

Test plan

• cargo test --release --lib prompt_cache::store::tests (28 passed, 5 new)
• cargo test --release --lib server::cli_input (98 passed, 5 new)
• cargo test --release --lib server::routes::cache (17 passed)
• cargo clippy --lib --tests -- -D warnings clean
• cargo fmt
• Real-checkpoint run on qwen3.5-0.8b-4bit (a supports_snapshot_reuse() family) against /v1/cache/stats, evidence below.

Real-checkpoint evidence

Server: mlxcel-server -m models/qwen3.5-0.8b-4bit --port 18146 --prompt-cache-snapshot-capacity-bytes 40000000.

The flag path works end to end. Startup logs snapshot_capacity_bytes=40000000 snapshot_max_entries=4096 snapshot_ttl_seconds=7200 and /v1/cache/stats reports "snapshot_capacity_bytes": 40000000, against the 512 MiB compiled default. The new snapshot_supersedes field is live on the same response.

One conversation snapshot for this model is 13246464 bytes and stays exactly that size as the conversation grows (turn vectors 72, 98, 124, 151 tokens all cost the same), confirming the sizing guidance: snapshot cost tracks model width, not prompt length, so the budget scales with concurrent conversations rather than with turns.

The reported thrash reproduces at HEAD. Under the 40 MB budget, a single four-turn conversation filled the store at three entries and began LRU-evicting from its fourth turn (snapshot_entries 1, 2, 3, 3 with snapshot_evictions_lru 0, 0, 0, 1), and a second conversation's turns then evicted the first's entries (snapshot_evictions_lru reaching 4 while snapshot_entries stayed pinned at 3).

Supersede did not fire on that path, and the cause is measured rather than assumed. Conversation A donated vectors of 72, 98, 124, and 151 tokens against following prompts of 96, 122, and 149 tokens, so the length precondition held throughout and the divergence is in the token ids at the turn boundary. A donated vector is prompt plus generated tokens and therefore ends with a generation prompt followed by generated content, which the next turn's re-rendered prompt does not reproduce at that offset, so no same-session strict extension reaches the store yet. That boundary is exactly what the sibling boundary-snapshot and warm-up-snapshot issues in epic #1148 address; this issue's own framing is that supersede is what holds their steady-state footprint at one entry per conversation. An identical-request repeat within one session was also checked as an attribution control: it advances snapshot_inserts and the LRU counter and leaves snapshot_supersedes at 0, confirming the new counter is not double-counting idempotent replacement.

Not measured: end-to-end chain collapse and two-conversation coexistence over HTTP, which are unreachable until the boundary snapshot lands. The supersede rule itself, including cross-session isolation and the two-entry coexistence case, is covered by the new store tests.

Closes #1146

[inureyes]

Security and performance review

Reviewed the full diff against the store's existing lock, byte-accounting, and resource-release patterns, plus the CLI/env parse paths. No CRITICAL or HIGH findings. One MEDIUM was fixed in ead5ac9; the rest are recorded below without changes.

Lock discipline and complexity: clean

The supersede scan holds the same write guard it was already holding and touches nothing but Inner. Inner::remove_snapshot is a plain &mut self map removal, so there is no nested acquisition of self.inner and no path back into a &self store method. The one re-acquisition in insert_snapshot (self.inner.read() for the oversized-error payload) is pre-existing and happens after drop(guard).

The scan is not a new complexity class. lookup_snapshot_prefix already walks every snapshot slot and runs slot.entry.tokens.as_slice() != &tokens[..len] on each same-session candidate, on the hotter lookup path; evict_oldest_snapshot is an O(N) min_by_key per eviction inside enforce_snapshot_caps, and sweep_snapshot_ttl collects O(N). The new filter also short-circuits well: starts_with runs only after the sessionless-bucket, session-key, and length tests have all passed, so the token comparison is reached by roughly the entries in the same conversation, not by all N.

Nit, not applied: the length test slot.entry.tokens.len() < new_tokens.len() is the cheapest predicate in the chain (one deref plus an integer compare) and currently runs third, behind two String comparisons that in a single-model server always run to full length on every entry. Moving it first would cut the common case for free.

Byte accounting: no underflow, no double-free

snapshot_bytes is mutated in exactly three places, all under the write lock: saturating_sub in remove_snapshot, saturating_add at insert, and the = 0 in clear next to snapshots.clear(). Underflow is structurally impossible, and the accounting matches remove_snapshot's existing pattern.

Double-free of one digest's bytes cannot happen for two independent reasons. Mechanically, remove_snapshot returns None on the second call for the same key, so a repeated removal subtracts nothing; the superseded vector also comes from a HashMap iteration and so carries unique digests. Logically the two paths are disjoint anyway: the scheduler builds the key from entry.tokens itself (compose_prompt_cache_key(&ctx, &key_tokens) with key_tokens = entry.tokens.clone()), so the incoming digest hashes the same token vector the supersede predicate tests, and a slot that is strictly shorter can never carry the incoming digest.

The guard.snapshots.insert(...) at the end is a raw map insert that would silently orphan an existing slot's bytes, but it is unconditionally preceded by remove_snapshot(&digest), which is what keeps that safe. Worth keeping the two lines adjacent in any future edit.

Ordering is right in the way that matters most: the oversized rejection returns before the supersede block, so a rejected insert never destroys the ancestor it would have replaced.

Resource release: no leak, nothing to stash

ModelSnapshotEntry wraps ModelStateSnapshot, which owns Vec<ModelStateTensor> where each tensor is a UniquePtr<MlxArray> produced by ffi::copy. Those are owned copies with no pool block pins, no file handles, and no borrowed model state, so Drop is the complete release and the stash_paged_pins_for_release treatment the paged detached-KV path needs does not apply here. Dropping under the write lock is what evict_oldest_snapshot and sweep_snapshot_ttl already do, so the supersede path introduces no new lock-held work of a kind the store did not already perform. A concurrent reader holding an Arc from lookup_snapshot_prefix keeps its snapshot alive past the supersede, so there is no use-after-free window.

Metrics do not drift either: record_snapshot_* are monotonic counters rather than gauges, and the scheduler's gauge update reads live store.bytes() / store.len() after the insert returns, so supersede-freed bytes are reflected without a hook.

CLI and env parsing: safe

No unwrap, expect, or panicking conversion on any user-controlled value. apply_optional_usize_env_fallback and the hand-rolled env_fallback_prompt_cache_snapshot_ttl both match on parse() and warn-log-and-ignore the error arm, leaving the value None so the compiled-in default applies. That is not a silent zero: unwrap_or(defaults.*) runs against PromptCacheConfig::default(). The TTL helper is a faithful copy of the existing env_fallback_prompt_cache_ttl, including the is_some() precedence branch and the info! note when both are supplied. Clap rejects a malformed --prompt-cache-snapshot-* value with its own usage error before any of this runs. Duration::from_secs accepts the whole u64 range, so no overflow path exists.

Two values are accepted without validation, matching the sibling knobs rather than diverging from them: --prompt-cache-snapshot-capacity-bytes 0 turns every insert into an oversized rejection (counted in snapshot_rejections_oversized, no panic), and --prompt-cache-snapshot-max-entries 0 evicts each entry immediately after admitting it. Both terminate; neither is a hang or a crash. Worth a range check across the whole family some day, not in this PR.

MEDIUM (fixed in ead5ac9): sizing guidance stated collapse as current behavior

docs/environment-variables.md said "Turns within one conversation cost nothing extra, since each turn's snapshot supersedes the previous one and snapshot_entries stays at one per conversation", and the --prompt-cache-snapshot-max-entries help text called the cap "effectively a concurrent-conversation limit". This PR's own measurement section says the opposite for the current donate path: the stored vector is prompt plus generated tokens and ends with a generation prompt the next turn does not reproduce, so no same-session strict extension reaches the store and snapshot_supersedes stays at 0.

That matters because it is the sizing advice an operator acts on. Reading it, they would budget one entry per conversation at two to three times a single snapshot, then hit exactly the thrash #1146 exists to remove. Both places now make the supersede rule conditional on a turn strictly extending the previous turn's vector, say the condition is unmet today, and point at snapshot_supersedes on /v1/cache/stats as the signal for when collapse begins.

LOW / informational, no change made

The None-session guard is unreachable in production. compose_prompt_cache_key always passes Some(ctx.session_key), and resolve_session_key falls back to ANONYMOUS_SESSION_SENTINEL when the request carries neither prompt_cache_key nor user. So the code comment's "a None session carries no conversation identity to chain on" describes a test-only branch, and in a deployment where clients set neither field, every conversation shares one session bucket. The strict-prefix predicate is what actually does the conversation identification there, and it holds: a second conversation's vector is a strict prefix of another only when the earlier conversation, generated tokens included, is a literal prefix, which in practice means it is the same conversation or an exact replay. The cost when it does misfire is the documented fork case, one re-prefill. No new trust boundary is crossed either, since the store already reuses prefixes across sessions by design and an attacker who knows a victim's session key could already force LRU eviction in that bucket. Suggest only that the comment name the sentinel so the next reader does not infer that unauthenticated traffic is excluded.

ServerStartupInput gains three public fields on a pub struct that is not #[non_exhaustive], which is technically semver-breaking for any downstream code constructing it as a literal. Same shape as how the sibling prompt-cache knobs were added, so this is consistency rather than regression.

Two copies of the same u64 env-fallback body now exist (env_fallback_prompt_cache_ttl and the new snapshot one). An apply_optional_u64_env_fallback mirroring the usize helper would remove the duplication.

No unsafe blocks in the diff. No injection surface, no secret in the warn-log path (the logged value is an operator-set numeric env var).

Verification

• cargo check --lib --tests clean
• cargo check --bin mlxcel-server clean
• cargo clippy --lib --tests -- -D warnings clean
• cargo test --release --lib prompt_cache::store 28 passed, including all five new supersede tests
• cargo test --release --lib server::cli_input 98 passed
• cargo fmt --check clean

[inureyes]

Implementation Review Summary

Intent

Session-chain supersede in insert_snapshot (before byte accounting), plus CLI/env wiring for the three previously unwired PromptCacheConfig snapshot fields (#1146, epic #1148).

Verified against the code, not the description

Supersede predicate (src/server/prompt_cache/store.rs:698-741). The filter is exactly the four conditions the issue asks for: key.session_key is Some, slot.bucket.session_key.as_deref() == Some(session_key), slot.sessionless == sessionless (model / lora / template_sig / mm_digest, per SessionlessBucketKey in types.rs:87), and slot.entry.tokens.len() < new_tokens.len() && new_tokens.starts_with(slot.entry.tokens). Equal-length and overlapping-but-diverging vectors are both excluded by construction. It also sits after the oversized-entry early return, so a rejected insert cannot destroy a chain it was never going to replace.

Byte-accounting order. Supersede runs before guard.snapshot_bytes = ...saturating_add(entry_bytes) and before enforce_snapshot_caps, and remove_snapshot subtracts size_bytes on the way out, so the freed bytes genuinely admit the successor. supersede_frees_bytes_before_the_capacity_check discriminates this correctly: under an 8000 + 9000 vs 10000 budget, an implementation that superseded after enforce_snapshot_caps would evict via LRU and trip the snapshot_evictions_lru == 0 assertion.

CLI/env parity across all four surfaces. Flag names, env names, and defaults line up between ServeArgs (src/main.rs:1770), ServerArgs (src/bin/mlx_server.rs:1155), both build_startup_input mappings, and cli_input.rs. Defaults quoted in help text match what is compiled: 512 MiB / 4096 / 7200 s, from DEFAULT_SNAPSHOT_CAPACITY_BYTES = DEFAULT_CAPACITY_BYTES / 4, DEFAULT_MAX_ENTRIES * 4, DEFAULT_TTL_SECONDS * 2. env_fallback_prompt_cache_snapshot_ttl is a line-for-line mirror of the existing env_fallback_prompt_cache_ttl, including the CLI-precedence info log and the warn-and-ignore on an unparseable value.

Counter is not dead. snapshot_supersedes runs Inner -> PromptCacheStats -> build_stats_response -> CacheStatsResponse, on both the enabled and disabled branches.

Independent end-to-end check of the env path. Started mlxcel-server from the built binary with all three MLXCEL_* variables set and no CLI flags. Startup logged snapshot_capacity_bytes=41943040 snapshot_max_entries=7 snapshot_ttl_seconds=123, and /v1/cache/stats reported the matching snapshot_capacity_bytes and snapshot_max_entries with snapshot_supersedes live on the payload. Both --help surfaces confirmed on freshly built binaries (mlxcel serve --help and mlxcel-server --help each list all three flags).

The "not validated end to end" claim is accurate. Read the five store tests rather than taking the body's word: they cover chain collapse, cross-session isolation, the None-session case, the fork case, and the byte accounting, and nothing in the PR claims HTTP-level supersede. That matches the epic, where #1143 is the load-bearing fix that makes the donated vector prefix-stable and is still open, so supersede cannot fire on the donate path yet. Deferring integration onto a sibling in the same epic is what the project convention allows.

Fixed on this branch

• 3efebe30 The --prompt-cache-snapshot-capacity-bytes help text on both binaries still told operators to budget two to three times a single conversation's snapshot. ead5ac9f had corrected that assumption for --prompt-cache-snapshot-max-entries and the docs but not here, so the capacity advice under-sized the store by roughly the turn count and pointed at the exact thrash this issue removes. Now: measure snapshot_bytes after one turn, multiply by the entries expected resident, count turns as well as concurrent sessions.
• 8af676e4 snapshot_supersedes was set on the /v1/cache/stats serialization fixture but absent from the asserted key list, leaving the wire name unpinned. Added. Also recorded in the supersede comment that "session" means the resolved key: callers setting neither prompt_cache_key nor user all land on ANONYMOUS_SESSION_SENTINEL and share one bucket here exactly as they already do in lookup_longest_prefix. Not a practical false-supersede risk, since it still needs one caller's whole stored vector to be a strict prefix of another's, but the previous wording implied per-caller isolation.
• PR body: the scope note justified deferring the model-aware default with "config resolution runs during CLI parsing, before any weights are read". True of build_prompt_cache_config, but it does not establish the constraint: the store is built in start_server (src/server/startup.rs:2006) with model_path in hand, twenty lines below a detect_hybrid_ssm_from_path call that already reads config.json and mutates the same config.prompt_cache to auto-disable APC. Rewritten to give the real reason, which is that per-family state-size formulas for thirteen families are their own piece of work.

Remaining, deliberately not changed

• Proposed Solution (b)'s model-aware default is deferred, now honestly justified. Worth a follow-up issue rather than growing this PR.
• Supersede emits no PromptCacheMetrics event, so Prometheus /metrics shows snapshot inserts climbing with no removal signal; only /v1/cache/stats carries the counter. Adding one means extending the trait and both impls, and nothing in the issue asks for it.
• The supersede scan is O(snapshot entries) per insert with string comparisons in the filter. At the 4096 default cap this is nothing against prefill cost, but it is the only full scan on that path.

Verification

• All stated acceptance criteria implemented (supersede rule, its boundaries, byte ordering, three knobs, tests beside the store)
• No placeholder or mock code; ModelSnapshotEntry::new_for_test is #[cfg(test)] and mirrors the existing CacheEntry::new_for_test
• Integrated into the production code flow: insert_snapshot is the donate path in scheduler.rs:2275 and compose_prompt_cache_key always supplies Some(session_key). Firing awaits feat: Donate a history-boundary snapshot during prefill for snapshot-only families #1143, per the epic's own dependency order
• Project conventions followed: _tests.rs siblings, no unwrap/expect added outside tests, no shared hotspot from docs/code-guidelines.md touched
• Existing modules reused: the env-fallback helpers, with_snapshot_limits, and the new_for_test pattern all follow their existing siblings
• No unintended structural changes; every edit is additive except the pub use re-export list
• Tests pass: cargo check --lib --tests, cargo clippy --lib --tests -- -D warnings clean, cargo fmt --check clean, store tests 28 passed, server::cli_input 98 passed, server::routes::cache 17 passed