[jvmenen]

Depends on #610 (segmented KV cache). This branch contains #610's commits; please review it after #610. The change itself is the last three commits.

What it changes

Segmented KV (MTPLX_SEGMENTED_KV=1) currently needs the tensor units of an M5-class GPU. On M1 to M4 the support check reports nax_unavailable and keeps the stock cache. This PR adds a second decode route so segments also work there. MTPLX_SEGMENTED_KV_DECODE=lse selects it by hand, and MTPLX_FORCE_GPU_FAMILY_FALLBACK=1 forces it on an M5 for testing. The M5 default is unchanged.

Why

Segments keep a follow-up turn from copying the whole history, which lowers peak memory, often by several GiB at long context. M1 to M4 Macs usually have less RAM, so they gain the most. The speed on those Macs is not known; users can try it and choose.

How it works

Without the tensor units there is no fused segment kernel. Instead, decode runs attention over each segment separately in float32 and also keeps each segment's logsumexp (how much softmax weight that segment carries). One exact merge then combines the pieces into the same result as attending over the whole history. The per-segment outputs stay float32 and the merged result is rounded to bf16 once, like a single attention call. A cache with a single segment still takes the stock route, prefill is unchanged, and rows that see no key at all get weight 0 instead of NaN. /health shows decode_route and lse_api.

Tested

• Unit and exactness tests for the new route, green normally and with MTPLX_FORCE_GPU_FAMILY_FALLBACK=1.
• Model runs on an M5 with MTPLX_FORCE_GPU_FAMILY_FALLBACK=1: Qwen3.8 27B, 20K and 50K context plus 4 follow-up turns, 200 greedy tokens per turn, order off, on, on, off.
• These runs used the commit before the final rebase onto Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610. The rebase only touched fully masked rows (the -inf handling).

Measured (M5 with forced fallback, NOT real M1-M4 speed)

Mean of the two runs per variant. Peak is GPU memory above the memory at the start of the turn.

Context	Turn	tok/s off	tok/s on	Peak above start off (GiB)	on (GiB)
20K	1	23.6	19.4	3.78	1.64
20K	2	29.1	20.1	4.04	1.22
20K	3	20.9	15.7	4.18	2.02
20K	4	24.8	18.9	5.56	2.16
50K	1	24.6	17.0	6.92	1.79
50K	2	21.4	14.6	5.94	2.05
50K	3	23.4	17.8	6.08	2.48
50K	4	18.4	10.6	7.49	2.76

Parity with the stock cache

Greedy output is not token-identical between off and on; each is deterministic run to run. I checked whether that is a bug or rounding, on real Qwen3.8 27B tensors (M5, forced fallback):

• Attention output (one 4-row verify window, all 16 attention layers, 3 segments): after the last commit, 12 to 23 % of the bf16 elements differ from the stock kernel's output. The exact float32 result, rounded once, differs from the stock kernel in the same 12 to 23 %, so the remaining difference is the stock kernel's own rounding. For comparison, the M5 segment kernel (Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610) differs from its stock kernel in 20 to 33 %.
• The fix (now in Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610, which this PR includes): each segment's output was rounded to bf16 before the merge, and the merged result was rounded again. That double rounding put 24 to 34 % of the elements one ulp away from stock. The per-segment outputs now stay float32. The fix lives in Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610's reference route and merge, so its prefill lse route benefits too; this PR keeps the decode-route tests. A test requires the route to equal the float32 attention rounded once (the old 2e-2 tolerance could not see this) and to be no further from stock SDPA than the exact result.
• Logits with forced identical tokens (200 reference tokens, windows of 4): the mean |logit difference| from stock is 0.035 to 0.048. The M5 segment kernel gives 0.040, and the two stock kernels (fallback vs tensor units) differ from each other by 0.041 on the same tokens. Top-1 differs in 1 to 4 of 199 rows, always where the stock top-2 margin is 0 or 0.125 (a bf16 tie or one ulp).
• Greedy runs (20K and 50K, 5 turns each): every first divergence, before and after the fix, is at a bf16 tie or one to two ulps (margins 0, 0.0625 or 0.125, checked by recomputing the stock logits at that position). That matches the M5 kernel route in Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) #610 (margins 0.0 / 0.125). This text has many exact ties (5 of 199 rows at 20K have margin 0), so no route that sums in a different order can be token-identical to stock.

The fix does not change speed or memory measurably (ms per verify within run-to-run noise, same peak).

Known risk

Per-segment float32 attention costs more than the fused kernel, and the cost grows with context and segment count. On the M5 it ran at 0.58 to 0.82 of the stock cache's tok/s. On M1 to M4 the stock route is slower too, so the ratio there is unknown.

Please test on M1-M4

If you have an M1-M4 Mac, please post decode tok/s and peak memory with MTPLX_SEGMENTED_KV=1 vs off, for example:

# off
mtplx serve --model <model>
# on
MTPLX_SEGMENTED_KV=1 mtplx serve --model <model>

Send a long prompt (20K+ tokens) and a few follow-up turns in the same conversation, then post the decode tok/s from the dashboard (or the response timings) and the peak memory per turn, plus the decode_route value from /health.