Repository navigation

Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV) - #610

Open

jvmenen wants to merge 38 commits into

youssofal:mainfrom

jvmenen:feat/segmented-kv

Open

Segmented KV cache: follow-up turns, forks and restores without copying the history (MTPLX_SEGMENTED_KV)#610
jvmenen wants to merge 38 commits into
youssofal:mainfrom
jvmenen:feat/segmented-kv

Commits on Oct 9, 2026
- SSD restore builds the request's cache layout (cache_factory) instead of runtime.make_cache()
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  0384434
- Segmented KV cache behind MTPLX_SEGMENTED_KV: segments, fused verify kernel, bank and admission
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  9235104
- Segmented KV only above the compiled-verify boundary (MTPLX_SEGMENTED_KV_MIN_TOKENS)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  353f696
- Segmented prefill: reshape the kernel's (B, Hq, Q, 1) logsumexp for the merge; real-kernel test
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  b0711b1
- Segmented verify windows of 6 to 32 rows: sub-windows over the segments instead of a gather; prefill route only above 32 rows; route counters per window size
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  9cc33b5
- Segmented decode route also serves single-row steps; seal and merge counters
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  b05703f
- Tests: filter the lse route counters
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  554ecc2
- Segmented KV: a segment made from stock rows owns exact row-contiguous storage (MLX copied the strided view on every kernel launch)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  8bb93dc
- SSD tier for segmented snapshots: the codec writes the stock block format of the same rows (no history gather), guards behind MTPLX_SEGMENTED_KV_SSD
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  738eb8a
- Segmented KV: spill hashes each sealed segment block once (digest memo on the segment), /health block, SSD counters
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  2e6edbe
- Segmented KV: the tiered merge also merges segments that bank snapshots hold (nothing merged before: every turn's snapshot holds its segments)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  d80590b
- Changelog: spill memo, merge of held segments, /health block
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  cac2908
- Segmented KV test: an SSD restore builds the request's cache layout with segments too (the fix itself is the base commit)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  2966d66
- Changelog: SSD restore layout
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  caaed25
- Segmented KV: models the segment kernel cannot serve (head_dim, GQA, attention not hooked, no flash route) keep the stock cache, log one line and report the reason on /health
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  b9e424f
- Segmented verify kernel at head_dim 128 (D is a template parameter), defensive lse fallback, tests (kernel tests pending a power-connected run)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  377aa86
- Segmented KV support check: read head_dim from the attention scale (mlx-lm Qwen3/Llama attention has no head_dim attribute)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  b7144b2
- Segmented KV: plain q/k-norm attention (Qwen3, head_dim 128) takes the hook body on a segmented cache (it would gather the history on every call otherwise); support check requires the layers to be …
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  07571bd
- Segmented KV: one path with the switch on (segments from the first token, no context fence, no MIN_TOKENS); single-segment calls take the stock route on a zero-copy view instead of counting as gath…
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  6da2451
- Segmented KV test: force the LSE prefill route so the no-gather test also holds on stock MLX
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  01a7db5
- Segmented KV: one log line when the MLX build has no return_lse (prefill over segments then gathers the history)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  eb59b70
- Segment kernel: the stock sdpa_nax_flash_dsplit.py is upstream again (no NOMASK parameter); the segment kernel switches the mask off in its own derived source
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  3c36a06
- Segmented KV: sub-windows respect the kernel's 10-row limit too (GQA 1 and 2 gathered windows of 11 to 32 rows)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  7095bd9
- Segmented KV: no verdict on the model means off (fail closed); every runtime load starts without one
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  4fc08ca
- Segmented KV review cleanup: no reference counts, unused helpers gone, one helper each for unique bytes and env switches, flag-off paths skip the segment walks
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  059b095
- Changelog: segmented KV entry as a heading with sub-bullets, current behaviour only (no 32K fence history, windows 1 to 32 rows)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  36f79fa
- Segmented KV: Llama-style attention (Llama, Mistral) through the split-attention hook
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  6fdd948
- Changelog: segmented KV serves Llama-style attention (Mistral 7B v0.3 numbers)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  8e15baa
- Segmented KV under any profile: the attention switches are kill switches only, the self-check probes the segment kernel
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  7155177
- Changelog: segmented KV under any profile (attention switches only as kill switches, self-check on the segment kernel, Qwen3.6-35B-A3B sustained numbers)
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  0ae05ae
Commits on Oct 10, 2026
- Segmented KV: no dense-decode ceiling, a long conversation stays on segments instead of repaging
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  7be91ff
- Changelog: segmented KV stays on segments past the dense-decode ceiling
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  31464db
Commits on Oct 13, 2026
- Segmented KV: use mx.fast.scaled_dot_product_attention_with_lse when present, keep return_lse as fallback
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  4348005
Commits on Oct 14, 2026
- Segmented KV lse route: keep per-segment outputs in float32, round once after the merge
  Jeroen van Menen
  committed
  257af12
- Segmented KV lse route: float32 per-segment output from the MLX kernel when the build offers it
  jvmenen
  committed
  0bc6803
- Segmented KV lse route: drop the float32_output probe, rely on scaled_dot_product_attention_with_lse
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  ab430ef
- Segmented KV test: the refused-shape fallback also refuses scaled_dot_product_attention_with_lse
  
  Jeroen van Menen
  authored and
  Jeroen van Menen
  committed
  261ef7e
Commits on Oct 20, 2026
- Segmented KV spill: segment digests first, kept through merges, memo for the rest
  jvmenen
  committed
  8270858

Commits

Commits on Oct 9, 2026

Commits on Oct 10, 2026

Commits on Oct 13, 2026

Commits on Oct 14, 2026

Commits on Oct 20, 2026