One snapshot per conversation: fixing prompt caching for a hybrid-attention model

Carmen, my assistant, and my coding sessions both talk to GLM-5.3-Flash on a Mac Studio M3 Ultra, through mlx-vlm's server. Every turn re-sends the whole conversation, so the server's prompt cache decides whether a turn starts in a second or in a minute. Last week I made decode about 5% faster (previous post). This week I looked at why the cache was working less well than it should, and the answer turned into a small redesign of how it stores things.

What the log said

Three things stood out in a morning of traffic.

Every request stored two full snapshots of the cache: one at the prompt minus 16 tokens, one at the full prompt. The second is what the next turn hits (its cached_tokens equals the previous prompt_tokens, every time). The first exists so that an identical resend of the same prompt can still be served: the server needs at least one uncached token to produce logits from, so a snapshot at exactly the prompt length is useless for a resend. In a normal turn flow that guard copy is never touched, and it costs a full clone, a 1 to 2.6 GB disk write and a cache slot every turn.

The RAM tier held 8 entries. Two per turn means 8 entries is four turns of one conversation. With a second conversation running (my coding tool spawns parallel sub-agents), the log showed 155 evictions in about an hour, and the disk tier held 939 GB of superseded chains that a daily LRU pass cycled through.

New sessions always started cold. The system prompt plus tool definitions is 9 to 12 thousand tokens for Carmen's agents and about 10 thousand for the coding tool; at ~420 tokens per second of prefill that is 22.8 seconds before the first token, paid by every fresh session even though every session shares that prefix with every other.

Why a hybrid model can't use the usual cache

For a plain transformer the prompt cache is a per-token KV cache, and a prefix cache is a trie of KV blocks: any prefix of any cached prompt is reusable by slicing. GLM-5.3-Flash has 11 layers of sparse attention with such a cache, but 34 layers of Kimi Delta Attention, whose "cache" is a recurrent state: a [64 heads, 128, 128] matrix per layer plus a four-token convolution window. You cannot slice a recurrent state back to an earlier token. The only reusable unit is a snapshot of the whole state at a token boundary, which mlx-vlm calls exact mode: whole-prompt snapshots, matched by exact token prefix.

Two numbers make the redesign obvious once you write them down.

• The KV part costs about 27 KB per token across the 11 attention layers: 2.6 GB at 95k tokens, but only 12 MB at 426 tokens.
• The recurrent part costs about 143 MB, full stop. It is the same size after 400 tokens and after 95,000.

So a checkpoint of the recurrent state is cheap and prompt-length independent, and the KV cache is append-only, so rows [:n] of a longer prompt's KV are exactly the KV at n tokens. A snapshot at an earlier point in a conversation does not need its own KV copy at all: it needs 143 MB of recurrent state plus a length, and can read the KV rows out of the conversation's latest snapshot. MLX makes holding those references safe: arrays are immutable once evaluated, and the in-place slice update the KV cache performs is copy-on-write when anything else still references the buffer.

The design

One entry per conversation lineage. The entry is the full snapshot at the latest prompt; it carries a list of checkpoints, each a set of recurrent states at an earlier length that share the entry's KV buffers.

Checkpoints come from four events, none periodic:

• guard: the prompt minus 16 tokens, captured during prefill (replaces the second full snapshot; 143 MB instead of a clone).
• final: when a new prompt extends a stored entry, the old entry is absorbed. Its own final state becomes a checkpoint at the turn boundary, its checkpoints are carried over, its RAM slot is freed and its disk file deleted once the new one is written.
• anchor: the boundary between the system prompt plus tools and the first user message, stored as a separate pinned entry, exempt from the conversation LRU and from absorption, once per distinct system prompt.
• divergence: where a prompt stops sharing tokens with its closest existing sibling, if no restore point covers that length and the gain is at least 256 tokens.

Lookup computes the longest common prefix between the prompt and each entry with numpy, then takes the entry's full length if the whole entry is a prefix, else its longest checkpoint at or before the divergence. The longest restore point across RAM wins; the disk index is consulted only for something longer. An edited last message, a regenerate, a retry after a dropped stream, a new session with the same system prompt: each now lands on a checkpoint instead of falling back to whatever older whole-prompt entry survived, or to a cold prefill.

