[jvmenen]

Why

A follow-up turn copies the whole conversation. The session bank keeps a snapshot of each layer's KV buffer, and that snapshot shares memory with the live buffer, so the first write of the next turn triggers a full copy: about 5 GB at 80K tokens on a 27B model, on top of what the turn itself needs. That copy is what pushes long conversations into 507 refusals and evictions sooner. A clone or lease restore, a fork, or a return to an earlier turn pays the same copy again.

How it works, in plain terms

Instead of one growing buffer per layer, a conversation is kept as a list of segments: one immutable segment per finished turn, plus a tail that grows during the current turn. A snapshot, a fork or a return to an earlier turn is then a list of references to segments that already exist, so nothing is copied. Attention reads all segments in one kernel launch and merges the partial results with their log-sum-exp, which gives the same result as attention over one buffer, within rounding. Small recent segments are merged now and then, so the list stays short (3 to 5 segments over 12 turns at 50K). On SSD a segmented conversation is written in exactly the existing block format, so dedupe, crash safety and old entries keep working, with the switch on or off.

Past the dense-decode ceiling. Without segments, a request's dense KV cache is one contiguous slab, so the auto layout caps it at a share of RAM (MTPLX_DENSE_DECODE_RAM_PERCENT, default 15 %) and repages the cache to the paged layout past that ceiling. With segments there is no single slab: each turn's segment is allocated on its own, so the reason for the ceiling is gone. Repaging a segmented cache is worse than useless: the paged cache is built from the concatenated segments, a full copy of the history at the moment memory is tightest, and the admission rightly prices that copy, so a follow-up turn just past the ceiling was refused with a 507 although it added only a few rows. With the switch on and an unquantized cache the auto layout now stays dense; the stock cache, a quantized cache and an explicit MTPLX_SUSTAINED_PREFILL_LAYOUT keep their repage.

What changes

Everything is behind MTPLX_SEGMENTED_KV=1; with the switch off nothing changes.

• mtplx/segmented_kv.py: the segment cache, sealing per turn, tiered merging, snapshots and rollback without copies, /health counters.
• mtplx/kernels/sdpa_segmented.py: a verify/decode kernel that reads up to 12 segments in one launch (head_dim 256 and 128), derived from the stock dsplit kernel without changing it; windows above 10 rows run as sub-windows.
• mtplx/attention_split.py: the attention hook routes a segmented cache to that kernel for gated, q/k-norm and Llama-style attention; one segment takes the stock route on a zero-copy view.
• mtplx/session_bank.py, cache_bank/codec.py, cold_tier.py: snapshots hold segments; the SSD writer reads segments as one virtual buffer and hashes each sealed block once.
• mtplx/kernel_selfcheck.py: when the switch is on, the startup self-check tests the segment kernel at the model's shape; a failure means the stock cache, with the reason on /health.
• Support check: a model the kernel cannot serve (head_dim, GQA, attention not hooked, no NAX) keeps the stock cache and says why on /health; a model loaded without a check is treated as off. MTPLX_GQA_PACKED_SDPA=0 or MTPLX_NAX_FLASH_ROUTE=0 act as kill switches; the feature works under every profile.
• Prefill over several segments uses the logsumexp output of MLX's fused SDPA when the MLX build has it (mx.fast.scaled_dot_product_attention_with_lse, or the return_lse=True keyword of a patched build; neither is in released MLX yet); otherwise it gathers the history for that chunk (about 5 % slower prefill on follow-up turns, decode and verify unaffected) and logs one line. Rows without a visible key (lse -inf) merge to zeros instead of NaN. The reference route (and so MTPLX_SEGMENTED_KV_PREFILL=lse without the MLX kernel) keeps each segment's output in float32 and the merge rounds once to the query dtype, instead of rounding every segment to bf16 first; the fused kernel still returns the query dtype per segment.
• mtplx/generation.py (_sustained_prefill_layout): with segments on and an unquantized cache, the auto layout stays contiguous_dense_decode past the dense-decode ceiling.

Results

Qwen3.8-27B, four follow-up turns, median of 3, segments (eager verify) against the stock cache (compiled verify):

• verify 3 to 12 % faster from 8K to 32K, 0 to 4 % at 50K;
• peak memory above the turn start 1 to 4.5 GiB lower at every context length;
• server test with long conversations and a bank capped at 7 GB, two runs each: no 507s and no memory sheds with segments, against two 507s and eleven guard interventions without; peak about 3 GB lower.

Qwen3.6-35B-A3B under the sustained profile (only the 10 attention layers are segmented): 24 to 30 % less time per verify, peak per follow-up turn 0.5 to 0.8 GiB instead of 2.7 to 3.3, prefill unchanged. Mistral 7B v0.3: peak per follow-up turn 0.7 to 1.7 GiB instead of 3.2 to 7.0, speed equal. Greedy output is identical to the stock cache except where two tokens tie within one bf16 step.

Tests

Nine test files (tests/test_segmented_kv*.py, tests/test_sdpa_segmented*.py): the cache against a plain buffer through turns, forks and rollback; the kernels against stock SDPA for 1 to 13 segments, head_dim 128 and 256, GQA 1 to 8; bank, SSD round trip in both directions, merging, support verdicts, self-check failure, and that no call gathers the history. Full suite passes.

Notes

Stacked on #608 (the SSD restore builds the request's cache layout); this PR shrinks by that commit once #608 is merged. Apple M5 only: on M1 to M4 the kernel is unavailable and the feature stays off. Gemma 4 is not covered yet (it stays off).

With an MLX build that has mx.fast.scaled_dot_product_attention_with_lse, the lse route takes the per-segment output in float32 from that function and the merge rounds once. The last three commits add the docstring for this and two tests: the route rounds once only with the function, and the real function is closer to the exact attention than rounding every segment first. They also make two existing tests independent of the build, because they mocked only the return_lse keyword. Taken together they change no code.

[jvmenen]

A note on how this PR meets #593 and #594. The segment kernel is derived at import from the source of the head-dim-split kernel. #593 (wide verify windows) and #594 (quarter head-dim split) turn the row blocks per threadgroup (TG_M) and the head-dim parts (NDH) into template parameters of that source. On a tree with all three, the segment launch passes neither, so a verify over two or more segments does not compile, and the segment tests here fail. This PR also runs every GQA-6 window past 5 rows as sub-windows of 5 rows, each a full walk over the keys, while #593 serves such a window in one launch on a contiguous cache.

#639 holds that coupling. It is stacked on this PR, #593 and #594, so none of the three has to know about the others. Once those are merged, its diff is only the segment-kernel launch (_route_shape, the key-block table, the TG_M/NDH template arguments) and its tests. With MTPLX_NAX_FLASH_WIDE and MTPLX_NAX_FLASH_DSPLIT4 off, the segment route is bit-identical to the one in this PR.