Segmented KV on Macs without tensor units (M1-M4): per-segment attention with a logsumexp merge

CHANGELOG.md

@@ -4,6 +4,23 @@ All notable user-facing changes to MTPLX. The format is based on
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow
 [Semantic Versioning](https://semver.org/).
 
+## [Unreleased]
+
+### Added
+
+- **Segmented KV cache, behind `MTPLX_SEGMENTED_KV=1` (default off).** The full-attention layers keep a conversation as sealed, exact-size KV segments plus one growing tail instead of one buffer that the session bank's snapshot aliases. A follow-up turn writes only into a new tail, so MLX's copy-on-write no longer copies the history: on Qwen3.8 27B a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K at the first write of the next turn, with segments +0.125 GiB.
+  - Verify windows of 1 to 32 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`); a window the kernel does not take whole runs as sub-windows over the segments. A cold prompt is one segment and runs the stock route on a zero-copy view of its rows. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers the history otherwise (one log line says so).
+  - With the switch on the full-attention layers are segment caches from the first token, for models the segment kernel serves (head_dim 128 or 256, the split-attention hook installed, a GPU with NAX TensorOps), under any profile: segments run their own kernel and do not need `MTPLX_GQA_PACKED_SDPA` or `MTPLX_NAX_FLASH_ROUTE`, which only switch them off when set to 0. The kernel self-check runs whenever the switch is on and tests the segment kernel at the model's head_dim before the first cache is built; a failure keeps the stock cache. On Qwen3.6-35B-A3B under the `sustained` profile (50K plus 4 turns, neither switch set) the peak above the turn start drops from 2.7 to 3.3 GiB to 0.45 to 0.76 GiB and ms per verify from 34 to 41 to 25 to 30, with no history gather. The hook serves the gated Qwen3.5/3.6/3.8 attention, Qwen3's q/k-norm attention and the Llama-style attention without norms (Llama, Mistral); on Mistral 7B v0.3 (AR path, 22K plus 4 turns) the peak above the turn start drops from 3.2 to 7.0 GiB to 0.7 to 1.3 GiB at the same ms per step, with greedy output identical. Any other model, and a model loaded without that check, keeps the stock cache; `/health` says why.
+  - The session bank stores the list of segments as the snapshot and counts a shared segment once; the admission gate prices no history copy for it. Small recent segments merge in geometric tiers, also when bank snapshots still hold them, which keeps a conversation at 3 to 5 segments over 12 follow-up turns at 50K.
+  - A conversation stays on segments past the dense-decode ceiling (157,286 tokens on a 64 GB Mac with the auto ceiling). The repage to the paged cache built its pages from a gather of every segment: on Qwen3.8 27B a follow-up turn of 161,835 tokens with 84 new ones was priced 11.6 GiB for that copy and refused with 10.1 GiB free; with segments it is priced for the new rows only, about 1.1 GiB. Quantized KV and an explicit `MTPLX_SUSTAINED_PREFILL_LAYOUT` keep the repage.
+  - Segmented entries persist to the SSD tier in the existing block format, without gathering the history (`MTPLX_SEGMENTED_KV_SSD=0` skips them): shared prefixes dedupe, entries restore with the switch off and stock entries restore with it on. A follow-up spill hashes only the new segment.
+  - Measured on Qwen3.8 27B against the stock cache with compiled verify: 3 to 12 % fewer ms per verify at 8K to 32K, 0 to 4 % at 50K, and 1 to 4.5 GiB less peak above the turn start at every context. Not covered: the compiled verify (segmented caches verify eager; `MTPLX_VERIFY_ASYNC_CHUNK_LAYERS=8` still helps), paged and quantized KV. `/health` gets a `segmented_kv` block (segments per entry, merges, route and SSD counters), absent when the switch is off.
+  - **Macs without the tensor units (M1 to M4) can use segments too.** The segment kernel needs the Metal 4 tensor units of an M5-class GPU, and on M1 to M4 the support check kept the stock cache (`nax_unavailable`). Without them (or with `MTPLX_SEGMENTED_KV_DECODE=lse`, or `MTPLX_FORCE_GPU_FAMILY_FALLBACK=1` to try it on an M5) decode attends per segment in float32, each with its logsumexp, and merges the pieces exactly; prefill and the single-segment case are unchanged, and the M5 default stays the kernel. The memory gain is the same; the speed is not measured on those Macs yet. `/health` shows `decode_route` and `lse_api`. On an M5 with the fallback forced (Qwen3.8 27B, 20K and 50K plus 4 turns), the peak above the turn start drops 2.1 to 5.1 GiB per follow-up turn and decode runs at 0.58 to 0.82 of the stock cache's tok/s there.
+
+### Fixed
+
+- **A conversation restored from the SSD tier gets the cache layout the request asked for**, the same one a restore from memory gets. The SSD restore built the runtime's default layout, which for a long context is the paged KV cache (`VllmMetalPagedKVCache`) instead of the dense cache: about 5x slower prefill and slower decode after every restore from disk. Measured on Qwen3.8 27B (with the segmented KV cache of a later change, which had the same omission), the first token after a restore from SSD at 50K and 80K came after 5 to 8 s instead of 16 to 48 s and decode went from 14 to 18 to 25 to 30 tok/s; a run without segments showed the same 14 to 17 tok/s after its SSD restore.
+
 ## [2.12.2] - 2026-10-03
 
 ### Fixed

mtplx/attention_split.py

@@ -19,6 +19,16 @@
     note_unowned_rotary_origin,
     rope_offset_of,
 )
+from .segmented_kv import (
+    VERIFY_WINDOW_MAX,
+    SegmentedKVCache,
+    attend_segments_lse,
+    decode_segments_attention,
+    gathered_view,
+    segmented_kv_requested,
+)
+from .segmented_kv import count_route as segmented_count
+from .segmented_kv import prefill_route as segmented_prefill_route
 
 
 def _env_enabled(name: str, *, default: bool = False) -> bool:
@@ -336,6 +346,39 @@ def _attention_has_gated_q_proj(attn: Any) -> bool:
     return int(weight.shape[0]) == expected
 
 
+def _attention_has_plain_q_proj(attn: Any) -> bool:
+    """Plain q/k-norm attention (Qwen3): q_proj has one output per head dimension, no gate."""
+    q_proj = getattr(attn, "q_proj", None)
+    q_norm = getattr(attn, "q_norm", None)
+    weight = getattr(q_proj, "weight", None)
+    norm_weight = getattr(q_norm, "weight", None)
+    if weight is None or norm_weight is None:
+        return False
+    num_heads = int(getattr(attn, "num_attention_heads", getattr(attn, "n_heads", 0)))
+    return int(weight.shape[0]) == num_heads * int(norm_weight.shape[0])
+
+
+def _attention_has_unnormed_q_proj(attn: Any) -> bool:
+    """Llama-style attention (Llama, Mistral): q/k/v projections and rope, no q/k norm, no gate."""
+    if getattr(attn, "q_norm", None) is not None or getattr(attn, "k_norm", None) is not None:
+        return False
+    if any(getattr(attn, name, None) is None for name in ("k_proj", "v_proj", "o_proj", "rope")):
+        return False
+    weight = getattr(getattr(attn, "q_proj", None), "weight", None)
+    head_dim = getattr(attn, "head_dim", None)
+    if not head_dim and getattr(attn, "scale", None):
+        head_dim = round(float(attn.scale) ** -2)
+    num_heads = int(getattr(attn, "num_attention_heads", getattr(attn, "n_heads", 0)))
+    if weight is None or not head_dim or not num_heads:
+        return False
+    return int(weight.shape[0]) == num_heads * int(head_dim)
+
+
+def _attention_has_ungated_q_proj(attn: Any) -> bool:
+    """Ungated attention the hook body serves on a segmented cache (Qwen3, Llama, Mistral)."""
+    return _attention_has_plain_q_proj(attn) or _attention_has_unnormed_q_proj(attn)
+
+
 def _install_split_attention_hook(attn: Any) -> bool:
     cls = type(attn)
     if getattr(cls, "_mtplx_split_full_attention_installed", False):
@@ -350,7 +393,11 @@ def split_call(
         cache: Any | None = None,
     ) -> mx.array:
         route_off = not getattr(self, "_mtplx_split_full_attention_enabled", False)
-        if route_off and not cache_owns_rotary_origin(cache):
+        if (
+            route_off
+            and not cache_owns_rotary_origin(cache)
+            and not isinstance(cache, SegmentedKVCache)
+        ):
             # No MTPLX attention route asked for, and the cache has no rotary
             # origin of its own: the stock forward, untouched. A cache that
             # owns one (an image request on a compiled route) rotates at it
@@ -371,27 +418,44 @@ def split_call(
             # standalone sigmoid eagerly and a fused one inside the compiled
             # verifier's trace, which differ in the last bit of some bfloat16
             # and float32 values (mtplx/attention_math.py).
-        if not _attention_has_gated_q_proj(self):
+            # A segmented cache always takes the body (the segment route
+            # below), under any profile: the stock forward would gather the
+            # history on every call.
+        gated_q_proj = _attention_has_gated_q_proj(self)
+        if not gated_q_proj and not (
+            isinstance(cache, SegmentedKVCache) and _attention_has_ungated_q_proj(self)
+        ):
+            # Only the gated (Qwen3-Next style) attention goes through this body, and, for
+            # MTPLX_SEGMENTED_KV, an ungated attention (Qwen3 with q/k norm, Llama and
+            # Mistral without) on a segmented cache: the stock forward would gather the
+            # history on every call.
             return original_call(self, x, mask=mask, cache=cache)
 
         from mlx_lm.models.base import scaled_dot_product_attention
 
+        n_heads = getattr(self, "num_attention_heads", None) or self.n_heads
+        n_kv_heads = getattr(self, "num_key_value_heads", None) or self.n_kv_heads
         B, L, _ = x.shape
         q_proj_output = self.q_proj(x)
-        queries, gate = mx.split(
-            q_proj_output.reshape(B, L, self.num_attention_heads, -1),
-            2,
-            axis=-1,
-        )
-        gate = gate.reshape(B, L, -1)
+        if gated_q_proj:
+            queries, gate = mx.split(
+                q_proj_output.reshape(B, L, n_heads, -1),
+                2,
+                axis=-1,
+            )
+            gate = gate.reshape(B, L, -1)
+        else:
+            queries, gate = q_proj_output.reshape(B, L, n_heads, -1), None
 
         keys = self.k_proj(x)
         values = self.v_proj(x)
-        queries = self.q_norm(queries).transpose(0, 2, 1, 3)
-        keys = self.k_norm(
-            keys.reshape(B, L, self.num_key_value_heads, -1)
-        ).transpose(0, 2, 1, 3)
-        values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+        keys = keys.reshape(B, L, n_kv_heads, -1)
+        if getattr(self, "q_norm", None) is not None:
+            queries = self.q_norm(queries)
+            keys = self.k_norm(keys)
+        queries = queries.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, n_kv_heads, -1).transpose(
             0,
             2,
             1,
@@ -425,11 +489,51 @@ def split_call(
             and int(B) == 1
             and can_slice_mask
         )
+        segmented_output = None
+        segmented_route = ""
         if cache is not None:
             rope_offset = _cache_rope_offset(cache)
             queries = self.rope(queries, offset=rope_offset)
             keys = self.rope(keys, offset=rope_offset)
-            if blockwise_enabled or vllm_metal_paged_enabled:
+            if isinstance(cache, SegmentedKVCache):
+                # Sealed segments plus a tail (MTPLX_SEGMENTED_KV): the new rows go to
+                # the tail. A prefill chunk over several segments takes the logsumexp
+                # route and a verify window attends straight over the segments; any
+                # other call runs the ladder below on one view of the rows (zero-copy
+                # for a single segment, gathered otherwise).
+                cache.append_rows(keys, values)
+                if (
+                    int(queries.shape[2]) > VERIFY_WINDOW_MAX
+                    and cache.segment_count > 1
+                    and segmented_prefill_route() == "lse"
+                    and (mask is None or (isinstance(mask, str) and mask == "causal"))
+                ):
+                    segmented_output = attend_segments_lse(
+                        queries, cache, scale=self.scale
+                    )
+                    segmented_route = "segments_lse"
+                    segmented_count("prefill_lse")
+                else:
+                    segmented_output = decode_segments_attention(
+                        queries,
+                        cache,
+                        scale=self.scale,
+                        mask=mask,
+                        packed_threshold=int(
+                            getattr(self, "_mtplx_gqa_packed_sdpa_threshold", 8192)
+                        ),
+                    )
+                    segmented_route = "nax_flash_segments"
+                if segmented_output is None:
+                    q_rows = int(queries.shape[2])
+                    segmented_count(
+                        f"{'single' if cache.segment_count == 1 else 'gather'}_q"
+                        f"{q_rows if q_rows <= VERIFY_WINDOW_MAX else 'prefill'}"
+                    )
+                    view = gathered_view(cache)
+                    keys, values = view.keys, view.values
+                    cache = view
+            elif blockwise_enabled or vllm_metal_paged_enabled:
                 cache.update_without_fetch(keys, values)
             else:
                 keys, values = cache.update_and_fetch(keys, values)
@@ -559,7 +663,10 @@ def split_call(
         )
         route = "sdpa"
         fallback = ""
-        if should_use_vllm_metal_paged:
+        if segmented_output is not None:
+            output = segmented_output
+            route = segmented_route
+        elif should_use_vllm_metal_paged:
             impl_override = (
                 "fast_sdpa_gather"
                 if getattr(self, "_mtplx_vllm_metal_exact_gather_layer", False)
@@ -799,6 +906,8 @@ def split_call(
             )
         _note_kv_attention(self, cache, queries, output, mask, route, fallback)
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        if gate is None:
+            return self.o_proj(output)
         return self.o_proj(attention_gate(output, gate))
 
     cls.__call__ = split_call
@@ -883,6 +992,7 @@ def configure_split_full_attention(
     }
     full_layers = list(_full_attention_layers(model))
     full_layer_count = len(full_layers)
+    hooked_layers: list[bool] = []
     for full_idx, attn in enumerate(full_layers):
         exact_gather_layer = bool(
             vllm_metal_paged
@@ -894,7 +1004,13 @@ def configure_split_full_attention(
                 )
             )
         )
-        stats["installed"] += int(_install_split_attention_hook(attn))
+        newly_hooked = bool(_install_split_attention_hook(attn))
+        # The hook is installed once per attention class: later instances report False.
+        hooked_layers.append(
+            (newly_hooked or bool(getattr(type(attn), "_mtplx_split_full_attention_installed", False)))
+            and (_attention_has_gated_q_proj(attn) or _attention_has_ungated_q_proj(attn))
+        )
+        stats["installed"] += int(newly_hooked)
         attn._mtplx_split_full_attention_enabled = bool(
             active or sdpa_2pass or vllm_metal_paged or gqa_packed
         )
@@ -916,4 +1032,11 @@ def configure_split_full_attention(
         attn._mtplx_split_full_attention_calls = 0
         stats["layers"] += 1
         stats["exact_gather_layers"] += int(exact_gather_layer)
+    if segmented_kv_requested():
+        # MTPLX_SEGMENTED_KV: decide once, from the hooked layers, whether this model can use
+        # segments without ever falling back to a gather (mtplx/segmented_kv.py).
+        from .segmented_kv import evaluate_model_support
+
+        verdict = evaluate_model_support(full_layers, hooked=hooked_layers)
+        stats["segmented_kv_supported"] = bool(verdict["supported"])
     return stats