The anchor needs the server's help. Rendering the chat template on the leading system messages plus the tool list, without the generation prompt, gives the exact text prefix that every session shares; the generator tokenizes it once per distinct text and passes the common-prefix length, minus two tokens of slack for BPE merges at the boundary, down to the batch generator as a capture point. GLM's template puts the tools block before the system text, which is why the slack matters: the last tokens of the system text are where per-request differences tend to live.

Thinning keeps at most four checkpoints per lineage, dropping old guard points first (only the newest is a live resend point), then the oldest turn boundaries, never anchors. Anchors have their own caps, 16 entries and 16 GB, inside the same 48 GB budget as the conversations; on disk, pinned anchor files are the last candidates for the 1 TB cap's eviction rather than exempt, because a system prompt with a timestamp inside it would otherwise pin a new file every minute forever. Everything is behind one switch, APC_LINEAGE=0, that restores the previous behaviour.

What it took

Five files in mlx-vlm's server and cache layers, a diff of about 2,100 lines, generated by a script of exact-match edits against pristine copies so the patch can be regenerated and reviewed. The unit tests run on the CPU with kilobyte-sized fake caches shaped like the model's layout, through an overlay package, so they exercise the staged code without installing it. Nothing touched the production process until the tests passed. Installing was four restarts of about 20 seconds each, with an idle check before every kill.

The tests caught two of my own wrong assumptions before anything ran against the model: a stored entry holds only its used KV rows, so the byte budget in a test was never exceeded; and an edited turn may legitimately snap to a checkpoint later than the turn boundary when the edit shares a few words with the original. The live checks caught a third: with a 426-token prompt the anchor point fell after the guard point, so the anchor was not pinned; my real traffic turned out to be exactly that shape, ten-thousand-token system blocks followed by seven-token questions, so the rule became "pin the anchor wherever it falls".

What it bought

On my own ~9.5k-token prompts, the same morning, before and after:

The first five real requests after the install were Carmen's scheduled agents: all served at nearly their full length from the previous day's disk files, and four anchors pinned in the process (0.37 to 0.47 GB each). Two of the coding tool's system prompts differ by four tokens; the second variant now reuses about 9.5 thousand of its 9.54 thousand tokens through a divergence checkpoint instead of paying the full 22.8 seconds.

Over the rest of the day the log kept telling the same story: conversations kept hitting their own lineage, new sessions landed on an anchor or on a sibling's checkpoint instead of starting cold, evictions all but stopped, and the disk tier shrank as superseded snapshots were removed. I'll let a week of traffic settle the rates before quoting them.

The cache design in Rapid-MLX, which I looked at before writing this, solves the same problem with a radix trie for KV and periodic recurrent checkpoints every 2,048 tokens, thinned to four per layer. Periodic placement handles arbitrary mid-prompt divergence more generally. Event-driven placement is cheaper for the agent pattern, where prompts diverge at message boundaries and at a timestamp inside the system prompt, and those points are known.

Where the other second went

The redesign landed, and the log still showed warm turns taking 2.6 to 4.6 seconds to the first token on 16–22k-token prompts, and 1.9 seconds on a 12.5k prompt with only 16 new tokens. The prefill itself was a fraction of that. The server's "prefill elapsed" figure only times the forward passes; the work around them was untimed, and that is where the time was.

Each warm turn was copying the whole KV cache four times: the restore clone out of the cache entry, the copy into the batch layout, the row extraction at commit, and a second clone of that extraction. At 21k tokens that is four copies of about 0.6 GB. Bytes alone would take tens of milliseconds. It took seconds because the batch generator returned MLX's buffer pool to the operating system after every prompt chunk and every prompt batch, so each copy allocated fresh memory and page-faulted it in at well under 1 GB/s. The old build's log broke it down: 1.0 to 2.1 seconds of admission plus 0.5 to 0.6 seconds of commit per warm turn at 16–22k tokens, and 1.4 to 2.6 seconds at 100k.

