[jvmenen]

Why

When a conversation has been evicted from memory to the SSD cache and comes back, MTPLX often does not restore it. It reuses only the system prompt it shares with another conversation still in memory and computes the rest again: about 145 s for a 50K conversation instead of a few seconds. Sometimes the request fails with a 507 instead, because a conversation read from SSD and then refused stays in memory during that long prefill.

How it works, in plain terms

Before the normal restore, a near-prefix lane looks for the best partial match. It asked the SSD tier, read and decoded the conversation's own full state, then refused it, because a full prefix belongs to the exact restore. The next candidate, a neighbour's system prompt in RAM, was accepted and the lane stopped, so the exact restore never ran. This PR changes three things:

• The SSD tier reports from its manifest, without reading tensors, how long a full prefix it holds; the near lane only accepts a match that beats it, so the exact restore takes over. A hybrid row whose recurrent checkpoint lies below the lane's floor is refused from payload.json, also before reading tensors.
• The near lane lets go of every candidate it tried or skipped before it starts the prefill, and of an SSD candidate once it has been cloned into the new cache, so neither stays in memory next to the prefill.
• The Gemma 4 near lane gets the same release.

What changes

Five commits:

1. An exact prefix on SSD wins over a neighbour's system prompt in RAM.
2. Refused and untried near candidates are released before the served prefill.
3. A served SSD candidate is released after the clone, before its suffix prefill.
4. The same release in the Gemma 4 near lane.
5. A hybrid SSD row below the caller's floor is refused before decoding (new miss reason ssd_prefix_boundary_below_floor, with a line in CacheExplanation.swift).

Results

Qwen3.8-27B, a 50K and an 80K conversation, eviction to SSD forced by the bank cap, two runs. The turn after eviction restored the 50K conversation from SSD in about 4 s in both runs, where it took 143 to 146 s (or failed) before. Each SSD restore now decodes the state once instead of twice. The release commits are checked by unit tests: refused, untried and served SSD candidates are gone during the prefill.

Tests

tests/test_near_lane_spilled_exact_prefix.py, tests/test_near_lane_releases_refused.py (including the served SSD candidate, checked with weak references), tests/test_gemma4_near_lane_releases.py and tests/test_cold_tier_restore_point_floor.py; each fails without its commit. Full suite passes.

🤖 Generated with Claude Code

[Vardaan97]

A data point for #608 and #609 together, on Flash-Next. Short version: no errors, the same cache hits as 2.12.2, and fully repeatable outputs.

Build

• main (9882703, 2.12.2) with one_copy: fall back to copying when a QSA resize would cut held bank rows #619, SSD restore builds the request's cache layout #608, Restore an evicted session from SSD, and release near-prefix candidates the lane does not keep #609, Prefill rungs carry the GDN states: lower memory peak on hybrid models (MTPLX_PREFILL_ASYNC_RUNGS_STATES, opt-in) #607 and Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610 merged in that order.
• Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610 stayed off (MTPLX_SEGMENTED_KV unset).
• Prefill rungs carry the GDN states: lower memory peak on hybrid models (MTPLX_PREFILL_ASYNC_RUNGS_STATES, opt-in) #607's flags were set in one of the three PR runs, but it doesn't install on Flash-Next (qwen4_exp): it hooks Qwen3_5TextModel only. All three PR runs therefore ran the same code.
• Write an idle conversation to SSD before the admission releases it #618 and Keep an unwritten idle conversation through a WARNING pressure trim #620 weren't included: neither merged cleanly on top of Restore an evicted session from SSD, and release near-prefix candidates the lane does not keep #609 in that order (conflict in mtplx/cache_bank/cold_tier.py).
• Same mlx 0.32.2 and mlx-lm 0.31.3 as the stock 2.12.2 it's compared with.

Setup

• M5 Max with 128 GB, macOS 27.0.1, Flash-Next Optimized-Speed.
• --profile turbo --depth 3, context 262144, --max-active-requests 1, MTPLX_ONE_COPY=0, bank 8 GB, MTPLX_MEMORY_LIMIT_BYTES=100G.
• Workload: 39 requests replayed at temperature 0 from a recorded coding-agent session (dsh, plain /v1/chat/completions, anon sessions). Prompts run from 3,968 to 165,052 tokens.
• Each run started from a fresh server, one server at a time. Run order: stock, PR, PR, PR, stock.

	stock 2.12.2 (2 runs)	PR build (3 runs)
requests OK	39/39, 39/39	39/39 ×3
prompt tokens from cache	95.9%	95.9%
restores per run: ssd_block_prefix_boundary_clone	12, 12	12, 12, 12
restores per run: ssd_clone	0, 0	1, 1, 1
median decode, requests ≥300 tokens (tok/s)	68.3, 62.4	73.1, 70.7, 70.2
peak GPU memory in use, system-wide (ioreg, 2 s samples), GB	102.8, 99.4	99.1, 100.9, 100.5

The restore counts are taken from request-log-8080.jsonl over the replayed requests only (38 or 39 of the 39 matched per run), leaving out the launcher's readiness probes. None of the matched requests restored cold in either build.

Repeatability (byte-identical reasoning, content and tool calls)

• PR build vs itself: 39/39 in all three runs.
• Stock vs itself: 38/39.
• PR build vs stock: 18/39 identical, with 29/39 matching tool calls.

So at T=0 the build changes outputs, but it's deterministic for a given request history.

About speed: decode followed run order more than build (the two stock runs were first and last), so I read it as no regression rather than a gain.

I didn't isolate #608 from #609, and with one-copy off this doesn't exercise #619's path. I can run a narrower A/B if that's useful.

[jvmenen]

Thanks for running this. On #607 not installing on Flash-Next: that's by design. qwen4_exp already does the same natively since 262592b: during wide prefill it flushes the hidden stream plus the finished group's recurrent states every four layers (MTPLX_QWEN4_PREFILL_MIDLOOP_EVAL, on by default). #607 brings that to the qwen3_5 family, which has no such loop of its own. I'll add a line to the description.