Three small changes removed the copies. The batched path now takes the stored row read-only — the KV entries become proxies truncated to the restore length — because the batch merge copies it anyway. The commit stores the row extraction without a second clone. And the per-chunk and per-batch pool clears are now opt-in, with the launcher's existing 16 GB cap and periodic release still bounding the pool; resident memory did not move (169.7 GB before and after). The safety argument is the same one from earlier in this post: MLX arrays are immutable once evaluated, and the cache's in-place slice update is copy-on-write when another reference exists, so sharing arrays is safe. Two timing lines were added to the log so the cost is visible per request.

Benchmark: one conversation grown to ~19.8k tokens, then short questions of 22 new tokens; the before figures come from the log of real traffic the same morning:

About one second per short turn remains, and it now sits inside the timed forwards: two small prefill chunks (the guard checkpoint splits the last 16 tokens off) plus building the generation batch. Copies were the lever this time; that remainder needs its own timing pass.

The last second

This is that timing pass. After the copies were gone, a short warm turn on a ~20k-token conversation still took 1.1 seconds to the first token, and a 13-token cold probe took 1.3. Two forwards of a few tokens each cost ~0.45 seconds apiece, regardless of token count — while a single-token decode step costs 0.033 seconds. Something in a short forward was eating fifteen decode steps' worth of time.

The first fix was cheap and structural. The guard checkpoint at prompt-minus-16 split every short suffix into two of those forwards. It is now captured only when the uncached suffix is at least 256 tokens; a verbatim resend of a shorter turn restores at the previous turn boundary, which absorption keeps, and re-prefills that short suffix. 1.1 seconds became 0.7.

Then a per-layer timer, gated by a flag file so it costs nothing in production. A 44-token forward at 19.9k tokens took 0.65 seconds: 0.43 of it was the 11 sparse-attention layers — 39 ms each — 0.17 the MoE blocks, and only 0.06 the 34 linear-attention layers whose recurrent kernels I had suspected. The cause: for chunks longer than one token but below the block-gather threshold (24.6k tokens), the sparse-attention layer took the unabsorbed MLA path, which expands the latent KV cache of the entire context into per-head keys and values on every forward — a 1.6 GB write at 20k tokens, 29 of those 39 milliseconds. The decode path (one token) already gathers each query's top-2048 latent keys and attends over them directly, at 1.6 ms per layer. Nobody had asked what happens between one token and two thousand.

The fix fills that gap: the decode shape batched over the chunk's queries, for chunks up to 128 tokens, and block-gather — already in production for long contexts — now also handles 128 to 512 tokens at any context length; the unabsorbed path keeps the long chunks. Measured on one real layer with random weights, no model load, milliseconds per forward at 19.9k tokens:

Equivalence: the relative L2 difference versus the previous path is about 1e-6 for every chunk size — the same as block-gather — and not bit-identical by construction, since it is a different association of the MLA absorb. A near-tie token can flip either way ("Four" versus "Four."), which the same prompt already does between two cache restore points.

On the live server, short warm turns at ~19.8k tokens, time to first token:

After the change, the sparse-attention share of a 21-token forward is 0.05 seconds; what remains is the MoE (0.09) and the linear-attention kernels (0.04), both proportional to real work.

Measuring one layer in isolation with random weights is the cheap way to find this kind of thing. A synthetic 100k-token context on a machine that is also serving is not small, though; I learned that twice this month.

What's left

A short warm turn now spends its time in real work — the MoE and the recurrent kernels for the new tokens — so the next gains are elsewhere: fewer tokens per turn, and the batch-size-1 efficiency of the expert matmuls. The single-sequence code path, which my server doesn't use, still writes its guard point as a full snapshot; it gets absorbed correctly, it is just not cheaper. And the exact-mode cache this extends belongs to mlx-vlm; the change is self-contained, off-switchable and tested, so it may become a pull request once a week of production traffic has confirmed the numbers above.