From 0384434bd3603c360a94a15d33739ddbb9a4c47f Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 16:11:33 +0200 Subject: [PATCH 01/37] SSD restore builds the request's cache layout (cache_factory) instead of runtime.make_cache() A restore from RAM builds the cache with the request's cache_factory; the SSD fallback (_restore_cold) ignored it and built the runtime's default layout, the paged KV cache for a long context: about 5x slower prefill and slower decode after every restore from disk. restore() now hands both factories to _restore_cold, which uses them exactly as the RAM restore does (meta state not imposed on the request's cache). Without a factory nothing changes. --- CHANGELOG.md | 6 ++ mtplx/session_bank.py | 27 ++++- tests/test_ssd_restore_cache_layout.py | 135 +++++++++++++++++++++++++ 3 files changed, 165 insertions(+), 3 deletions(-) create mode 100644 tests/test_ssd_restore_cache_layout.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5039769..bd422c787 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ 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] + +### 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 diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py index 458b38a55..61759f0a0 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -2361,6 +2361,8 @@ def cold_fallback() -> SessionBankRestore | None: draft_head_identity=draft_head_identity, policy_fingerprint=policy_fingerprint, ram_refusal=ram_refusal, + cache_factory=cache_factory, + mtp_cache_factory=mtp_cache_factory, ) entry = self.longest_prefix(token_ids) @@ -3419,7 +3421,12 @@ def _restore_cold( draft_head_identity: str | None, policy_fingerprint: str | None, ram_refusal: dict[str, Any] | None = None, + cache_factory: Callable[[], list[Any]] | None = None, + mtp_cache_factory: Callable[[], list[Any]] | None = None, ) -> SessionBankRestore | None: + # cache_factory / mtp_cache_factory: the request's cache layout, the same + # factories restore() hands to its RAM restore. + # # ram_refusal: restore() refused this lookup's longest RAM prefix # (ram_miss_reason, ram_refused_prefix_len, ram_refused_session_id). # It stays in the diagnostic whatever this lookup finds. @@ -3519,14 +3526,28 @@ def _restore_cold( ): self.last_miss_reason = CacheMissReason.SNAPSHOT_DESYNC.value return None - cache = runtime.make_cache() - restore_cache(cache, entry.cache_snapshot) + # The request's cache layout, exactly as a RAM restore builds it (restore() above). + # runtime.make_cache() is the default layout, which for a long context is the paged + # KV cache: about 5x slower prefill and slower decode than the dense cache the request + # asked for (Qwen3.8 27B, 2026-10-08: first token after an SSD restore at 50K and 80K + # 16 to 48 s with the default layout, 5 to 8 s with the request's). As in the RAM + # restore, the stored meta state is not imposed on the request's cache. + cache = cache_factory() if cache_factory is not None else runtime.make_cache() + restore_cache( + cache, + entry.cache_snapshot, + restore_meta_state=cache_factory is None, + ) # Read back from disk, the entry's floor is its restored caches'. entry.restore_floor_tokens = _cache_restore_floor(cache) entry.window_nbytes = _cache_window_nbytes(cache) mtp_history_cache = None if entry.mtp_history_snapshot is not None: - mtp_history_cache = runtime.make_mtp_cache() + mtp_history_cache = ( + mtp_cache_factory() + if mtp_cache_factory is not None + else runtime.make_mtp_cache() + ) restore_cache(mtp_history_cache, entry.mtp_history_snapshot) entry.hits += 1 entry.last_access_s = time.time() diff --git a/tests/test_ssd_restore_cache_layout.py b/tests/test_ssd_restore_cache_layout.py new file mode 100644 index 000000000..8de42ce85 --- /dev/null +++ b/tests/test_ssd_restore_cache_layout.py @@ -0,0 +1,135 @@ +"""An SSD restore builds the request's cache layout, the same one a RAM restore builds. + +Before, the SSD restore called ``runtime.make_cache()`` and ignored the request's +``cache_factory``: for a long context that default is the paged KV cache, with about 5x slower +prefill and slower decode than the dense cache the request asked for (Qwen3.8 27B, 2026-10-08: +first token after an SSD restore at 50K and 80K 16 to 48 s instead of 5 to 8 s). +""" + +from __future__ import annotations + +from pathlib import Path + +import mlx.core as mx +import numpy as np +from mlx_lm.models.cache import KVCache + +from mtplx.cache_bank.cold_tier import SessionBankColdTier +from mtplx.cache_state import snapshot_cache +from mtplx.session_bank import SessionBank + +HEADS, DIM, ROWS = 2, 16, 700 + + +class _DefaultLayoutKV(KVCache): + """Stands in for the runtime's default layout (the paged cache in production).""" + + +class _RequestLayoutKV(KVCache): + """Stands in for the layout the request asks for (the dense cache in production).""" + + +class _Runtime: + model_path = Path("models/example") + + def __init__(self, *, mtp: bool = False) -> None: + self.mtp_enabled = mtp + self.made = 0 + self.made_mtp = 0 + + def make_cache(self): + self.made += 1 + return [_DefaultLayoutKV() for _ in range(2)] + + def make_mtp_cache(self): + self.made_mtp += 1 + return [_DefaultLayoutKV()] + + +def _rand(seed: int, n: int) -> mx.array: + rng = np.random.default_rng(seed) + return mx.array(rng.standard_normal((1, HEADS, n, DIM)).astype(np.float32)).astype(mx.bfloat16) + + +def _filled(layers: int, seed: int) -> list[KVCache]: + cache = [KVCache() for _ in range(layers)] + for i, c in enumerate(cache): + c.update_and_fetch(_rand(seed + i, ROWS), _rand(seed + 50 + i, ROWS)) + return cache + + +def _tier(path: Path) -> SessionBankColdTier: + return SessionBankColdTier(base_dir=path, mode="on", min_prefix_tokens=2, block_size=256) + + +def _bank_on_disk(tmp_path: Path, runtime: _Runtime, *, mtp: bool = False): + """A conversation written to the SSD tier, then a fresh bank (empty RAM) on the same store.""" + tier = _tier(tmp_path / "tier") + tokens = list(range(ROWS)) + cache = _filled(2, 1) + kwargs = {} + if mtp: + kwargs = { + "mtp_history_snapshot": snapshot_cache(_filled(1, 9)), + "snapshot_epoch": ROWS, + "mtp_snapshot_epoch": ROWS, + } + SessionBank(cold_tier=tier).put( + runtime=runtime, + token_ids=tokens, + cache=cache, + logits=mx.zeros((1, 8)), + hidden=mx.ones((1, 4), dtype=mx.bfloat16) if mtp else None, + session_id="s", + **kwargs, + ) + assert tier.flush(timeout_s=30.0) + return SessionBank(cold_tier=tier), tokens, cache + + +def test_ssd_restore_builds_the_requests_cache_layout(tmp_path) -> None: + runtime = _Runtime() + bank, tokens, written = _bank_on_disk(tmp_path, runtime) + made_before = runtime.made + restored = bank.restore( + runtime, + tokens + [ROWS + 1, ROWS + 2], + session_id="s", + cache_factory=lambda: [_RequestLayoutKV() for _ in range(2)], + ) + assert restored is not None and restored.restore_mode == "ssd_clone" + assert bank.last_restore_source == "ssd" + assert all(type(c) is _RequestLayoutKV for c in restored.cache) + assert runtime.made == made_before # the default layout is not built at all + for got, want in zip(restored.cache, written): + assert got.offset == ROWS + gk, gv = got.state + wk, wv = want.state + assert mx.array_equal(gk, wk).item() and mx.array_equal(gv, wv).item() + + +def test_ssd_restore_without_a_factory_keeps_the_default_layout(tmp_path) -> None: + runtime = _Runtime() + bank, tokens, _ = _bank_on_disk(tmp_path, runtime) + restored = bank.restore(runtime, tokens + [ROWS + 1], session_id="s") + assert restored is not None and restored.restore_mode == "ssd_clone" + assert all(type(c) is _DefaultLayoutKV for c in restored.cache) + assert restored.cache[0].offset == ROWS + + +def test_ssd_restore_builds_the_mtp_history_with_the_requests_factory(tmp_path) -> None: + runtime = _Runtime(mtp=True) + bank, tokens, _ = _bank_on_disk(tmp_path, runtime, mtp=True) + made_mtp_before = runtime.made_mtp + restored = bank.restore( + runtime, + tokens + [ROWS + 1], + session_id="s", + cache_factory=lambda: [_RequestLayoutKV() for _ in range(2)], + mtp_cache_factory=lambda: [_RequestLayoutKV()], + ) + assert restored is not None and restored.restore_mode == "ssd_clone" + assert restored.mtp_history_cache is not None + assert type(restored.mtp_history_cache[0]) is _RequestLayoutKV + assert restored.mtp_history_cache[0].offset == ROWS + assert runtime.made_mtp == made_mtp_before From 923510415ba2ae7b49afeaab540373df667525b5 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 01:50:40 +0200 Subject: [PATCH 02/37] Segmented KV cache behind MTPLX_SEGMENTED_KV: segments, fused verify kernel, bank and admission --- CHANGELOG.md | 4 + mtplx/attention_split.py | 52 +- mtplx/cache_state.py | 6 + mtplx/graphbank.py | 7 + mtplx/kernels/sdpa_nax_flash_dsplit.py | 10 +- mtplx/kernels/sdpa_segmented.py | 337 +++++++++++ mtplx/segmented_kv.py | 745 +++++++++++++++++++++++++ mtplx/server/openai.py | 15 + mtplx/session_bank.py | 48 +- tests/test_sdpa_segmented.py | 132 +++++ tests/test_segmented_kv.py | 450 +++++++++++++++ tests/test_segmented_kv_attention.py | 133 +++++ tests/test_segmented_kv_bank.py | 263 +++++++++ 13 files changed, 2197 insertions(+), 5 deletions(-) create mode 100644 mtplx/kernels/sdpa_segmented.py create mode 100644 mtplx/segmented_kv.py create mode 100644 tests/test_sdpa_segmented.py create mode 100644 tests/test_segmented_kv.py create mode 100644 tests/test_segmented_kv_attention.py create mode 100644 tests/test_segmented_kv_bank.py diff --git a/CHANGELOG.md b/CHANGELOG.md index bd422c787..5fc81e93b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ All notable user-facing changes to MTPLX. The format is based on ## [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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Not covered: the compiled verify (segmented caches verify eager), the SSD tier (segmented entries are not persisted), paged and quantized KV. + ### 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. diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index 289e1d2ca..27c0aa5b0 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -14,6 +14,14 @@ ) from .attention_math import attention_gate from .compile_state import in_compiled_step_body, is_compile_trace_error +from .segmented_kv import ( + SegmentedKVCache, + attend_segments_lse, + decode_segments_attention, + gathered_view, + _count as segmented_count, + prefill_route as segmented_prefill_route, +) from .rope_origin import ( cache_owns_rotary_origin, note_unowned_rotary_origin, @@ -425,11 +433,48 @@ 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 verify window attends straight over the segments; every + # other call runs the ladder below on one gathered view of the rows. + cache.append_rows(keys, values) + if ( + int(queries.shape[2]) > 8 + 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") + elif ( + getattr(self, "_mtplx_gqa_packed_sdpa_enabled", False) + and not blockwise_enabled + and not vllm_metal_paged_enabled + ): + 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: + 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 +604,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) diff --git a/mtplx/cache_state.py b/mtplx/cache_state.py index 6e7a2546a..a859660af 100644 --- a/mtplx/cache_state.py +++ b/mtplx/cache_state.py @@ -4676,6 +4676,12 @@ def configure_tail_owned_attention_kv_cache(cache: list[Any]) -> dict[str, int | turboquant_config=turboquant_config, kv_quant_config=kv_quant_config_from_env(), ) + from .segmented_kv import install_segmented_attention_kv_cache, segmented_kv_enabled + + if segmented_kv_enabled(): + # MTPLX_SEGMENTED_KV=1: sealed segments plus a growing tail per full-attention + # layer (mtplx/segmented_kv.py). Paged/quantized KV (above) keeps its own layout. + return install_segmented_attention_kv_cache(cache) raw = os.environ.get("MTPLX_OWNED_ATTN_KV") or "" normalized = raw.strip().lower().replace("-", "_") if normalized not in { diff --git a/mtplx/graphbank.py b/mtplx/graphbank.py index 9a66cf2ea..260e69bae 100644 --- a/mtplx/graphbank.py +++ b/mtplx/graphbank.py @@ -1510,6 +1510,13 @@ def promote_kv_cache_offsets( cache[idx] = TensorOffsetVllmMetalPagedKVCache.from_paged_cache(entry) promoted += 1 continue + from .segmented_kv import SegmentedKVCache + + if isinstance(entry, SegmentedKVCache): + # A segmented KV cache (MTPLX_SEGMENTED_KV) has no single buffer for the + # tensor-offset adapter to adopt; the verify runs eager on it. + failures["segmented_kv_cache"] = failures.get("segmented_kv_cache", 0) + 1 + continue offset = getattr(entry, "offset", None) if not isinstance(offset, int): continue diff --git a/mtplx/kernels/sdpa_nax_flash_dsplit.py b/mtplx/kernels/sdpa_nax_flash_dsplit.py index 627ab45ef..3baa7acfb 100644 --- a/mtplx/kernels/sdpa_nax_flash_dsplit.py +++ b/mtplx/kernels/sdpa_nax_flash_dsplit.py @@ -34,7 +34,12 @@ # its bails) — the one-line receipt the #459 reports needed. nax_flash_dsplit_dispatch_counts: dict[str, int] = {} -# Template params: InT, PartT, D, QL, GQA_F. +# Template params: InT, PartT, D, QL, GQA_F, NOMASK. +# NOMASK=1 makes every key below ``offset`` visible to every query row (no tail-causal +# mask): a sealed segment of a segmented KV cache sits entirely before the rows being +# verified (mtplx/kernels/sdpa_segmented.py). The kernel derives ``tail_lo`` from the same +# ``offset`` that bounds the keys, so the mask cannot be switched off through the inputs. +# NOMASK=0 (every other caller) compiles to the previous kernel. _SOURCE = r""" constexpr int TK = 32; constexpr int MROWS = 16; @@ -172,7 +177,7 @@ for (short j = 0; j < kElemCols; j++) { const int gp = t0 + hh * 16 + sc.x + j; const float raw = s_raw[hh][i * kElemCols + j]; - const bool vis = live && gp < kv_end && gp <= row_limit; + const bool vis = live && gp < kv_end && (NOMASK != 0 || gp <= row_limit); const float s = vis ? raw * scale : -1e38f; s_p[hh][i * kElemCols + j] = s; tmax[i] = metal::max(tmax[i], s); @@ -366,6 +371,7 @@ def sdpa_nax_flash_dsplit( ("D", d), ("QL", q_len), ("GQA_F", gqa_factor), + ("NOMASK", 0), ], grid=(hk * nthreads, 1, blocks), threadgroup=(nthreads, 1, 1), diff --git a/mtplx/kernels/sdpa_segmented.py b/mtplx/kernels/sdpa_segmented.py new file mode 100644 index 000000000..accdfc71a --- /dev/null +++ b/mtplx/kernels/sdpa_segmented.py @@ -0,0 +1,337 @@ +"""Decode/verify attention over a list of KV segments (segmented KV cache). + +A conversation's full-attention KV is a list of sealed segments plus one growing tail +(``mtplx/segmented_kv.py``). Attention over the list is split-softmax-merge, the way +``sdpa_nax_flash_dsplit`` already splits one buffer into key blocks: every block yields an +unnormalised partial output, a row sum and a row max, and one reduce merges them exactly +(global max, exp-sum, weighted sum). A segment is just more blocks. + +Two kernels: + +* the partial kernel is the head-dim-split TensorOps kernel of ``sdpa_nax_flash_dsplit``, + derived from its source at import (``_fused_source``), launched ONCE for up to + ``MAX_FUSED_SEGMENTS`` segments: grid.z walks the blocks of all segments, a ``meta`` array + (block starts, row counts, capacities) selects the segment and its K/V buffers inside the + kernel. Segments before the last one run with every key visible (no tail-causal mask); + only the last segment, which holds the rows being verified, keeps the mask; +* the reduce kernel is ``_paged_reduce_kernel`` generalised to a list of partial buffers + with arbitrary block counts (blocks are addressed through a ``starts`` prefix table). + +Receipts (phase 1, 2026-10-07, Qwen3.8-27B shapes: 4 KV heads, 24 query heads, head_dim 256, +bf16; docs/onderzoek/2026-10-07-segmented-kv-prototype.md in the research branch): +segmented output equals the contiguous kernel's error against an fp32 reference (rel L2 +2.8e-3, 1 segment bit-identical, more segments within one bf16 ulp of the contiguous +kernel); cost per extra segment at q_len 4 is 0 to 0.005 ms/layer fused up to 12 segments +(0.012 to 0.016 with one launch per segment); a verify forward over 6 or 20 segments at 50K +and 80K costs 0 to 1.1 % more than the contiguous one. + +Routes that cannot take segments yet fall back to ``segmented_kv.gather_attention`` (one +contiguous copy of the history per call, then the stock kernel): the wide (q_len 6 to 16) +and dsplit4 variants are not in this tree, and the head-dim-split kernel serves q_len 1 to 5 +at GQA 6 (``gqa_factor * q_len <= 32``). +""" + +from __future__ import annotations + +import os +from collections.abc import Sequence +from functools import lru_cache + +import mlx.core as mx + +from ..nax_verify import nax_available +from . import sdpa_nax_flash_dsplit as _dsplit +from .sdpa_2pass import unnormalized_partials_dtype +from .sdpa_nax_flash import _HEADER + +#: Metal allows 31 buffer bindings: queries + 2 per segment + meta + scale + 3 outputs. +MAX_FUSED_SEGMENTS = 12 +#: The reduce kernel binds one partial buffer per launch group plus sums, maxs, starts, out. +MAX_GROUPS = 20 + +segmented_bail_counts: dict[str, int] = {} +segmented_dispatch_counts: dict[str, int] = {} + + +def _bail(reason: str): + segmented_bail_counts[reason] = segmented_bail_counts.get(reason, 0) + 1 + if os.environ.get("MTPLX_NAX_FLASH_DEBUG"): + print(f"[nax-flash-segments bail] {reason}") + + +_OLD_HEAD = """ const int block_idx = int(threadgroup_position_in_grid.z); + const int n_blocks = int(blocks[0]); + + const int n_kv = static_cast(offset[0]);""" + + +def _fused_source(nseg: int) -> str: + """The partial kernel over ``nseg`` segment buffers in one launch. + + Same body as ``sdpa_nax_flash_dsplit._SOURCE``; only the prologue (segment lookup, + buffer selection, per-segment mask flag) and the partial/stat indexing change. + ``meta`` = bstart[nseg+1] | n[nseg] | cap[nseg] (int32). + """ + src = _dsplit._SOURCE + if _OLD_HEAD not in src: + raise RuntimeError("sdpa_nax_flash_dsplit source changed: update sdpa_segmented") + sel = "".join(f" if (seg == {i}) {{ keys = K{i}; values = V{i}; }}\n" for i in range(1, nseg)) + new_head = f""" const int gblock = int(threadgroup_position_in_grid.z); + const int B_total = meta[{nseg}]; + int seg = 0; + while (seg + 1 < {nseg} && gblock >= meta[seg + 1]) seg++; + const int block_idx = gblock - meta[seg]; + const int n_blocks = meta[seg + 1] - meta[seg]; + const int n_kv = meta[{nseg + 1} + seg]; + const int kcap = meta[{2 * nseg + 1} + seg]; + const bool nomask_seg = (seg < NLAST); + const device InT* keys = K0; + const device InT* values = V0; +{sel}""" + src = src.replace(_OLD_HEAD, new_head) + for old, new in ( + ("(NOMASK != 0 || gp <= row_limit)", "(nomask_seg || gp <= row_limit)"), + ("((size_t)hq_row * n_blocks + block_idx) * D + db", "((size_t)hq_row * B_total + gblock) * D + db"), + ("sums[hq_row * n_blocks + block_idx]", "sums[hq_row * B_total + gblock]"), + ("maxs[hq_row * n_blocks + block_idx]", "maxs[hq_row * B_total + gblock]"), + ): + if old not in src: + raise RuntimeError(f"sdpa_nax_flash_dsplit source changed ({old!r}): update sdpa_segmented") + src = src.replace(old, new) + return src + + +@lru_cache(maxsize=16) +def _fused_kernel(nseg: int): + try: + return mx.fast.metal_kernel( + name=f"mtplx_nax_flash_dsplit_segments_{nseg}", + input_names=["queries"] + [f"K{i}" for i in range(nseg)] + [f"V{i}" for i in range(nseg)] + ["meta", "scale"], + output_names=["partials", "sums", "maxs"], + header=_HEADER, + source=_fused_source(nseg), + ) + except Exception: # noqa: BLE001 — toolchain without Metal4/MPP support + return None + + +def _reduce_source(ngroups: int) -> str: + """Merge kernel over ``ngroups`` partial buffers with per-buffer block counts. + + Same arithmetic as ``mtplx_sdpa_2pass_paged_reduce`` (global max over all blocks, global + exp-sum, then the sum of ``exp(m_b - m) * partial_b`` divided by the exp-sum); only the + block loop is generalised: ``starts`` (int32[ngroups+1]) holds the prefix sums of the + per-buffer block counts, so counts need not be multiples of 32. + """ + acc = [] + for s in range(ngroups): + acc.append(f""" + {{ + const int b0 = starts[{s}]; + const int nb = starts[{s + 1}] - b0; + const auto pbase = partials{s} + (size_t)q_offset * nb * V + lane * ept; + for (int lb = sg; lb < nb; lb += BN) {{ + U factor = fast::exp(maxs[b0 + lb] - max_score); + for (int i = 0; i < ept; ++i) {{ + o[i] += factor * static_cast(pbase[(size_t)lb * V + i]); + }} + }} + }}""") + head = """ + constexpr int BN = 32; + constexpr int BD = 32; + constexpr int ept = V / BD; + typedef float U; + thread U o[ept] = {0}; + threadgroup U outputs[BN * BD]; + const int head_idx = threadgroup_position_in_grid.x; + const int q_seq_idx = threadgroup_position_in_grid.y; + const int q_offset = head_idx * threadgroups_per_grid.y + q_seq_idx; + const int sg = simdgroup_index_in_threadgroup; + const int lane = thread_index_in_simdgroup; + const int btot = starts[NGROUPS]; + sums += (size_t)q_offset * btot; + maxs += (size_t)q_offset * btot; + out += q_offset * V + sg * ept; + U sum_exp_score = 0.0f; + U max_score = Limits::finite_min; + for (int b = lane; b < btot; b += BN) max_score = metal::max(max_score, maxs[b]); + max_score = simd_max(max_score); + for (int b = lane; b < btot; b += BN) + sum_exp_score += fast::exp(maxs[b] - max_score) * sums[b]; + sum_exp_score = simd_sum(sum_exp_score); + """.replace("NGROUPS", str(ngroups)) + tail = """ + for (int i = 0; i < ept; ++i) { + outputs[lane * BD + sg] = o[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + o[i] = simd_sum(outputs[sg * BD + lane]); + o[i] = sum_exp_score == 0.0f ? o[i] : (o[i] / sum_exp_score); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lane == 0) { + for (int i = 0; i < ept; ++i) out[i] = static_cast(o[i]); + } + """ + return head + "".join(acc) + tail + + +@lru_cache(maxsize=32) +def _reduce_kernel(ngroups: int): + try: + return mx.fast.metal_kernel( + name=f"mtplx_nax_flash_segments_reduce_{ngroups}", + input_names=[f"partials{s}" for s in range(ngroups)] + ["sums", "maxs", "starts"], + output_names=["out"], + source=_reduce_source(ngroups), + ) + except Exception: # noqa: BLE001 + return None + + +def segment_block_counts(rows: Sequence[int]) -> list[int]: + """Key blocks per segment: the chunk the contiguous dispatcher uses for the same total.""" + total = sum(int(n) for n in rows) + override = int(os.environ.get("MTPLX_NAX_FLASH_DSPLIT_BLOCKS", "0") or 0) + base = override or _dsplit._default_dsplit_blocks(total) + base = max(32, (base // 32) * 32) + chunk = (-(-total // base) + 31) // 32 * 32 + return [max(1, -(-int(n) // chunk)) for n in rows] + + +def sdpa_nax_flash_dsplit_segments( + *, + queries: mx.array, + segments: Sequence[tuple[mx.array, mx.array, int]], + scale: float, +) -> mx.array | None: + """Exact attention of ``queries`` over ``segments`` = [(keys, values, n), ...], oldest first. + + ``keys``/``values`` are (1, Hk, capacity, 256) buffers of which the first ``n`` rows are + live. Every segment but the last is fully visible; the last holds the rows being verified + as its final ``q_len`` rows and keeps the tail-causal mask. None when the contract is not + met (the caller falls back to ``segmented_kv.gather_attention``). + """ + if os.environ.get("MTPLX_NAX_FLASH_DSPLIT", "1") == "0": + return _bail("env_disabled") + if not mx.metal.is_available(): + return _bail("metal_unavailable") + if not nax_available(): + return _bail("gpu_family_or_os") + if queries.ndim != 4: + return _bail("ndim") + bsz, hq, q_len, d = (int(x) for x in queries.shape) + if bsz != 1 or d != 256: + return _bail("shape_gate") + if not segments: + return _bail("no_segments") + hk = int(segments[0][0].shape[1]) + if hk <= 0 or hq % hk: + return _bail("gqa_heads") + gqa_factor = hq // hk + if gqa_factor * q_len > 32: + return _bail("m_rows_gt_32") + if q_len < 1 or q_len > 10: + return _bail("q_len") + if queries.dtype not in (mx.bfloat16, mx.float16): + return _bail("query_dtype") + for keys, values, n in segments: + if keys.dtype != queries.dtype or values.dtype != queries.dtype: + return _bail("kv_dtype_mismatch") + if ( + keys.ndim != 4 + or int(keys.shape[1]) != hk + or tuple(values.shape) != tuple(keys.shape) + or int(keys.shape[3]) != d + ): + return _bail("kv_layout_mismatch") + if int(n) <= 0 or int(n) > int(keys.shape[2]): + return _bail("segment_rows") + if int(segments[-1][2]) < q_len: + return _bail("tail_shorter_than_q") + + if len(segments) == 1: + keys, values, n = segments[0] + return _dsplit.sdpa_nax_flash_dsplit( + queries=queries, keys=keys, values=values, offset=int(n), scale=scale + ) + + groups = [segments[i:i + MAX_FUSED_SEGMENTS] for i in range(0, len(segments), MAX_FUSED_SEGMENTS)] + if len(groups) > MAX_GROUPS: + return _bail("too_many_segments") + rows = [int(n) for _, _, n in segments] + blocks = segment_block_counts(rows) + reduce_kernel = _reduce_kernel(len(groups)) + if reduce_kernel is None: + return _bail("kernel_unavailable") + + nsgm = (gqa_factor * q_len + 15) // 16 + nthreads = 32 * nsgm * 2 + pdt = unnormalized_partials_dtype(queries.dtype) + parts, sums, maxs, counts = [], [], [], [] + first = 0 + for gi, group in enumerate(groups): + nseg = len(group) + kernel = _fused_kernel(nseg) + if kernel is None: + return _bail("kernel_unavailable") + nbs = blocks[first:first + nseg] + first += nseg + bstart = [0] + for nb in nbs: + bstart.append(bstart[-1] + nb) + total_blocks = bstart[-1] + meta = mx.array( + bstart + [int(n) for _, _, n in group] + [int(k.shape[2]) for k, _, _ in group], + dtype=mx.int32, + ) + last_group = gi == len(groups) - 1 + try: + p, s_, m_ = kernel( + inputs=[queries] + [k for k, _, _ in group] + [v for _, v, _ in group] + [meta, float(scale)], + template=[ + ("InT", queries.dtype), + ("PartT", pdt), + ("D", d), + ("QL", q_len), + ("GQA_F", gqa_factor), + ("NOMASK", 0), + ("NLAST", nseg - 1 if last_group else nseg), + ], + grid=(hk * nthreads, 1, total_blocks), + threadgroup=(nthreads, 1, 1), + output_shapes=[ + (bsz, hq, q_len, total_blocks, d), + (bsz, hq, q_len, total_blocks), + (bsz, hq, q_len, total_blocks), + ], + output_dtypes=[pdt, mx.float32, mx.float32], + ) + except Exception as exc: # noqa: BLE001 — dispatch/compile failure => gather fallback + return _bail(f"dispatch_failed: {type(exc).__name__}: {str(exc)[:2000]}") + parts.append(p) + sums.append(s_) + maxs.append(m_) + counts.append(total_blocks) + + sums_all = sums[0] if len(sums) == 1 else mx.concatenate(sums, axis=3) + maxs_all = maxs[0] if len(maxs) == 1 else mx.concatenate(maxs, axis=3) + starts = [0] + for c in counts: + starts.append(starts[-1] + c) + (out,) = reduce_kernel( + inputs=parts + [sums_all, maxs_all, mx.array(starts, dtype=mx.int32)], + template=[("InT", queries.dtype), ("V", d)], + grid=(bsz * hq * 1024, q_len, 1), + threadgroup=(1024, 1, 1), + output_shapes=[(bsz, hq, q_len, d)], + output_dtypes=[queries.dtype], + ) + segmented_dispatch_counts["dispatched"] = segmented_dispatch_counts.get("dispatched", 0) + 1 + return out + + +def segments_supported(q_len: int, gqa_factor: int, head_dim: int) -> bool: + """Static part of the contract (no device work), for the route gate in attention_split.""" + return head_dim == 256 and 1 <= q_len <= 10 and gqa_factor * q_len <= 32 + diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py new file mode 100644 index 000000000..06a28cb0a --- /dev/null +++ b/mtplx/segmented_kv.py @@ -0,0 +1,745 @@ +"""Segmented KV cache for the full-attention layers (MTPLX_SEGMENTED_KV=1, default off). + +A conversation's attention KV is a list of immutable SEGMENTS plus one growing TAIL, instead +of one contiguous buffer that the session bank aliases with a lazy snapshot: + +* a sealed segment has its exact capacity and is never written again; +* the tail is the only buffer a decode, a verify forward or a prefill chunk writes into; + verify rollback (``trim``) only moves the tail's row count; +* ``seal()`` closes the tail at the end of a turn (a copy of the tail only, never of the + history) and a new tail starts empty at the next write; +* a snapshot of the cache (``state``) IS the list of segment references: it seals the tail + and hands out ``SegmentedKVState`` (segment, rows) references. Nothing a later turn writes + can alias them, so MLX's copy-on-write never copies the history. A fork or a return into the + middle of a segment is a reference ``(segment, n')`` with n' below the segment's rows: + no copy, and the rows past n' belong to whoever else holds the segment. + +Why (receipts, 2026-10-06/07, Qwen3.8-27B, 16 full-attention layers, 64 KiB of KV per token): +the session bank's snapshot views alias the per-layer KV buffers, so the first write of the +next turn copies every buffer: +4.95 GiB at 80K (+3.1 GiB at 50K), whether the restore is a +clone or a reference lease (docs/onderzoek/2026-10-07-lease-vs-clone.md). With segments the +same follow-up turn costs the new segment only (+0.125 GiB for 1,024 tokens plus slack, +2026-10-07 phase 1 measurement) and a branch from the same history costs 0. + +Attention over the segments: ``mtplx/kernels/sdpa_segmented.py`` (exact split-softmax merge, +one launch for up to 12 segments) for the verify/decode window; everything else gathers the +rows into one contiguous array per layer and call (``gather``, correct and slower for long +histories) and runs the stock route on it. The prefill over existing segments has its own +interface, ``attend_segments_lse`` (per segment ``(out, lse)`` plus one merge), so that an +MLX kernel which returns the logsumexp can replace ``segment_sdpa_lse``. + +GatedDeltaNet (recurrent) layers keep their state objects as they are. + +Reference counts. A segment counts its holders (live caches and ``SegmentedKVState`` +snapshots) in a weak set: ``refcount`` is the number of holders still alive, an evicted or +dropped bank entry releases its segments by going away, and the merge policy only merges +segments that no snapshot holds. +""" + +from __future__ import annotations + +import itertools +import os +import weakref +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any + +import mlx.core as mx + +_SEGMENT_IDS = itertools.count(1) + + +def segmented_kv_enabled() -> bool: + """MTPLX_SEGMENTED_KV (default off), read on every call.""" + return os.environ.get("MTPLX_SEGMENTED_KV", "").strip().lower() in {"1", "true", "yes", "on"} + + +def _int_env(name: str, default: int) -> int: + raw = os.environ.get(name, "").strip() + try: + return int(raw) if raw else default + except ValueError: + return default + + +#: Receipts: which route served segmented prefill chunks and which per-segment attention ran. +route_counts: dict[str, int] = {} + + +def _count(route: str) -> None: + route_counts[route] = route_counts.get(route, 0) + 1 + + +_SDPA_LSE: bool | None = None + + +def _probe_sdpa_lse() -> bool: + """Whether this MLX build's fused SDPA can return the logsumexp (fork branch sdpa-lse). + + The one place that knows the API: ``mx.fast.scaled_dot_product_attention(..., return_lse=True)`` + returning ``(out, lse)``. Stock released MLX raises TypeError on the keyword, which is the + "absent" answer; any other failure also counts as absent. + """ + try: + q = mx.zeros((1, 1, 1, 64), dtype=mx.bfloat16) + result = mx.fast.scaled_dot_product_attention(q, q, q, scale=1.0, return_lse=True) + mx.eval(result) + return isinstance(result, (tuple, list)) and len(result) == 2 + except Exception: # noqa: BLE001 + return False + + +def sdpa_lse_available() -> bool: + """Feature detection, once per process (first use).""" + global _SDPA_LSE + if _SDPA_LSE is None: + _SDPA_LSE = _probe_sdpa_lse() + return _SDPA_LSE + + +def prefill_route() -> str: + """How a prefill chunk attends over existing segments. + + ``lse``: per segment (out, lse) plus one merge. Default when the MLX kernel with LSE output + is present; ``MTPLX_SEGMENTED_KV_PREFILL=lse`` forces it (exact reference route without the + kernel). ``gather`` (default on stock MLX): one contiguous copy per layer and call, then the + stock SDPA. + """ + raw = os.environ.get("MTPLX_SEGMENTED_KV_PREFILL", "").strip().lower() + if raw in {"lse", "gather"}: + return raw + return "lse" if sdpa_lse_available() else "gather" + + +class KVSegment: + """One K/V buffer pair, (1, Hk, capacity, D), shared by reference. + + ``sealed`` segments hold exactly ``rows == capacity`` live rows and are never written. + The tail of a cache is an unsealed segment: only its owner writes it. + """ + + __slots__ = ("__weakref__", "_holders", "id", "keys", "sealed", "values") + + def __init__(self, keys: mx.array, values: mx.array, *, sealed: bool = False) -> None: + self.id = next(_SEGMENT_IDS) + self.keys = keys + self.values = values + self.sealed = bool(sealed) + self._holders: weakref.WeakSet[Any] = weakref.WeakSet() + + @property + def capacity(self) -> int: + return int(self.keys.shape[2]) + + @property + def nbytes(self) -> int: + return int(self.keys.nbytes) + int(self.values.nbytes) + + @property + def refcount(self) -> int: + """Holders alive right now: live caches and snapshot states.""" + return len(self._holders) + + def hold(self, holder: Any) -> None: + self._holders.add(holder) + + def release(self, holder: Any) -> None: + self._holders.discard(holder) + + +@dataclass(frozen=True) +class SegRef: + """The first ``n`` rows of a segment.""" + + segment: KVSegment + n: int + + +class SegmentedKVState: + """An immutable snapshot of a segmented cache: references, no buffers of its own. + + This is what the session bank stores as the cache snapshot of a segmented layer. It holds + its segments (reference counts) for as long as it lives. + """ + + def __init__(self, refs: Iterable[SegRef]) -> None: + self.refs: tuple[SegRef, ...] = tuple(refs) + for ref in self.refs: + ref.segment.hold(self) + + @property + def rows(self) -> int: + return sum(int(ref.n) for ref in self.refs) + + @property + def nbytes(self) -> int: + """Bytes of the segments this state references, counted as if it were alone.""" + return sum(ref.segment.nbytes for ref in _unique_refs(self.refs)) + + def segments(self) -> list[KVSegment]: + return [ref.segment for ref in self.refs] + + def prefix(self, rows: int) -> SegmentedKVState: + """The first ``rows`` rows as a new state (shortens the last reference, no copy).""" + return SegmentedKVState(_prefix_refs(self.refs, rows)) + + def to_arrays(self) -> tuple[mx.array, mx.array] | None: + """One contiguous (keys, values) pair, the state a stock KV cache would give.""" + return _gather_refs(self.refs) + + def __len__(self) -> int: + return len(self.refs) + + +def _unique_refs(refs: Iterable[SegRef]) -> list[SegRef]: + seen: set[int] = set() + out: list[SegRef] = [] + for ref in refs: + if ref.segment.id not in seen: + seen.add(ref.segment.id) + out.append(ref) + return out + + +def _prefix_refs(refs: Sequence[SegRef], rows: int) -> list[SegRef]: + rows = max(0, int(rows)) + out: list[SegRef] = [] + left = rows + for ref in refs: + if left <= 0: + break + take = min(int(ref.n), left) + out.append(SegRef(ref.segment, take) if take != ref.n else ref) + left -= take + return out + + +def _rows_view(array: mx.array, n: int) -> mx.array: + return array if int(array.shape[2]) == int(n) else array[..., : int(n), :] + + +def _gather_refs(refs: Sequence[SegRef]) -> tuple[mx.array, mx.array] | None: + if not refs: + return None + keys = [_rows_view(ref.segment.keys, ref.n) for ref in refs] + values = [_rows_view(ref.segment.values, ref.n) for ref in refs] + if len(refs) == 1: + return keys[0], values[0] + return mx.concatenate(keys, axis=2), mx.concatenate(values, axis=2) + + +def _own_copy(array: mx.array) -> mx.array: + """An evaluated array that owns its buffer (a slice of a larger buffer is copied).""" + owned = mx.contiguous(array) + mx.eval(owned) + return owned + + +class SegmentedKVCache: + """KV cache of one full-attention layer as sealed segments plus a growing tail. + + Drop-in for the stock ``KVCache`` where the model code needs it: ``update_and_fetch``, + ``offset`` (the TOTAL number of rows, which is also the rotary offset and what every + route gate must read instead of a buffer capacity), ``make_mask``, ``trim``, ``state``. + """ + + step = 256 + + def __init__(self, *, step: int = 256) -> None: + self.step = int(step) + self._sealed: list[SegRef] = [] + self._tail: KVSegment | None = None + self._tail_n = 0 + self.seal_copies = 0 + self.seal_copy_rows = 0 + self.merges = 0 + self.merge_rows = 0 + + # ---- shape ---------------------------------------------------------------------------- + + @property + def offset(self) -> int: + return sum(int(ref.n) for ref in self._sealed) + self._tail_n + + @offset.setter + def offset(self, value: int) -> None: + # Moving the offset is a trim (the stock cache has no other use for the setter here). + value = int(value) + current = self.offset + if value > current: + raise ValueError("a segmented KV cache cannot grow its offset without rows") + self.trim(current - value) + + @property + def tail_rows(self) -> int: + return self._tail_n + + @property + def segment_count(self) -> int: + return len(self._sealed) + (1 if self._tail_n else 0) + + def size(self) -> int: + return self.offset + + def empty(self) -> bool: + return self.offset == 0 + + def is_trimmable(self) -> bool: + return True + + @property + def nbytes(self) -> int: + """Bytes this cache keeps allocated: its unique segments plus the tail buffer.""" + total = sum(ref.segment.nbytes for ref in _unique_refs(self._sealed)) + if self._tail is not None: + total += self._tail.nbytes + return total + + @property + def meta_state(self) -> str: + return "" + + @meta_state.setter + def meta_state(self, value: Any) -> None: + del value + + def make_mask(self, *args, **kwargs): + from mlx_lm.models.cache import create_attention_mask + + return create_attention_mask(*args, offset=self.offset, **kwargs) + + # ---- writing -------------------------------------------------------------------------- + + def append_rows(self, keys: mx.array, values: mx.array) -> None: + """Write new rows into the tail (the only buffer this cache ever writes).""" + if self._tail is None: + self.compact() + prev = self._tail_n + steps = int(keys.shape[2]) + tail = self._tail + if tail is None or prev + steps > tail.capacity: + b, n_kv_heads, _, k_head_dim = keys.shape + v_head_dim = values.shape[3] + n_steps = (self.step + steps - 1) // self.step + new_k = mx.zeros((b, n_kv_heads, n_steps * self.step, k_head_dim), keys.dtype) + new_v = mx.zeros((b, n_kv_heads, n_steps * self.step, v_head_dim), values.dtype) + if tail is not None: + old_k, old_v = tail.keys, tail.values + if prev % self.step != 0: + old_k, old_v = old_k[..., :prev, :], old_v[..., :prev, :] + new_k = mx.concatenate([old_k, new_k], axis=2) + new_v = mx.concatenate([old_v, new_v], axis=2) + tail = KVSegment(new_k, new_v) + self._tail = tail + tail.keys[..., prev : prev + steps, :] = keys + tail.values[..., prev : prev + steps, :] = values + self._tail_n = prev + steps + + def update_and_fetch(self, keys: mx.array, values: mx.array) -> tuple[mx.array, mx.array]: + self.append_rows(keys, values) + gathered = self.gather() + assert gathered is not None + return gathered + + # ---- reading -------------------------------------------------------------------------- + + def attention_segments(self) -> list[tuple[mx.array, mx.array, int]]: + """(keys, values, rows) per segment, oldest first; the last is the tail when it has rows.""" + out = [(ref.segment.keys, ref.segment.values, int(ref.n)) for ref in self._sealed] + if self._tail is not None and self._tail_n: + out.append((self._tail.keys, self._tail.values, self._tail_n)) + return out + + def gather(self) -> tuple[mx.array, mx.array] | None: + """One contiguous (keys, values) over all rows (lazy; a view when there is one segment).""" + parts = self.attention_segments() + if not parts: + return None + keys = [_rows_view(k, n) for k, _, n in parts] + values = [_rows_view(v, n) for _, v, n in parts] + if len(parts) == 1: + return keys[0], values[0] + return mx.concatenate(keys, axis=2), mx.concatenate(values, axis=2) + + @property + def keys(self) -> mx.array | None: + gathered = self.gather() + return None if gathered is None else gathered[0] + + @property + def values(self) -> mx.array | None: + gathered = self.gather() + return None if gathered is None else gathered[1] + + # ---- rollback ------------------------------------------------------------------------- + + def trim(self, n: int) -> int: + """Drop the last ``n`` rows. Verify rollback only ever moves the tail's row count.""" + n = min(self.offset, max(0, int(n))) + left = n + take = min(self._tail_n, left) + self._tail_n -= take + left -= take + if self._tail_n == 0 and self._tail is not None and left > 0: + self._tail = None + while left > 0 and self._sealed: + ref = self._sealed[-1] + if ref.n <= left: + self._sealed.pop() + if not any(r.segment is ref.segment for r in self._sealed): + ref.segment.release(self) + left -= int(ref.n) + else: + self._sealed[-1] = SegRef(ref.segment, int(ref.n) - left) + left = 0 + return n + + # ---- sealing and state ---------------------------------------------------------------- + + def seal(self) -> bool: + """Close the tail: a segment of exactly its rows, then a fresh empty tail at the next write. + + Copies the tail's rows only (never a sealed segment); a tail that already has exactly + its capacity is adopted without a copy. Returns True when a segment was added. + """ + tail, rows = self._tail, self._tail_n + if tail is None: + return False + if rows == 0: + self._tail = None + return False + if rows == tail.capacity: + keys, values = tail.keys, tail.values + mx.eval(keys, values) + else: + keys = _own_copy(tail.keys[..., :rows, :]) + values = _own_copy(tail.values[..., :rows, :]) + self.seal_copies += 1 + self.seal_copy_rows += rows + self._tail = None + self._tail_n = 0 + segment = KVSegment(keys, values, sealed=True) + segment.hold(self) + self._sealed.append(SegRef(segment, rows)) + return True + + @property + def state(self) -> SegmentedKVState: + """The snapshot of this cache: seals the tail, then the list of references.""" + self.seal() + return SegmentedKVState(self._sealed) + + @state.setter + def state(self, value: Any) -> None: + for ref in self._sealed: + ref.segment.release(self) + self._sealed = [] + self._tail = None + self._tail_n = 0 + if value is None: + return + if isinstance(value, SegmentedKVState): + refs = list(value.refs) + elif isinstance(value, (tuple, list)) and len(value) == 2 and value[0] is not None: + # A contiguous pair (an SSD restore, a stock snapshot): one sealed segment. + keys, values = value + segment = KVSegment(keys, values, sealed=True) + refs = [SegRef(segment, int(keys.shape[2]))] + else: + raise ValueError(f"unsupported segmented KV state: {type(value).__name__}") + for ref in refs: + ref.segment.hold(self) + self._sealed = refs + + # ---- merge policy --------------------------------------------------------------------- + + def compact(self, *, max_segments: int | None = None) -> int: + """Merge small recent sealed segments among themselves (geometric tiers). + + Two adjacent sealed segments merge when no snapshot holds either (reference count 1: + this cache), the older is at most ``MTPLX_SEGMENTED_KV_TIER`` times (default 2) the + newer one and the result stays within ``MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS`` (default + 16384) rows. The big history segment is never merged here: a full copy would bring the + duplicate back while a snapshot, a subagent or a lease references it. Copying two small + segments costs 0.8 ms (512 + 512 rows) to 4.1 ms (4096 + 4096 rows) for 16 layers (phase 1 + measurement). Returns the number of merges. + """ + tier = max(1, _int_env("MTPLX_SEGMENTED_KV_TIER", 2)) + cap_rows = _int_env("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", 16384) + limit = max_segments if max_segments is not None else _int_env("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", 12) + merged = 0 + while len(self._sealed) >= 2: + older, newer = self._sealed[-2], self._sealed[-1] + alone = older.segment.refcount <= 1 and newer.segment.refcount <= 1 + small = int(older.n) + int(newer.n) <= cap_rows + geometric = int(older.n) <= tier * int(newer.n) + over = len(self._sealed) > limit and alone + if not (alone and ((small and geometric) or over)): + break + self._merge_last_two() + merged += 1 + return merged + + def _merge_last_two(self) -> None: + older, newer = self._sealed[-2], self._sealed[-1] + keys = mx.concatenate( + [_rows_view(older.segment.keys, older.n), _rows_view(newer.segment.keys, newer.n)], axis=2 + ) + values = mx.concatenate( + [_rows_view(older.segment.values, older.n), _rows_view(newer.segment.values, newer.n)], axis=2 + ) + mx.eval(keys, values) + segment = KVSegment(keys, values, sealed=True) + segment.hold(self) + for ref in (older, newer): + ref.segment.release(self) + self._sealed[-2:] = [SegRef(segment, int(older.n) + int(newer.n))] + self.merges += 1 + self.merge_rows += int(older.n) + int(newer.n) + + def stats(self) -> dict[str, int]: + return { + "segments": self.segment_count, + "sealed": len(self._sealed), + "tail_rows": self._tail_n, + "rows": self.offset, + "bytes": self.nbytes, + "seal_copies": self.seal_copies, + "seal_copy_rows": self.seal_copy_rows, + "merges": self.merges, + "merge_rows": self.merge_rows, + } + + +class GatheredKVView: + """The contiguous (keys, values, offset) of a segmented cache, for the stock route ladder. + + The attention ladder in ``attention_split`` reads ``cache.keys``/``values``/``offset``; + when no segment route takes a call it runs on this view of one gathered pair, so the + fallback is the stock route on the same numbers. + """ + + def __init__(self, keys: mx.array, values: mx.array, offset: int) -> None: + self.keys = keys + self.values = values + self.offset = int(offset) + + def make_mask(self, *args, **kwargs): + from mlx_lm.models.cache import create_attention_mask + + return create_attention_mask(*args, offset=self.offset, **kwargs) + + +def gathered_view(cache: SegmentedKVCache) -> GatheredKVView: + gathered = cache.gather() + if gathered is None: + raise ValueError("empty segmented KV cache") + return GatheredKVView(gathered[0], gathered[1], cache.offset) + + +# ---- prefill over segments: per segment (out, lse) plus one merge ------------------------ + + +def _segment_sdpa_lse_reference( + queries: mx.array, + keys: mx.array, + values: mx.array, + n: int, + *, + scale: float, + causal: bool, + q_chunk_bytes: int = 256 * 1024 * 1024, +) -> tuple[mx.array, mx.array]: + """Exact attention of ``queries`` over the first ``n`` rows of one segment, with its logsumexp. + + ``queries`` (1, Hq, Q, D); returns (out in the query dtype, lse float32 (1, Hq, Q)). + ``causal=True`` marks the segment holding the Q new rows as its last Q rows (query row i + sees keys up to ``n - Q + i``); otherwise every key is visible. Scores run in float32 and + are bounded by chunking the query rows: this is the slow, exact reference route (QK^T plus + logsumexp; phase 1 estimated about 3 times the contiguous prefill at 80K). An MLX attention + kernel that returns the logsumexp replaces this function (``segment_sdpa_lse``). + """ + b, hq, q_len, d = (int(x) for x in queries.shape) + hk = int(keys.shape[1]) + gqa = hq // hk + k = _rows_view(keys, n).astype(mx.float32) + v = _rows_view(values, n).astype(mx.float32) + per_row = hk * gqa * int(n) * 4 + chunk = max(1, min(q_len, q_chunk_bytes // max(1, per_row))) + outs, lses = [], [] + for start in range(0, q_len, chunk): + stop = min(q_len, start + chunk) + qc = queries[:, :, start:stop, :].astype(mx.float32) + rows = stop - start + qg = qc.reshape(b, hk, gqa * rows, d) + scores = (qg @ k.transpose(0, 1, 3, 2)) * scale + if causal: + row = mx.arange(gqa * rows) % rows + start + limit = n - q_len + row + visible = mx.arange(int(n))[None, :] <= limit[:, None] + scores = mx.where(visible[None, None], scores, -mx.inf) + lse = mx.logsumexp(scores, axis=-1) + p = mx.exp(scores - lse[..., None]) + out = (p @ v).reshape(b, hq, rows, d) + outs.append(out.astype(queries.dtype)) + lses.append(lse.reshape(b, hq, rows)) + if len(outs) == 1: + return outs[0], lses[0] + return mx.concatenate(outs, axis=2), mx.concatenate(lses, axis=2) + + +def _segment_sdpa_lse_fast(queries, keys, values, n, *, scale, causal): + """Per-segment attention on the MLX kernel that returns the logsumexp (fork only).""" + k, v = _rows_view(keys, n), _rows_view(values, n) + kwargs = {"mask": "causal"} if causal else {} + out, lse = mx.fast.scaled_dot_product_attention( + queries, k, v, scale=scale, return_lse=True, **kwargs + ) + return out, lse.astype(mx.float32) + + +def segment_sdpa_lse(queries, keys, values, n, *, scale, causal): + """Per segment (out, lse): the MLX kernel when present, else the exact reference route.""" + if sdpa_lse_available(): + _count("sdpa_lse_kernel") + return _segment_sdpa_lse_fast(queries, keys, values, n, scale=scale, causal=causal) + _count("sdpa_lse_reference") + return _segment_sdpa_lse_reference(queries, keys, values, n, scale=scale, causal=causal) + + +def merge_segment_outputs(outs: Sequence[mx.array], lses: Sequence[mx.array]) -> mx.array: + """Exact merge of per-segment (out, lse): sum_s exp(lse_s - L) * out_s with L = logsumexp(lse).""" + if len(outs) == 1: + return outs[0] + stacked = mx.stack([l.astype(mx.float32) for l in lses], axis=0) + total = mx.logsumexp(stacked, axis=0) + merged = None + for out, lse in zip(outs, lses): + weight = mx.exp(lse.astype(mx.float32) - total)[..., None] + term = out.astype(mx.float32) * weight + merged = term if merged is None else merged + term + return merged.astype(outs[0].dtype) + + +def attend_segments_lse(queries: mx.array, cache: SegmentedKVCache, *, scale: float) -> mx.array: + """Attention of the cache's last Q rows (``queries``) over all of its segments, via (out, lse).""" + segments = cache.attention_segments() + outs, lses = [], [] + for index, (keys, values, n) in enumerate(segments): + out, lse = segment_sdpa_lse( + queries, keys, values, n, scale=scale, causal=(index == len(segments) - 1) + ) + outs.append(out) + lses.append(lse) + return merge_segment_outputs(outs, lses) + + +# ---- the decode/verify route used by attention_split ----------------------------------- + + +def decode_segments_attention( + queries: mx.array, + cache: SegmentedKVCache, + *, + scale: float, + mask: Any, + packed_threshold: int, +) -> mx.array | None: + """Verify/decode attention straight over the segments, or None for the gather fallback. + + Engages under the same conditions as the packed route over a dense cache, with the total + length where that route reads the buffer capacity: MTPLX_NAX_FLASH_ROUTE on, a verify + window of 2 to 8 rows, total length at least ``packed_threshold``, a causal or no mask. + The head-dim-split kernel serves windows of up to 5 rows at GQA 6 (wide and dsplit4 are not + in this tree); a longer window, an unsupported shape or a kernel bail returns None and the + call runs on the gathered rows through the stock ladder. + """ + from .kernel_selfcheck import lane_disabled + from .kernels.sdpa_segmented import ( + sdpa_nax_flash_dsplit_segments, + segments_supported, + ) + + if os.environ.get("MTPLX_NAX_FLASH_ROUTE", "").strip().lower() not in {"1", "true", "yes", "on"}: + return None + if lane_disabled("nax_flash_dsplit_sdpa"): + return None + q_len = int(queries.shape[2]) + if not 2 <= q_len <= 8: + return None + if cache.offset < int(packed_threshold): + return None + if not (mask is None or (isinstance(mask, str) and mask == "causal")): + return None + hq, d = int(queries.shape[1]), int(queries.shape[3]) + segments = cache.attention_segments() + if not segments: + return None + hk = int(segments[0][0].shape[1]) + if hk <= 0 or hq % hk or not segments_supported(q_len, hq // hk, d): + return None + return sdpa_nax_flash_dsplit_segments(queries=queries, segments=segments, scale=scale) + + +# ---- install -------------------------------------------------------------------------------- + + +def is_segmented(entry: Any) -> bool: + return isinstance(entry, SegmentedKVCache) + + +def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = None) -> dict[str, int | str]: + """Replace the empty stock full-attention KV caches of a fresh cache list by segmented ones. + + Only plain dense ``KVCache`` entries that hold nothing yet are replaced (a prefill or a + restore fills them afterwards); recurrent layers, rotating/indexed caches and anything with + rows are left alone. + """ + stats: dict[str, int | str] = {"enabled": 1, "mode": "segmented", "entries": 0, "skipped": 0} + from .cache_state import _is_trimmable + + for idx, entry in enumerate(cache or []): + if entry is None: + stats["skipped"] = int(stats["skipped"]) + 1 + continue + if isinstance(entry, SegmentedKVCache): + stats["entries"] = int(stats["entries"]) + 1 + continue + if ( + not _is_trimmable(entry) + or getattr(entry, "_idx", None) is not None + or getattr(entry, "keys", None) is not None + or int(getattr(entry, "offset", 0) or 0) != 0 + or type(entry).__name__ != "KVCache" + ): + stats["skipped"] = int(stats["skipped"]) + 1 + continue + cache[idx] = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) + stats["entries"] = int(stats["entries"]) + 1 + return stats + + +# ---- accounting ----------------------------------------------------------------------------- + + +def unique_state_bytes(states: Iterable[Any]) -> int: + """Bytes of the segments referenced by ``states`` (SegmentedKVState objects), each counted once.""" + seen: dict[int, int] = {} + for state in states: + if isinstance(state, SegmentedKVState): + for ref in state.refs: + seen.setdefault(ref.segment.id, ref.segment.nbytes) + return sum(seen.values()) + + +def walk_segmented_states(value: Any) -> Iterable[SegmentedKVState]: + """Every SegmentedKVState inside a snapshot tree (tuples, lists, dicts).""" + if isinstance(value, SegmentedKVState): + yield value + elif isinstance(value, (list, tuple)): + for item in value: + yield from walk_segmented_states(item) + elif isinstance(value, dict): + for item in value.values(): + yield from walk_segmented_states(item) diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 30fb37fb0..3af799f70 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -82,6 +82,7 @@ from mtplx.adaptive import AdaptiveDepthPolicy, ExpectedValueDepthPolicy from mtplx.attention_context import attention_phase from mtplx.cache_state import snapshot_cache +from mtplx.segmented_kv import segmented_kv_enabled from mtplx.one_copy import ( anchor_nbytes as one_copy_anchor_nbytes, one_copy_runtime, @@ -21777,6 +21778,14 @@ def _admission_restore_copies_prefix( if entry is None: return True + from mtplx.session_bank import entry_has_segmented_kv + + if entry_has_segmented_kv(entry): + # A segmented snapshot is a list of references to sealed segments: the restored + # cache writes its turn into a new tail, so the history is never copied + # (mtplx/segmented_kv.py; +4.95 GiB at 80K with a stock snapshot, +0.125 GiB with + # segments, phase 1 measurement). + return False if str(restore_mode or "clone") != "reference": return True if getattr(entry, "cache_ref", None) is None: @@ -22094,6 +22103,7 @@ def _admission_growth( restore_fixed_bytes: int | None = None, publish_bytes: int | None = None, slack_rows: int = 0, + segmented_kv: bool = False, ) -> dict[str, Any]: """New memory one request needs at its peak, on top of what is measured. @@ -22282,6 +22292,10 @@ def _admission_growth( publish_copy = 0 elif publish_bytes is not None: publish_copy = max(0, int(publish_bytes)) + elif segmented_kv and not paged_live: + # Segmented KV: the snapshot is a list of references and seals the tail, a copy of + # the rows this request writes (one layer at a time), never of the history. + publish_copy = (M + max(0, int(output_tokens))) * live_w else: publish_copy = live_total decode_start = live_decode + publish_copy @@ -23071,6 +23085,7 @@ def growth( ), publish_bytes=anchor_bytes, slack_rows=slack_rows, + segmented_kv=segmented_kv_enabled(), ) model["scratch_source"] = scratch_source calibration = getattr(runtime, "prefill_scratch_calibration", None) diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py index 61759f0a0..cbc4850d2 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -35,6 +35,7 @@ from .checkpoint_anchors import AnchorPlan, default_tail_backoff, retain_checkpoints from .runtime import MTPLXRuntime from .runtime_options import block_prefix_restore_enabled +from .segmented_kv import SegmentedKVCache, SegmentedKVState, walk_segmented_states def _policy_uses_committed_history(policy: str | None) -> bool: @@ -390,6 +391,10 @@ def _tree_nbytes(value: Any) -> int: return _tree_nbytes(value.states) + _tree_nbytes(value.meta_states) if isinstance(value, mx.array): return int(value.nbytes) + if isinstance(value, SegmentedKVState): + # The segments it references, counted as if it were alone; segments that several + # entries share are taken out once in ``SessionBank.total_nbytes``. + return int(value.nbytes) if isinstance(value, (list, tuple)): return sum(_tree_nbytes(item) for item in value) if isinstance(value, dict): @@ -397,6 +402,34 @@ def _tree_nbytes(value: Any) -> int: return 0 +def _entry_segmented_states(entry: Any) -> list[SegmentedKVState]: + """The segmented KV states a bank entry's cache snapshot holds (empty for a stock snapshot).""" + snapshot = getattr(entry, "cache_snapshot", None) + if snapshot is None: + return [] + return list(walk_segmented_states(snapshot.states)) + + +def entry_has_segmented_kv(entry: Any) -> bool: + return bool(_entry_segmented_states(entry)) + + +def _cache_has_segmented_kv(cache: list[Any] | None) -> bool: + return any(isinstance(layer, SegmentedKVCache) for layer in cache or ()) + + +def _shared_segment_overcount(entries: list[Any]) -> int: + """Bytes counted more than once because several snapshots reference the same segment.""" + naive = 0 + unique: dict[int, int] = {} + for entry in entries: + for state in _entry_segmented_states(entry): + naive += int(state.nbytes) + for ref in state.refs: + unique.setdefault(ref.segment.id, int(ref.segment.nbytes)) + return max(0, naive - sum(unique.values())) + + def _snapshot_nbytes(snapshot: CacheSnapshot) -> int: return _tree_nbytes(snapshot.states) + _tree_nbytes(snapshot.meta_states) @@ -1266,7 +1299,9 @@ def total_nbytes(self) -> int: # "nothing sheddable", shrink_to_bytes never entered its loop, and the # dynamic ceiling (working set = active - weights - this total) took # the lease for working set and evicted the durable snapshots instead. - return sum(entry.held_nbytes for entry in list(self._entries.values())) + entries = list(self._entries.values()) + # A segment that several entries reference is held once (mtplx/segmented_kv.py). + return sum(entry.held_nbytes for entry in entries) - _shared_segment_overcount(entries) @property def lease_entries(self) -> int: @@ -3115,6 +3150,13 @@ def _enqueue_cold_entry( cold["enabled"] = False cold["skip_reason"] = "no_cold_tier" return + if entry_has_segmented_kv(entry): + # The SSD tier encodes contiguous KV arrays; a segmented snapshot is a list of + # references. Not persisted in this phase (MTPLX_SEGMENTED_KV, mtplx/segmented_kv.py). + if cold is not None: + cold["enabled"] = False + cold["skip_reason"] = "segmented_kv" + return put_entry = getattr(self.cold_tier, "put_entry", None) if not callable(put_entry): if cold is not None: @@ -3280,6 +3322,8 @@ def _schedule_live_ref_spill(self, entry: SessionBankEntry) -> None: cold = self.cold_tier if cold is None or not callable(getattr(cold, "spill_entry", None)): return + if _cache_has_segmented_kv(entry.cache_ref): + return if self.cold_enqueue_dispatch is None: self.eviction_log.append( { @@ -3342,6 +3386,8 @@ def run_live_ref_spill( spill = getattr(cold, "spill_entry", None) if cold is not None else None if not callable(spill): return False + if _cache_has_segmented_kv(entry.cache_ref): + return False try: snapshot = snapshot_cache_lazy_hybrid(entry.cache_ref) # Either form of the draft history goes to disk: without it a diff --git a/tests/test_sdpa_segmented.py b/tests/test_sdpa_segmented.py new file mode 100644 index 000000000..c42d7bd6a --- /dev/null +++ b/tests/test_sdpa_segmented.py @@ -0,0 +1,132 @@ +"""Segmented-KV decode/verify kernel (GPU, TensorOps required): exactness against the contiguous +kernel and an fp32 reference, at the real Qwen3.8-27B head geometry.""" + +from __future__ import annotations + +import math + +import mlx.core as mx +import pytest + +from mtplx.kernels import sdpa_segmented as S +from mtplx.kernels.sdpa_nax_flash_dsplit import sdpa_nax_flash_dsplit +from mtplx.nax_verify import nax_available + +HQ, HK, D = 24, 4, 256 +SCALE = 1.0 / math.sqrt(D) + +pytestmark = pytest.mark.skipif( + not mx.metal.is_available() or not nax_available(), + reason="TensorOps kernel unavailable on this machine", +) + + +def _data(total, q_len, peaked=False, seed=3): + mx.random.seed(seed) + qs = 2.5 if peaked else 0.5 + q = (mx.random.normal((1, HQ, q_len, D)) * qs).astype(mx.bfloat16) + k = (mx.random.normal((1, HK, total + 64, D)) * 0.5).astype(mx.bfloat16) + v = (mx.random.normal((1, HK, total + 64, D)) * 0.5).astype(mx.bfloat16) + mx.eval(q, k, v) + return q, k, v + + +def _split(k, v, lens, spare=0): + segs, a = [], 0 + for n in lens: + kk = mx.contiguous(k[:, :, a : a + n + spare]) + vv = mx.contiguous(v[:, :, a : a + n + spare]) + mx.eval(kk, vv) + segs.append((kk, vv, n)) + a += n + return segs + + +def _lens(total, nseg): + lens = [total // nseg] * nseg + lens[-1] += total - sum(lens) + return lens + + +def _ref(q, k, v, total): + q_len = q.shape[2] + kf = mx.repeat(k[:, :, :total].astype(mx.float32), HQ // HK, axis=1) + vf = mx.repeat(v[:, :, :total].astype(mx.float32), HQ // HK, axis=1) + s = (q.astype(mx.float32) @ kf.transpose(0, 1, 3, 2)) * SCALE + vis = mx.arange(total)[None, :] <= (total - q_len + mx.arange(q_len))[:, None] + s = mx.where(vis[None, None], s, -mx.inf) + return mx.softmax(s, axis=-1) @ vf + + +def _ulps(a, b): + """Max |a - b| in bf16 ulps of the output's largest magnitude (small elements carry the + cancellation error of their terms, so a per-element ulp would be meaningless).""" + a, b = a.astype(mx.float32), b.astype(mx.float32) + top = float(mx.maximum(mx.abs(a).max(), mx.abs(b).max()).item()) + ulp = 2.0 ** (math.floor(math.log2(top)) - 7) + return float(mx.abs(a - b).max().item()) / ulp + + +@pytest.mark.parametrize("q_len", [1, 2, 4, 5]) +def test_one_segment_is_bit_identical_to_the_contiguous_kernel(q_len) -> None: + total = 9000 + q, k, v = _data(total, q_len) + ref = sdpa_nax_flash_dsplit(queries=q, keys=k, values=v, offset=total, scale=SCALE) + out = S.sdpa_nax_flash_dsplit_segments(queries=q, segments=[(k, v, total)], scale=SCALE) + assert mx.array_equal(out, ref).item() + + +@pytest.mark.parametrize("nseg", [2, 3, 6, 13, 30]) +@pytest.mark.parametrize("q_len", [1, 4, 5]) +@pytest.mark.parametrize("peaked", [False, True]) +def test_segments_match_contiguous_within_bf16_ulps_and_the_fp32_reference(nseg, q_len, peaked) -> None: + total = 12000 + q, k, v = _data(total, q_len, peaked) + contiguous = sdpa_nax_flash_dsplit(queries=q, keys=k, values=v, offset=total, scale=SCALE) + out = S.sdpa_nax_flash_dsplit_segments( + queries=q, segments=_split(k, v, _lens(total, nseg)), scale=SCALE + ) + assert out is not None and bool(mx.all(mx.isfinite(out)).item()) + assert _ulps(out, contiguous) <= 3 + ref = _ref(q, k, v, total) + err = float(mx.abs(out.astype(mx.float32) - ref).max()) + err_contig = float(mx.abs(contiguous.astype(mx.float32) - ref).max()) + assert err <= max(2 * err_contig, 1e-3) + # argmax over a fixed random projection agrees (the logits' argmax in miniature) + proj = mx.random.normal((HQ * D, 512), key=mx.random.key(1)) + a = mx.argmax(out.astype(mx.float32).transpose(0, 2, 1, 3).reshape(q_len, -1) @ proj, axis=-1) + b = mx.argmax(contiguous.astype(mx.float32).transpose(0, 2, 1, 3).reshape(q_len, -1) @ proj, axis=-1) + assert mx.array_equal(a, b).item() + + +def test_a_segment_reference_shorter_than_its_buffer_reads_only_its_rows() -> None: + """A fork inside a segment: (buffer, n') with live rows past n' that must stay invisible.""" + total, q_len = 9000, 4 + q, k, v = _data(total, q_len) + contiguous = sdpa_nax_flash_dsplit(queries=q, keys=k, values=v, offset=total, scale=SCALE) + lens = _lens(total, 3) + segs = _split(k, v, lens, spare=0) + # make the first segment's buffer longer than its reference, with junk past n' + n0 = lens[0] + big_k = mx.concatenate([segs[0][0], mx.ones((1, HK, 300, D), mx.bfloat16) * 9], axis=2) + big_v = mx.concatenate([segs[0][1], mx.ones((1, HK, 300, D), mx.bfloat16) * 9], axis=2) + segs[0] = (big_k, big_v, n0) + out = S.sdpa_nax_flash_dsplit_segments(queries=q, segments=segs, scale=SCALE) + assert _ulps(out, contiguous) <= 3 + + +def test_only_the_last_segment_is_tail_causal() -> None: + """An old segment is fully visible to every query row, the last one keeps the mask.""" + total, q_len = 9000, 5 + q, k, v = _data(total, q_len, peaked=True, seed=11) + ref = _ref(q, k, v, total) + out = S.sdpa_nax_flash_dsplit_segments(queries=q, segments=_split(k, v, _lens(total, 4)), scale=SCALE) + assert float(mx.abs(out.astype(mx.float32) - ref).max()) < 5e-3 + + +def test_unsupported_shapes_bail_without_a_kernel_launch() -> None: + q, k, v = _data(9000, 6) # 6 rows x GQA 6 = 36 > 32 + before = dict(S.segmented_bail_counts) + assert S.sdpa_nax_flash_dsplit_segments(queries=q, segments=_split(k, v, _lens(9000, 2)), scale=SCALE) is None + assert S.segmented_bail_counts.get("m_rows_gt_32", 0) == before.get("m_rows_gt_32", 0) + 1 + assert S.segments_supported(5, 6, 256) and not S.segments_supported(6, 6, 256) diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py new file mode 100644 index 000000000..efeeef269 --- /dev/null +++ b/tests/test_segmented_kv.py @@ -0,0 +1,450 @@ +"""Segmented KV cache (MTPLX_SEGMENTED_KV): data structure, rollback, seal, merge, snapshots. + +Small shapes on the GPU; the kernel and the attention route have their own test files +(test_sdpa_segmented.py, test_segmented_kv_attention.py) and the session bank flow is in +test_segmented_kv_bank.py. +""" + +from __future__ import annotations + +import gc + +import mlx.core as mx +import pytest +from mlx_lm.models.cache import KVCache + +from mtplx.cache_state import ( + configure_tail_owned_attention_kv_cache, + restore_cache, + snapshot_cache, + snapshot_cache_lazy_hybrid, +) +from mtplx.segmented_kv import ( + KVSegment, + SegmentedKVCache, + SegmentedKVState, + SegRef, + attend_segments_lse, + install_segmented_attention_kv_cache, + merge_segment_outputs, + segment_sdpa_lse, + segmented_kv_enabled, + unique_state_bytes, +) + +HEADS = 2 +DIM = 8 + + +@pytest.fixture(autouse=True) +def _no_merge_by_default(monkeypatch): + # Merge-policy tests opt back in; the others count segments. + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") + + +def _rows(seed: int, n: int) -> tuple[mx.array, mx.array]: + k = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed)).astype(mx.bfloat16) + v = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed + 50_000)).astype(mx.bfloat16) + return k, v + + +def _fill(cache, chunks: list[tuple[int, int]]) -> None: + for seed, n in chunks: + cache.update_and_fetch(*_rows(seed, n)) + + +def _same(cache: SegmentedKVCache, control: KVCache) -> bool: + k, v = cache.gather() + ck, cv = control.state + return bool(mx.array_equal(k, ck).item() and mx.array_equal(v, cv).item()) + + +def test_switch_defaults_off_and_is_read_per_call(monkeypatch) -> None: + monkeypatch.delenv("MTPLX_SEGMENTED_KV", raising=False) + assert segmented_kv_enabled() is False + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + assert segmented_kv_enabled() is True + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "0") + assert segmented_kv_enabled() is False + + +def test_switch_off_leaves_the_stock_caches(monkeypatch) -> None: + monkeypatch.delenv("MTPLX_SEGMENTED_KV", raising=False) + monkeypatch.delenv("MTPLX_VLLM_METAL_PAGED_ATTN", raising=False) + monkeypatch.delenv("MTPLX_OWNED_ATTN_KV", raising=False) + cache = [KVCache(), KVCache()] + configure_tail_owned_attention_kv_cache(cache) + assert all(type(entry) is KVCache for entry in cache) + + +def test_switch_on_replaces_only_empty_plain_kv_caches(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + monkeypatch.delenv("MTPLX_VLLM_METAL_PAGED_ATTN", raising=False) + filled = KVCache() + filled.update_and_fetch(*_rows(1, 4)) + + class Recurrent: + def is_trimmable(self) -> bool: + return False + + cache = [KVCache(), Recurrent(), filled] + stats = configure_tail_owned_attention_kv_cache(cache) + assert isinstance(cache[0], SegmentedKVCache) + assert isinstance(cache[1], Recurrent) + assert cache[2] is filled + assert stats["entries"] == 1 and stats["skipped"] == 2 + + +def test_single_segment_matches_the_stock_cache_bit_for_bit() -> None: + cache, control = SegmentedKVCache(), KVCache() + for seed, n in [(1, 300), (2, 1), (3, 7), (4, 600)]: + got = cache.update_and_fetch(*_rows(seed, n)) + want = control.update_and_fetch(*_rows(seed, n)) + assert cache.offset == control.offset + assert mx.array_equal(got[0], want[0]).item() and mx.array_equal(got[1], want[1]).item() + assert cache.segment_count == 1 + + +def test_multiple_segments_gather_equals_the_stock_contents() -> None: + cache, control = SegmentedKVCache(), KVCache() + for turn, (seed, n) in enumerate([(1, 300), (2, 5), (3, 700), (4, 9)]): + for c in (cache, control): + c.update_and_fetch(*_rows(seed, n)) + cache.seal() # end of turn + assert cache.segment_count == 4 + assert cache.offset == control.offset == 1014 + assert _same(cache, control) + + +def test_make_mask_matches_the_stock_cache() -> None: + cache = SegmentedKVCache() + cache.update_and_fetch(*_rows(1, 12)) + control = KVCache() + control.update_and_fetch(*_rows(1, 12)) + kw = {"return_array": False, "window_size": None} + assert cache.make_mask(4, **kw) == control.make_mask(4, **kw) == "causal" + assert cache.make_mask(1, **kw) is None + + +def test_verify_rollback_stays_in_the_tail_and_never_touches_sealed_segments() -> None: + cache = SegmentedKVCache() + _fill(cache, [(1, 400)]) + cache.seal() + sealed = cache._sealed[0].segment + before = (sealed.keys, sealed.values) + _fill(cache, [(2, 10)]) + cache.update_and_fetch(*_rows(3, 4)) # a verify window + assert cache.offset == 414 + assert cache.trim(3) == 3 # accept one token, roll back three + assert cache.offset == 411 and cache.tail_rows == 11 + cache.update_and_fetch(*_rows(4, 4)) + assert cache.offset == 415 + assert cache._sealed[0].segment is sealed + assert sealed.keys is before[0] and sealed.values is before[1] + assert cache.segment_count == 2 + # the control cache saw the same history + control = KVCache() + for seed, n in [(1, 400), (2, 10), (3, 4)]: + control.update_and_fetch(*_rows(seed, n)) + control.trim(3) + control.update_and_fetch(*_rows(4, 4)) + assert _same(cache, control) + + +def test_trim_across_sealed_segments_shortens_references_without_a_copy() -> None: + cache = SegmentedKVCache() + for seed, n in [(1, 100), (2, 50), (3, 20)]: + _fill(cache, [(seed, n)]) + cache.seal() + keys_before = [ref.segment.keys for ref in cache._sealed] + assert cache.trim(60) == 60 # drops the 20-row segment, shortens the 50-row one + assert cache.offset == 110 + assert [ref.n for ref in cache._sealed] == [100, 10] + assert cache._sealed[1].segment.keys is keys_before[1] # no copy + assert cache.trim(10_000) == 110 + assert cache.offset == 0 and cache.empty() + + +def test_seal_copies_only_the_tail_and_gives_it_its_exact_capacity() -> None: + cache = SegmentedKVCache() + _fill(cache, [(1, 300)]) + assert cache._tail.capacity > 300 # growth slack + assert cache.seal() is True + assert cache.seal() is False # nothing left to seal + ref = cache._sealed[0] + assert ref.segment.sealed and ref.segment.capacity == ref.n == 300 + assert cache.seal_copies == 1 and cache.seal_copy_rows == 300 + # the history is not copied by a later seal + _fill(cache, [(2, 40)]) + history = cache._sealed[0].segment.keys + cache.seal() + assert cache._sealed[0].segment.keys is history + assert cache.seal_copy_rows == 340 + + +def test_state_is_a_list_of_references_and_a_snapshot_never_aliases_the_tail() -> None: + cache = SegmentedKVCache() + _fill(cache, [(1, 300)]) + state = cache.state + assert isinstance(state, SegmentedKVState) and state.rows == 300 + snapshot_k = state.refs[0].segment.keys + # a later turn writes only into a new tail + _fill(cache, [(2, 50)]) + assert state.rows == 300 and state.refs[0].segment.keys is snapshot_k + control = KVCache() + control.update_and_fetch(*_rows(1, 300)) + assert mx.array_equal(state.to_arrays()[0], control.state[0]).item() + + +def test_restore_from_a_state_then_extend_equals_a_never_snapshotted_cache() -> None: + live = SegmentedKVCache() + control = KVCache() + for c in (live, control): + c.update_and_fetch(*_rows(1, 500)) + snapshot = snapshot_cache_lazy_hybrid([live]) + for seed in (10, 11): # the live cache diverges after the snapshot + live.update_and_fetch(*_rows(seed, 7)) + restored = SegmentedKVCache() + restore_cache([restored], snapshot, clone_states=False) + assert restored.offset == 500 + for c in (restored, control): + c.update_and_fetch(*_rows(20, 30)) + assert _same(restored, control) + # the snapshot still serves the original rows + again = SegmentedKVCache() + restore_cache([again], snapshot) + ref = KVCache() + ref.update_and_fetch(*_rows(1, 500)) + assert again.offset == 500 and _same(again, ref) + + +def test_fork_inside_a_segment_is_a_reference_without_a_copy() -> None: + owner = SegmentedKVCache() + _fill(owner, [(1, 400)]) + state = owner.state + segment = state.refs[0].segment + fork = SegmentedKVCache() + fork.state = state + fork.trim(150) # return into the middle of the segment + assert fork.offset == 250 + assert fork._sealed[0].segment is segment and fork._sealed[0].n == 250 + fork.update_and_fetch(*_rows(7, 20)) # divergent continuation + control = KVCache() + control.update_and_fetch(*_rows(1, 400)) + control.trim(150) + control.update_and_fetch(*_rows(7, 20)) + assert _same(fork, control) + # the owner's rows are untouched + ref = KVCache() + ref.update_and_fetch(*_rows(1, 400)) + assert _same(owner, ref) + assert segment.refcount >= 3 # owner, state, fork + + +def test_state_setter_accepts_a_contiguous_pair() -> None: + k, v = _rows(5, 64) + cache = SegmentedKVCache() + cache.state = (k, v) + assert cache.offset == 64 and cache.segment_count == 1 + cache.update_and_fetch(*_rows(6, 3)) + assert cache.offset == 67 + cache.state = None + assert cache.offset == 0 + + +def test_refcount_follows_the_holders_and_eviction_frees_the_segment() -> None: + cache = SegmentedKVCache() + _fill(cache, [(1, 300)]) + state = cache.state + segment = state.refs[0].segment + assert segment.refcount == 2 # the cache and the snapshot + other = SegmentedKVState(state.refs) + assert segment.refcount == 3 + del other + gc.collect() + assert segment.refcount == 2 + cache.state = None + assert segment.refcount == 1 + import weakref + + probe = weakref.ref(segment) + del segment, state + gc.collect() + assert probe() is None # nobody holds it: the buffers are free + + +def test_tiered_merge_only_merges_small_recent_unshared_segments(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "16384") + cache = SegmentedKVCache() + _fill(cache, [(1, 4096)]) + cache.seal() # the big history + for seed in (2, 3, 4, 5): + _fill(cache, [(seed, 300)]) # starting a tail runs compact() + cache.seal() + # history stays alone; the small segments merged geometrically + rows = [ref.n for ref in cache._sealed] + assert rows[0] == 4096 + assert sum(rows) == 4096 + 1200 + assert len(rows) < 5 + control = KVCache() + for seed, n in [(1, 4096), (2, 300), (3, 300), (4, 300), (5, 300)]: + control.update_and_fetch(*_rows(seed, n)) + assert _same(cache, control) + + +def test_merge_leaves_segments_that_a_snapshot_holds(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "16384") + cache = SegmentedKVCache() + for seed in (1, 2): + _fill(cache, [(seed, 300)]) + cache.seal() + held = cache.state # a bank entry references both segments + before = [ref.segment for ref in cache._sealed] + _fill(cache, [(3, 300)]) # new tail: compact() must not merge held segments + assert [ref.segment for ref in cache._sealed] == before + assert held.rows == 600 + + +def test_nbytes_counts_each_shared_segment_once() -> None: + cache = SegmentedKVCache() + _fill(cache, [(1, 256)]) + a = cache.state + b = cache.state + assert unique_state_bytes([a, b]) == a.nbytes + _fill(cache, [(2, 256)]) + c = cache.state + assert unique_state_bytes([a, b, c]) < a.nbytes + b.nbytes + c.nbytes + assert unique_state_bytes([c]) == c.nbytes == 2 * a.nbytes + + +def test_eager_snapshot_of_a_segmented_cache_costs_no_copy() -> None: + cache = SegmentedKVCache() + _fill(cache, [(1, 256)]) + snapshot = snapshot_cache([cache]) + assert isinstance(snapshot.states[0], SegmentedKVState) + assert snapshot.states[0].refs[0].segment.keys is cache._sealed[0].segment.keys + + +def test_install_skips_a_cache_with_rows_and_keeps_existing_segmented() -> None: + seg = SegmentedKVCache() + cache = [seg, KVCache()] + stats = install_segmented_attention_kv_cache(cache) + assert cache[0] is seg and isinstance(cache[1], SegmentedKVCache) + assert stats["entries"] == 2 + + +def test_lse_merge_equals_attention_over_the_concatenated_rows() -> None: + """The prefill interface: per segment (out, lse) plus one merge is exact (fp32 reference).""" + heads_q, dim, q_len = 4, 16, 6 + cache = SegmentedKVCache() + pieces = [] + for seed, n in [(1, 200), (2, 70), (3, 40)]: + k = mx.random.normal((1, HEADS, n, dim), key=mx.random.key(seed)).astype(mx.bfloat16) + v = mx.random.normal((1, HEADS, n, dim), key=mx.random.key(seed + 99)).astype(mx.bfloat16) + cache.update_and_fetch(k, v) + cache.seal() + pieces.append((k, v)) + # the last q_len rows of the cache are the new tokens: put a tail with them + k_new = mx.random.normal((1, HEADS, q_len, dim), key=mx.random.key(7)).astype(mx.bfloat16) + v_new = mx.random.normal((1, HEADS, q_len, dim), key=mx.random.key(8)).astype(mx.bfloat16) + cache.update_and_fetch(k_new, v_new) + q = mx.random.normal((1, heads_q, q_len, dim), key=mx.random.key(9)).astype(mx.bfloat16) + scale = dim ** -0.5 + got = attend_segments_lse(q, cache, scale=scale) + keys, values = cache.gather() + want = mx.fast.scaled_dot_product_attention(q, keys, values, scale=scale, mask="causal") + diff = mx.abs(got.astype(mx.float32) - want.astype(mx.float32)) + assert float(diff.max()) < 0.05 # bf16 output rounding of both routes + # and exactly equal to an fp32 reference + kf = mx.repeat(keys.astype(mx.float32), heads_q // HEADS, axis=1) + vf = mx.repeat(values.astype(mx.float32), heads_q // HEADS, axis=1) + s = (q.astype(mx.float32) @ kf.transpose(0, 1, 3, 2)) * scale + total = keys.shape[2] + vis = mx.arange(total)[None, :] <= (total - q_len + mx.arange(q_len))[:, None] + s = mx.where(vis[None, None], s, -mx.inf) + ref = mx.softmax(s, axis=-1) @ vf + assert float(mx.abs(got.astype(mx.float32) - ref).max()) < 0.02 + + +def test_merge_of_one_segment_is_the_segment_output() -> None: + out = mx.ones((1, 2, 3, 4), dtype=mx.bfloat16) + lse = mx.zeros((1, 2, 3)) + assert merge_segment_outputs([out], [lse]) is out + + +def test_segment_sdpa_lse_is_replaceable() -> None: + import mtplx.segmented_kv as module + + assert module.segment_sdpa_lse is not None + assert callable(segment_sdpa_lse) + + +def test_kvsegment_and_segref_are_plain_references() -> None: + k, v = _rows(1, 4) + segment = KVSegment(k, v, sealed=True) + ref = SegRef(segment, 3) + assert ref.n == 3 and ref.segment.capacity == 4 + assert segment.nbytes == k.nbytes + v.nbytes + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) + + +def _lse_cache(q_len=4, dim=16): + cache = SegmentedKVCache() + for seed, n in [(1, 120), (2, 60)]: + cache.update_and_fetch(*(mx.random.normal((1, HEADS, n, dim), key=mx.random.key(seed + o)).astype(mx.bfloat16) for o in (0, 99))) + cache.seal() + cache.update_and_fetch(*(mx.random.normal((1, HEADS, q_len, dim), key=mx.random.key(7 + o)).astype(mx.bfloat16) for o in (0, 5))) + q = mx.random.normal((1, 4, q_len, dim), key=mx.random.key(9)).astype(mx.bfloat16) + return cache, q + + +def test_stock_mlx_has_no_lse_kernel_and_falls_back_to_the_reference_route(monkeypatch) -> None: + import mtplx.segmented_kv as module + + monkeypatch.setattr(module, "_SDPA_LSE", None) + monkeypatch.delenv("MTPLX_SEGMENTED_KV_PREFILL", raising=False) + assert module.sdpa_lse_available() is False # stock released MLX in the venv + assert module.prefill_route() == "gather" + module.route_counts.clear() + cache, q = _lse_cache() + attend_segments_lse(q, cache, scale=0.25) + assert module.route_counts == {"sdpa_lse_reference": 3} + + +def test_lse_kernel_path_plumbing_with_mocked_detection(monkeypatch) -> None: + import mtplx.segmented_kv as module + + monkeypatch.setattr(module, "_SDPA_LSE", True) + monkeypatch.delenv("MTPLX_SEGMENTED_KV_PREFILL", raising=False) + assert module.prefill_route() == "lse" + calls = [] + + def fake_sdpa(q, k, v, *, scale, return_lse=False, mask=None): + calls.append((int(k.shape[2]), mask)) + out, lse = _segment_sdpa_lse_reference_for_test(q, k, v, scale, mask) + return out, lse + + def _segment_sdpa_lse_reference_for_test(q, k, v, scale, mask): + return module._segment_sdpa_lse_reference(q, k, v, int(k.shape[2]), scale=scale, causal=mask == "causal") + + monkeypatch.setattr(mx.fast, "scaled_dot_product_attention", fake_sdpa) + module.route_counts.clear() + cache, q = _lse_cache() + got = attend_segments_lse(q, cache, scale=0.25) + assert module.route_counts == {"sdpa_lse_kernel": 3} + assert [c[1] for c in calls] == [None, None, "causal"] # only the last segment is causal + monkeypatch.setattr(module, "_SDPA_LSE", False) + want = attend_segments_lse(q, cache, scale=0.25) + assert float(mx.abs(got.astype(mx.float32) - want.astype(mx.float32)).max()) < 1e-2 + + +def test_prefill_route_env_override(monkeypatch) -> None: + import mtplx.segmented_kv as module + + monkeypatch.setattr(module, "_SDPA_LSE", False) + monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "lse") + assert module.prefill_route() == "lse" diff --git a/tests/test_segmented_kv_attention.py b/tests/test_segmented_kv_attention.py new file mode 100644 index 000000000..fe0407fee --- /dev/null +++ b/tests/test_segmented_kv_attention.py @@ -0,0 +1,133 @@ +"""Segmented KV through the real attention layer (attention_split ladder), GPU. + +Qwen3-Next attention at the 27B head geometry (24 query / 4 KV heads, head_dim 256) with a small +hidden size: a stock KVCache and a SegmentedKVCache see the same turns, the same verify windows and +the same rollbacks and must agree. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import mlx.core as mx +import pytest +from mlx.utils import tree_map +from mlx_lm.models.cache import KVCache +from mlx_lm.models.qwen3_next import Qwen3NextAttention + +from mtplx.attention_split import configure_split_full_attention +from mtplx.kernels import sdpa_segmented as S +from mtplx.nax_verify import nax_available +from mtplx.segmented_kv import SegmentedKVCache + +HIDDEN = 128 + + +def _model(): + args = SimpleNamespace( + hidden_size=HIDDEN, num_attention_heads=24, num_key_value_heads=4, head_dim=256, + attention_bias=False, rms_norm_eps=1e-6, partial_rotary_factor=0.25, + rope_theta=10000.0, rope_scaling=None, max_position_embeddings=65536, + ) + mx.random.seed(0) + attn = Qwen3NextAttention(args) + attn.update(tree_map(lambda p: p.astype(mx.bfloat16), attn.parameters())) + mx.eval(attn.parameters()) + layer = SimpleNamespace(is_linear=False, self_attn=attn) + model = SimpleNamespace(model=SimpleNamespace(layers=[layer])) + return model, attn + + +def _x(seed: int, n: int) -> mx.array: + return (mx.random.normal((1, n, HIDDEN), key=mx.random.key(seed)) * 0.5).astype(mx.bfloat16) + + +def _top_ulps(a, b) -> float: + a, b = a.astype(mx.float32), b.astype(mx.float32) + top = float(mx.maximum(mx.abs(a).max(), mx.abs(b).max()).item()) + return float(mx.abs(a - b).max().item()) / 2.0 ** (math.floor(math.log2(top)) - 7) + + +def _run_turns(attn, cache, route_env: bool): + """Turn 1 (two prefill chunks), snapshot at turn end, turn 2 on a restored cache, then verify + windows with rollback. Returns every output.""" + outs = [] + for chunk, (seed, n) in enumerate([(1, 4608), (2, 4608)]): + outs.append(attn(_x(seed, n), mask="causal", cache=cache)) + mx.eval(outs) + if isinstance(cache, SegmentedKVCache): + state = cache.state # the snapshot: seals the tail + cache = SegmentedKVCache() + cache.state = state + outs.append(attn(_x(3, 300), mask="causal", cache=cache)) # turn 2 prefill over history + for seed in (4, 5, 6): + outs.append(attn(_x(seed, 4), mask="causal", cache=cache)) # verify window + cache.trim(2) # accept two, roll back two + outs.append(attn(_x(7, 1), mask=None, cache=cache)) + mx.eval(outs) + return outs, cache + + +@pytest.fixture +def lane(monkeypatch): + monkeypatch.setenv("MTPLX_GQA_PACKED_SDPA", "1") + monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") + monkeypatch.delenv("MTPLX_SEGMENTED_KV_PREFILL", raising=False) + model, attn = _model() + configure_split_full_attention(model) + return attn + + +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +def test_segmented_attention_matches_the_stock_cache_through_turns_and_rollback(lane) -> None: + stock, _ = _run_turns(lane, KVCache(), True) + S.segmented_dispatch_counts.clear() + seg, cache = _run_turns(lane, SegmentedKVCache(), True) + assert len(stock) == len(seg) + for a, b in zip(stock, seg): + assert bool(mx.all(mx.isfinite(b)).item()) + assert _top_ulps(a, b) <= 3 + # the verify windows over history + tail went through the segment kernel + assert S.segmented_dispatch_counts.get("dispatched", 0) >= 3 + assert cache.segment_count == 2 + + +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +def test_single_segment_is_bit_identical_to_the_stock_cache(lane) -> None: + def run(cache): + outs = [lane(_x(1, 9000), mask="causal", cache=cache)] + for seed in (4, 5): + outs.append(lane(_x(seed, 4), mask="causal", cache=cache)) + cache.trim(2) + mx.eval(outs) + return outs + + for a, b in zip(run(KVCache()), run(SegmentedKVCache())): + assert mx.array_equal(a, b).item() + + +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +def test_without_the_route_the_gather_fallback_is_the_stock_result(monkeypatch) -> None: + """No packed lane configured: segments gather into one array and the stock SDPA runs on it.""" + for name in ("MTPLX_GQA_PACKED_SDPA", "MTPLX_NAX_FLASH_ROUTE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + model, attn = _model() + configure_split_full_attention(model) + stock, _ = _run_turns(attn, KVCache(), False) + seg, _ = _run_turns(attn, SegmentedKVCache(), False) + for a, b in zip(stock, seg): + assert mx.array_equal(a, b).item() + + +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +def test_lse_prefill_route_matches_the_gather_route(lane, monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "gather") + a, _ = _run_turns(lane, SegmentedKVCache(), True) + monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "lse") + b, _ = _run_turns(lane, SegmentedKVCache(), True) + for x, y in zip(a, b): + assert _top_ulps(x, y) <= 8 # fp32 reference route vs bf16 SDPA on the same rows diff --git a/tests/test_segmented_kv_bank.py b/tests/test_segmented_kv_bank.py new file mode 100644 index 000000000..cc5242da1 --- /dev/null +++ b/tests/test_segmented_kv_bank.py @@ -0,0 +1,263 @@ +"""Segmented KV in the session bank and the admission gate (MTPLX_SEGMENTED_KV). + +Follow-up turns, a fork, a return to an earlier turn, memory that grows only with new tokens while +snapshots live, eviction, and the admission pricing. The same flow on stock KVCache objects is the +control that shows the duplicate this removes. +""" + +from __future__ import annotations + +import gc +from pathlib import Path + +import mlx.core as mx +from mlx_lm.models.cache import KVCache + +from mtplx.segmented_kv import SegmentedKVCache, SegmentedKVState +from mtplx.session_bank import SessionBank, entry_has_segmented_kv + +LAYERS, HEADS, DIM = 4, 4, 256 +BASE, TURN = 4096, 512 +ROW_BYTES = LAYERS * HEADS * DIM * 2 * 2 # K and V, bf16 + + +class Runtime: + model_path = Path("models/example") + mtp_enabled = False + + def __init__(self, kind: str) -> None: + self.kind = kind + + def make_cache(self): + return [SegmentedKVCache() if self.kind == "segmented" else KVCache() for _ in range(LAYERS)] + + +def _rows(seed: int, n: int): + k = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed)).astype(mx.bfloat16) + v = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed + 7)).astype(mx.bfloat16) + return k, v + + +def _write(cache, seed: int, n: int) -> None: + for layer, entry in enumerate(cache): + entry.update_and_fetch(*_rows(seed * 100 + layer, n)) + mx.eval([mx.array(0)]) + for entry in cache: + k, v = entry.state if not isinstance(entry, SegmentedKVCache) else (entry.keys, entry.values) + mx.eval(k, v) + + +def _contents(cache): + out = [] + for entry in cache: + k, v = (entry.keys, entry.values) if isinstance(entry, SegmentedKVCache) else entry.state + out.append((k[..., : entry.offset, :], v[..., : entry.offset, :])) + return out + + +def _bank(monkeypatch) -> SessionBank: + monkeypatch.setenv("MTPLX_SESSION_LAZY_SNAPSHOT", "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") + return SessionBank(max_entries=64, max_bytes=1 << 40, per_session_max_bytes=1 << 40) + + +def _turn(bank, runtime, tokens: list[int], new: int, seed: int, session: str, *, prev=None): + """One turn: restore the previous prompt, write ``new`` rows, bank the result.""" + if prev is None: + cache = runtime.make_cache() + else: + restored = bank.restore(runtime, prev, session_id=session) + assert restored is not None + cache = restored.cache + assert cache[0].offset == len(prev) + _write(cache, seed, new) + full = list(tokens) + entry = bank.put(runtime=runtime, token_ids=full, cache=cache, logits=None, hidden=None, session_id=session) + assert entry is not None + return cache, entry + + +def _growth(kind: str, monkeypatch, turns: int = 6) -> tuple[int, int, int]: + """Active-memory growth and peak growth over ``turns`` follow-ups, every snapshot kept alive + (the entries are held here); returns (growth, peak growth, bytes of the new tokens).""" + bank = _bank(monkeypatch) + runtime = Runtime(kind) + tokens = list(range(BASE)) + cache, entry = _turn(bank, runtime, tokens, BASE, 1, "s") + held = [entry] + del cache + gc.collect() + mx.clear_cache() + before = mx.get_active_memory() + mx.reset_peak_memory() + prev = tokens + for t in range(turns): + tokens = prev + [BASE * 10 + t * TURN + i for i in range(TURN)] + cache, entry = _turn(bank, runtime, tokens, TURN, 10 + t, "s", prev=prev) + held.append(entry) + prev = tokens + del cache + gc.collect() + mx.clear_cache() + return mx.get_active_memory() - before, mx.get_peak_memory() - before, turns * TURN * ROW_BYTES + + +def test_six_follow_up_turns_grow_memory_only_with_the_new_tokens(monkeypatch) -> None: + growth, peak, new_bytes = _growth("segmented", monkeypatch) + slack = 8 * 1024 * 1024 + assert growth <= new_bytes * 1.35 + slack, (growth, new_bytes) + assert peak <= new_bytes * 1.35 + 2 * TURN * ROW_BYTES + slack, (peak, new_bytes) + + +def test_the_stock_cache_duplicates_the_history_in_the_same_flow(monkeypatch) -> None: + """Control: a stock KVCache bank copies the history on a follow-up while a snapshot aliases it.""" + _g, peak, _new = _growth("stock", monkeypatch, turns=2) + assert peak > 0.9 * BASE * ROW_BYTES # at least one full copy of the history at the peak + + +def test_follow_up_content_equals_a_cache_that_never_went_through_the_bank(monkeypatch) -> None: + bank = _bank(monkeypatch) + runtime = Runtime("segmented") + control = [KVCache() for _ in range(LAYERS)] + tokens = list(range(BASE)) + _, _ = _turn(bank, runtime, tokens, BASE, 1, "s") + _write(control, 1, BASE) + prev = tokens + for t in range(6): + tokens = prev + [90_000 + t * TURN + i for i in range(TURN)] + cache, _ = _turn(bank, runtime, tokens, TURN, 10 + t, "s", prev=prev) + _write(control, 10 + t, TURN) + prev = tokens + for (ck, cv), (kk, kv) in zip(_contents(cache), _contents(control)): + assert mx.array_equal(ck, kk).item() and mx.array_equal(cv, kv).item() + assert cache[0].segment_count == 7 + + +def test_the_bank_entry_is_a_list_of_references_and_counts_shared_segments_once(monkeypatch) -> None: + bank = _bank(monkeypatch) + runtime = Runtime("segmented") + tokens = list(range(BASE)) + _, first = _turn(bank, runtime, tokens, BASE, 1, "s") + assert entry_has_segmented_kv(first) + assert all(isinstance(state, SegmentedKVState) for state in first.cache_snapshot.states) + tokens2 = tokens + list(range(BASE, BASE + TURN)) + _, second = _turn(bank, runtime, tokens2, TURN, 2, "s", prev=tokens) + naive = first.nbytes + second.nbytes + assert bank.total_nbytes < naive # the history segments are shared + assert bank.total_nbytes >= (BASE + TURN) * ROW_BYTES + + +def _borrow(bank, runtime, entry, matched: int): + borrowed = bank.restore_entry_prefix_cache(runtime, entry, matched) + assert borrowed is not None + return borrowed[0] + + +def test_fork_shares_the_history_and_branches_do_not_see_each_other(monkeypatch) -> None: + bank = _bank(monkeypatch) + runtime = Runtime("segmented") + tokens = list(range(BASE)) + cache, root = _turn(bank, runtime, tokens, BASE, 1, "main") + del cache + gc.collect() + mx.clear_cache() + before = mx.get_active_memory() + a = _borrow(bank, runtime, root, BASE) + b = _borrow(bank, runtime, root, BASE) + fork_at = a[0].offset + assert fork_at == b[0].offset and fork_at >= BASE - 8 # a seed slot may be left to re-forward + _write(a, 2, TURN) + _write(b, 3, TURN) + gc.collect() + mx.clear_cache() + # two new tails (512 rows plus the 256-row growth step), no second copy of the 4096-row history + assert mx.get_active_memory() - before <= 2 * (TURN + 256) * ROW_BYTES + 4 * 1024 * 1024 + assert not mx.array_equal(a[0].keys[..., fork_at:, :], b[0].keys[..., fork_at:, :]).item() + assert mx.array_equal(a[0].keys[..., :fork_at, :], b[0].keys[..., :fork_at, :]).item() + assert a[0]._sealed[0].segment is b[0]._sealed[0].segment # one history buffer + reference = [KVCache() for _ in range(LAYERS)] + _write(reference, 1, BASE) + for (ak, _), (rk, _) in zip(_contents(a), _contents(reference)): + assert mx.array_equal(ak[..., :fork_at, :], rk[..., :fork_at, :]).item() + + +def test_return_to_an_earlier_turn_restores_a_prefix_reference_without_a_copy(monkeypatch) -> None: + bank = _bank(monkeypatch) + runtime = Runtime("segmented") + t0 = list(range(BASE)) + _turn(bank, runtime, t0, BASE, 1, "s") + t1 = t0 + [50_000 + i for i in range(TURN)] + cache, _e1 = _turn(bank, runtime, t1, TURN, 2, "s", prev=t0) + t2 = t1 + [60_000 + i for i in range(TURN)] + cache, last = _turn(bank, runtime, t2, TURN, 3, "s", prev=t1) + del cache + gc.collect() + mx.clear_cache() + before = mx.get_active_memory() + # go back to the end of turn 1 (a client edited its last message) + cache = _borrow(bank, runtime, last, len(t1)) + back_at = cache[0].offset + assert len(t1) - 8 <= back_at <= len(t1) + assert cache[0]._sealed[-1].n <= TURN # a (segment, n') reference into turn 2's segment + assert mx.get_active_memory() - before <= 8 * 1024 * 1024 # no copy of the history + reference = [KVCache() for _ in range(LAYERS)] + _write(reference, 1, BASE) + _write(reference, 2, TURN) + for (ck, cv), (rk, rv) in zip(_contents(cache), _contents(reference)): + assert mx.array_equal(ck, rk[..., :back_at, :]).item() + assert mx.array_equal(cv, rv[..., :back_at, :]).item() + + +def test_eviction_releases_the_segments(monkeypatch) -> None: + bank = _bank(monkeypatch) + runtime = Runtime("segmented") + tokens = list(range(BASE)) + cache, entry = _turn(bank, runtime, tokens, BASE, 1, "s") + del cache + gc.collect() + mx.clear_cache() + held = mx.get_active_memory() + segments = [state.refs[0].segment for state in entry.cache_snapshot.states] + assert all(segment.refcount >= 1 for segment in segments) + assert bank.clear(session_id="s") == 1 + del entry, segments + gc.collect() + mx.clear_cache() + assert held - mx.get_active_memory() >= BASE * ROW_BYTES * 0.9 + + +def test_ssd_tier_skips_segmented_entries_and_leaves_stock_ones_alone(monkeypatch) -> None: + class Cold: + puts = 0 + + def put_entry(self, *a, **k): + Cold.puts += 1 + + monkeypatch.setenv("MTPLX_SESSION_LAZY_SNAPSHOT", "1") + bank = SessionBank(cold_tier=Cold()) + timing: dict = {} + for kind in ("segmented", "stock"): + runtime = Runtime(kind) + cache = runtime.make_cache() + _write(cache, 1, 600) + bank.put( + runtime=runtime, token_ids=list(range(600)), cache=cache, logits=None, hidden=None, + session_id=kind, timing_out=timing, + ) + if kind == "segmented": + assert timing["cold_enqueue"]["skip_reason"] == "segmented_kv" + assert Cold.puts == 0 + assert Cold.puts == 1 # the stock entry still goes to the tier + + +def test_admission_prices_no_history_copy_for_a_segmented_entry(monkeypatch) -> None: + from mtplx.server.openai import _admission_restore_copies_prefix + + bank = _bank(monkeypatch) + _, seg_entry = _turn(bank, Runtime("segmented"), list(range(BASE)), BASE, 1, "x") + stock_bank = _bank(monkeypatch) + _, stock_entry = _turn(stock_bank, Runtime("stock"), list(range(BASE)), BASE, 1, "y") + assert _admission_restore_copies_prefix(seg_entry, "clone", "x") is False + assert _admission_restore_copies_prefix(seg_entry, "reference", "x") is False + assert _admission_restore_copies_prefix(stock_entry, "clone", "y") is True From 353f6960314fc1f9572c5f24cf4a7a0eec06bba8 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 07:59:35 +0200 Subject: [PATCH 03/37] Segmented KV only above the compiled-verify boundary (MTPLX_SEGMENTED_KV_MIN_TOKENS) --- CHANGELOG.md | 2 +- mtplx/cache_state.py | 9 +- mtplx/segmented_kv.py | 130 +++++++++++++++++--- tests/test_segmented_kv_boundary.py | 184 ++++++++++++++++++++++++++++ 4 files changed, 308 insertions(+), 17 deletions(-) create mode 100644 tests/test_segmented_kv_boundary.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fc81e93b..3a810431c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ All notable user-facing changes to MTPLX. The format is based on ### 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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Not covered: the compiled verify (segmented caches verify eager), the SSD tier (segmented entries are not persisted), paged and quantized KV. +- **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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Not covered: the compiled verify (segmented caches verify eager), the SSD tier (segmented entries are not persisted), paged and quantized KV. ### Fixed diff --git a/mtplx/cache_state.py b/mtplx/cache_state.py index a859660af..4454237c5 100644 --- a/mtplx/cache_state.py +++ b/mtplx/cache_state.py @@ -5021,7 +5021,14 @@ def restore_cache( restore_meta_state: bool = True, clone_states: bool = True, ) -> None: - for entry, state, meta_state in zip(cache, snapshot.states, snapshot.meta_states): + for idx, (entry, state, meta_state) in enumerate(zip(list(cache), snapshot.states, snapshot.meta_states)): + if state is not None: + from .segmented_kv import adapt_layer_for_restore, segmented_kv_enabled + + if segmented_kv_enabled(): + # MTPLX_SEGMENTED_KV: the layer takes the layout its rows call for (above the + # compiled-verify boundary segments, below it the stock cache). + entry, state = adapt_layer_for_restore(cache, idx, state) if state is not None: install_view = not clone_states and _is_trimmable(entry) _restore_state_preserving_container(entry, state, clone=not install_view) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 06a28cb0a..bce9a4ff6 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -689,16 +689,106 @@ def is_segmented(entry: Any) -> bool: return isinstance(entry, SegmentedKVCache) -def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = None) -> dict[str, int | str]: - """Replace the empty stock full-attention KV caches of a fresh cache list by segmented ones. +def segmented_kv_min_tokens() -> int: + """Context length above which a layer is kept as segments (MTPLX_SEGMENTED_KV). + + The segment cache cannot be promoted to the compiled verify (graphbank: + ``promotion_failure:segmented_kv_cache``), so below the compiled-verify context ceiling the + stock cache stays: production runs compiled there and the switch would only slow it down. + ``MTPLX_SEGMENTED_KV_MIN_TOKENS`` sets the boundary explicitly (0 = always segmented). + Otherwise it is the compiled-verify ceiling itself (``MTPLX_COMPILED_VERIFY_MAX_CONTEXT``, + 32768 in the turbo profile, 6144 without a profile), read per call; with the compiled verify + off the boundary is 0, and with the compiled verify unlimited (ceiling 0) it is infinite, + because a compiled verify that runs at every length leaves the segments nothing to win. + A layer is segmented when its rows are strictly above the boundary, the same comparison the + compiled verify uses to fall back to eager. + """ + raw = os.environ.get("MTPLX_SEGMENTED_KV_MIN_TOKENS", "").strip() + if raw: + try: + return max(0, int(raw)) + except ValueError: + pass + from .graphbank import _compiled_verify_max_context, compiled_verify_mode + + if compiled_verify_mode() == "off": + return 0 + ceiling = _compiled_verify_max_context() + return ceiling if ceiling > 0 else (1 << 62) + + +def _is_stock_plain_kv(entry: Any) -> bool: + return type(entry).__name__ == "KVCache" and getattr(entry, "_idx", None) is None + - Only plain dense ``KVCache`` entries that hold nothing yet are replaced (a prefill or a - restore fills them afterwards); recurrent layers, rotating/indexed caches and anything with - rows are left alone. +def _segmentable(entry: Any) -> bool: + """A stock full-attention layer of the target cache (tagged at install; MTP caches never are).""" + return bool(getattr(entry, "_mtplx_segmentable", False)) and _is_stock_plain_kv(entry) + + +def _segmented_from_stock(entry: Any) -> SegmentedKVCache: + """The stock cache's rows as one sealed segment, without a copy. + + The segment views the stock buffers and is never written again; the stock object is dropped + by the caller, so nothing writes into those buffers any more. """ - stats: dict[str, int | str] = {"enabled": 1, "mode": "segmented", "entries": 0, "skipped": 0} + new = SegmentedKVCache(step=int(getattr(entry, "step", 256))) + rows = int(getattr(entry, "offset", 0) or 0) + if rows and entry.keys is not None: + new.state = (entry.keys[..., :rows, :], entry.values[..., :rows, :]) + return new + + +def adapt_layer_for_restore(cache: list[Any], idx: int, state: Any) -> tuple[Any, Any]: + """Pick the layout of layer ``idx`` for a restore of ``state``; returns (entry, state). + + Above ``segmented_kv_min_tokens()`` the layer is a segment cache (a stock snapshot becomes + one sealed segment, no copy); at or below it the layer is a stock cache (a segmented + snapshot is gathered into contiguous rows), so a conversation that crosses the boundary + changes layout at its next restore. Layers that were never tagged (MTP caches, recurrent + layers) are left alone. The entry is replaced in ``cache`` when its class changes. + """ + entry = cache[idx] + if not (isinstance(entry, SegmentedKVCache) or _segmentable(entry)): + return entry, state + if isinstance(state, SegmentedKVState): + rows = state.rows + elif isinstance(state, (tuple, list)) and len(state) == 2 and state[0] is not None: + rows = int(state[0].shape[2]) + else: + return entry, state + floor = segmented_kv_min_tokens() + want = rows > floor if floor > 0 else isinstance(entry, SegmentedKVCache) + if want and not isinstance(entry, SegmentedKVCache): + entry = SegmentedKVCache(step=int(getattr(entry, "step", 256))) + entry._mtplx_segmentable = True + cache[idx] = entry + elif not want and isinstance(entry, SegmentedKVCache): + from mlx_lm.models.cache import KVCache + + entry = KVCache() + entry._mtplx_segmentable = True + cache[idx] = entry + if isinstance(state, SegmentedKVState) and not isinstance(entry, SegmentedKVCache): + state = state.to_arrays() + return entry, state + + +def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = None) -> dict[str, int | str]: + """Choose the layout of the full-attention KV layers of a target cache list. + + Plain dense ``KVCache`` layers are tagged as segmentable. A layer without rows becomes a + segment cache straight away when the boundary is 0 (``segmented_kv_min_tokens``); otherwise + it stays stock (compiled verify below the boundary) and is converted after its prefill when + its rows are above the boundary (this function runs again from ``repage_target_prefill_cache``) + or at the next restore (``adapt_layer_for_restore``). A stock layer with rows above the + boundary becomes one sealed segment (no copy). Recurrent layers, rotating/indexed caches and + caches of other lists (MTP) are left alone. + """ + stats: dict[str, int | str] = {"enabled": 1, "mode": "segmented", "entries": 0, "skipped": 0, "converted": 0} from .cache_state import _is_trimmable + floor = segmented_kv_min_tokens() for idx, entry in enumerate(cache or []): if entry is None: stats["skipped"] = int(stats["skipped"]) + 1 @@ -706,17 +796,27 @@ def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = if isinstance(entry, SegmentedKVCache): stats["entries"] = int(stats["entries"]) + 1 continue - if ( - not _is_trimmable(entry) - or getattr(entry, "_idx", None) is not None - or getattr(entry, "keys", None) is not None - or int(getattr(entry, "offset", 0) or 0) != 0 - or type(entry).__name__ != "KVCache" - ): + if not _is_trimmable(entry) or not _is_stock_plain_kv(entry): stats["skipped"] = int(stats["skipped"]) + 1 continue - cache[idx] = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) - stats["entries"] = int(stats["entries"]) + 1 + rows = int(getattr(entry, "offset", 0) or 0) + if rows == 0 and getattr(entry, "keys", None) is None: + if floor == 0: + new = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) + new._mtplx_segmentable = True + cache[idx] = new + stats["entries"] = int(stats["entries"]) + 1 + else: + entry._mtplx_segmentable = True + stats["skipped"] = int(stats["skipped"]) + 1 + continue + if getattr(entry, "_mtplx_segmentable", False) and rows > floor > 0: + cache[idx] = _segmented_from_stock(entry) + cache[idx]._mtplx_segmentable = True + stats["entries"] = int(stats["entries"]) + 1 + stats["converted"] = int(stats["converted"]) + 1 + continue + stats["skipped"] = int(stats["skipped"]) + 1 return stats diff --git a/tests/test_segmented_kv_boundary.py b/tests/test_segmented_kv_boundary.py new file mode 100644 index 000000000..ab90b223f --- /dev/null +++ b/tests/test_segmented_kv_boundary.py @@ -0,0 +1,184 @@ +"""Segmented KV only above the compiled-verify boundary (MTPLX_SEGMENTED_KV_MIN_TOKENS).""" + +from __future__ import annotations + +import mlx.core as mx +import pytest +from mlx_lm.models.cache import KVCache + +from mtplx.cache_state import ( + configure_tail_owned_attention_kv_cache, + restore_cache, + snapshot_cache, +) +from mtplx.segmented_kv import ( + SegmentedKVCache, + SegmentedKVState, + segmented_kv_min_tokens, +) + +HEADS = 2 +DIM = 8 + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") + monkeypatch.delenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", raising=False) + monkeypatch.delenv("MTPLX_COMPILED_VERIFY", raising=False) + monkeypatch.delenv("MTPLX_COMPILED_VERIFY_MAX_CONTEXT", raising=False) + + +def _rows(seed: int, n: int): + k = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed)).astype(mx.bfloat16) + v = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed + 50_000)).astype(mx.bfloat16) + return k, v + + +def _same(a, b) -> bool: + ka, va = a + kb, vb = b + return bool(mx.array_equal(ka, kb).item() and mx.array_equal(va, vb).item()) + + +def _target_cache(n_layers: int = 2) -> list: + cache = [KVCache() for _ in range(n_layers)] + configure_tail_owned_attention_kv_cache(cache) + return cache + + +def test_boundary_follows_the_compiled_verify_settings(monkeypatch) -> None: + assert segmented_kv_min_tokens() == 0 # compiled verify off: segments from the first row + monkeypatch.setenv("MTPLX_COMPILED_VERIFY", "1") + assert segmented_kv_min_tokens() == 6144 # compiled verify's own default ceiling + monkeypatch.setenv("MTPLX_COMPILED_VERIFY_MAX_CONTEXT", "32768") + assert segmented_kv_min_tokens() == 32768 + monkeypatch.setenv("MTPLX_COMPILED_VERIFY_MAX_CONTEXT", "0") + assert segmented_kv_min_tokens() > 10**12 # compiled at every length: never segmented + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "123") + assert segmented_kv_min_tokens() == 123 + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "0") + assert segmented_kv_min_tokens() == 0 + + +def test_boundary_zero_installs_segments_at_creation() -> None: + cache = _target_cache() + assert all(isinstance(c, SegmentedKVCache) for c in cache) + + +def test_below_the_boundary_the_cache_stays_stock(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + cache = _target_cache() + assert all(type(c) is KVCache for c in cache) + cache[0].update_and_fetch(*_rows(1, 60)) + configure_tail_owned_attention_kv_cache(cache) # what repage runs after a prefill + assert type(cache[0]) is KVCache + + +def test_a_prefill_above_the_boundary_becomes_one_sealed_segment(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + cache = _target_cache() + control = KVCache() + for seed, n in ((1, 70), (2, 70)): + cache[0].update_and_fetch(*_rows(seed, n)) + control.update_and_fetch(*_rows(seed, n)) + stats_before = type(cache[0]) + configure_tail_owned_attention_kv_cache(cache) + assert stats_before is KVCache and isinstance(cache[0], SegmentedKVCache) + assert type(cache[1]) is KVCache # empty layer: nothing to convert + assert cache[0].offset == 140 and cache[0].segment_count == 1 + assert _same(cache[0].gather(), control.state) + # growth after the conversion goes to a new tail; contents equal the stock cache + cache[0].update_and_fetch(*_rows(3, 33)) + control.update_and_fetch(*_rows(3, 33)) + assert cache[0].segment_count == 2 and _same(cache[0].gather(), control.state) + + +def test_restore_picks_the_layout_by_rows(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + long_src = _target_cache(1) + long_src[0].update_and_fetch(*_rows(1, 150)) + configure_tail_owned_attention_kv_cache(long_src) + assert isinstance(long_src[0], SegmentedKVCache) + long_snap = snapshot_cache(long_src) + assert isinstance(long_snap.states[0], SegmentedKVState) + + short_src = _target_cache(1) + short_src[0].update_and_fetch(*_rows(2, 80)) + short_snap = snapshot_cache(short_src) + assert type(short_src[0]) is KVCache + + # segmented snapshot, long -> segments (a reference, no copy of the rows) + dest = _target_cache(1) + restore_cache(dest, long_snap, clone_states=False) + assert isinstance(dest[0], SegmentedKVCache) + assert dest[0].attention_segments()[0][0] is long_src[0].attention_segments()[0][0] + # stock snapshot, short -> stays stock + dest = _target_cache(1) + restore_cache(dest, short_snap) + assert type(dest[0]) is KVCache and dest[0].offset == 80 + # stock snapshot above the boundary -> one sealed segment + stock_long = _target_cache(1) + stock_long[0].update_and_fetch(*_rows(5, 130)) + dest = _target_cache(1) + restore_cache(dest, snapshot_cache(stock_long)) + assert isinstance(dest[0], SegmentedKVCache) and dest[0].offset == 130 + assert _same(dest[0].gather(), stock_long[0].state) + # segmented snapshot of a short conversation (boundary raised) -> stock rows + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "500") + dest = _target_cache(1) + restore_cache(dest, long_snap) + assert type(dest[0]) is KVCache and dest[0].offset == 150 + assert _same(dest[0].state, long_src[0].gather()) + + +def test_restored_segments_extend_like_a_stock_cache(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + control = KVCache() + src = _target_cache(1) + for seed, n in ((1, 90), (2, 90)): + src[0].update_and_fetch(*_rows(seed, n)) + control.update_and_fetch(*_rows(seed, n)) + snap = snapshot_cache(src) # stock, 180 rows + dest = _target_cache(1) + restore_cache(dest, snap) + assert isinstance(dest[0], SegmentedKVCache) + dest[0].update_and_fetch(*_rows(9, 17)) + control.update_and_fetch(*_rows(9, 17)) + assert _same(dest[0].gather(), control.state) + dest[0].trim(5) + control.trim(5) + assert _same(dest[0].gather(), control.state) + + +def test_untagged_caches_are_never_converted(monkeypatch) -> None: + # An MTP history cache is a list of stock KVCache that never went through the install. + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + mtp = [KVCache()] + mtp[0].update_and_fetch(*_rows(1, 150)) + snap = snapshot_cache(mtp) + dest = [KVCache()] + restore_cache(dest, snap) + assert type(dest[0]) is KVCache + configure_tail_owned_attention_kv_cache(mtp) + assert type(mtp[0]) is KVCache + + +def test_conversion_and_snapshot_do_not_copy_the_history(monkeypatch) -> None: + rows = 4096 + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + cache = _target_cache(1) + cache[0].update_and_fetch(*_rows(1, rows)) + mx.eval(cache[0].keys, cache[0].values) + history = 2 * HEADS * rows * DIM * 2 + mx.synchronize() + before = mx.get_active_memory() + configure_tail_owned_attention_kv_cache(cache) + snap = snapshot_cache(cache) + cache[0].append_rows(*_rows(2, 64)) + mx.eval(cache[0]._tail.keys, cache[0]._tail.values) + grew = mx.get_active_memory() - before + assert isinstance(snap.states[0], SegmentedKVState) + assert grew < history // 4 From b0711b10a8441a9644a0150357da8727f9d1d2a0 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 08:06:03 +0200 Subject: [PATCH 04/37] Segmented prefill: reshape the kernel's (B, Hq, Q, 1) logsumexp for the merge; real-kernel test --- mtplx/segmented_kv.py | 3 ++- tests/test_segmented_kv.py | 19 ++++++++++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index bce9a4ff6..b297e975b 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -596,7 +596,8 @@ def _segment_sdpa_lse_fast(queries, keys, values, n, *, scale, causal): out, lse = mx.fast.scaled_dot_product_attention( queries, k, v, scale=scale, return_lse=True, **kwargs ) - return out, lse.astype(mx.float32) + # The kernel's logsumexp is (B, Hq, Q, 1); the merge works on (B, Hq, Q). + return out, lse.astype(mx.float32).reshape(out.shape[:-1]) def segment_sdpa_lse(queries, keys, values, n, *, scale, causal): diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py index efeeef269..9022dbdb6 100644 --- a/tests/test_segmented_kv.py +++ b/tests/test_segmented_kv.py @@ -407,6 +407,9 @@ def test_stock_mlx_has_no_lse_kernel_and_falls_back_to_the_reference_route(monke monkeypatch.setattr(module, "_SDPA_LSE", None) monkeypatch.delenv("MTPLX_SEGMENTED_KV_PREFILL", raising=False) + if module.sdpa_lse_available(): + pytest.skip("this MLX build has return_lse") + monkeypatch.setattr(module, "_SDPA_LSE", None) assert module.sdpa_lse_available() is False # stock released MLX in the venv assert module.prefill_route() == "gather" module.route_counts.clear() @@ -426,7 +429,7 @@ def test_lse_kernel_path_plumbing_with_mocked_detection(monkeypatch) -> None: def fake_sdpa(q, k, v, *, scale, return_lse=False, mask=None): calls.append((int(k.shape[2]), mask)) out, lse = _segment_sdpa_lse_reference_for_test(q, k, v, scale, mask) - return out, lse + return out, lse[..., None] # the MLX kernel returns (B, Hq, Q, 1) def _segment_sdpa_lse_reference_for_test(q, k, v, scale, mask): return module._segment_sdpa_lse_reference(q, k, v, int(k.shape[2]), scale=scale, causal=mask == "causal") @@ -448,3 +451,17 @@ def test_prefill_route_env_override(monkeypatch) -> None: monkeypatch.setattr(module, "_SDPA_LSE", False) monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "lse") assert module.prefill_route() == "lse" + + +def test_real_lse_kernel_matches_the_reference_route() -> None: + import mtplx.segmented_kv as module + + if not module.sdpa_lse_available(): + pytest.skip("MLX build without return_lse") + cache, q = _lse_cache() + module.route_counts.clear() + got = attend_segments_lse(q, cache, scale=0.25) + assert module.route_counts.get("sdpa_lse_kernel") == len(cache.attention_segments()) + keys, values = cache.gather() + want = mx.fast.scaled_dot_product_attention(q, keys, values, scale=0.25, mask="causal") + assert float(mx.abs(got.astype(mx.float32) - want.astype(mx.float32)).max()) < 2e-2 From 9cc33b586a866ce4f14d171a34ec20c1a7c4729d Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 08:38:56 +0200 Subject: [PATCH 05/37] 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 --- mtplx/attention_split.py | 4 +++- mtplx/segmented_kv.py | 36 +++++++++++++++++++++++++--- tests/test_segmented_kv_attention.py | 28 ++++++++++++++++++++++ 3 files changed, 64 insertions(+), 4 deletions(-) diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index 27c0aa5b0..e8c2f248f 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -16,6 +16,7 @@ from .compile_state import in_compiled_step_body, is_compile_trace_error from .segmented_kv import ( SegmentedKVCache, + VERIFY_WINDOW_MAX, attend_segments_lse, decode_segments_attention, gathered_view, @@ -445,7 +446,7 @@ def split_call( # other call runs the ladder below on one gathered view of the rows. cache.append_rows(keys, values) if ( - int(queries.shape[2]) > 8 + 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")) @@ -471,6 +472,7 @@ def split_call( ) segmented_route = "nax_flash_segments" if segmented_output is None: + segmented_count(f"gather_q{int(queries.shape[2]) if int(queries.shape[2]) <= 32 else 'prefill'}") view = gathered_view(cache) keys, values = view.keys, view.values cache = view diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index b297e975b..3ffb844d8 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -639,6 +639,10 @@ def attend_segments_lse(queries: mx.array, cache: SegmentedKVCache, *, scale: fl # ---- the decode/verify route used by attention_split ----------------------------------- +#: Verify windows (rows per call) the segmented decode route serves; bigger calls are prefill chunks. +VERIFY_WINDOW_MAX = 32 + + def decode_segments_attention( queries: mx.array, cache: SegmentedKVCache, @@ -667,7 +671,7 @@ def decode_segments_attention( if lane_disabled("nax_flash_dsplit_sdpa"): return None q_len = int(queries.shape[2]) - if not 2 <= q_len <= 8: + if not 2 <= q_len <= VERIFY_WINDOW_MAX: return None if cache.offset < int(packed_threshold): return None @@ -678,9 +682,35 @@ def decode_segments_attention( if not segments: return None hk = int(segments[0][0].shape[1]) - if hk <= 0 or hq % hk or not segments_supported(q_len, hq // hk, d): + if hk <= 0 or hq % hk: + return None + gqa = hq // hk + if segments_supported(q_len, gqa, d): + out = sdpa_nax_flash_dsplit_segments(queries=queries, segments=segments, scale=scale) + if out is not None: + _count(f"fused_q{q_len}") + return out + # A window the kernel does not take whole (6 to 32 rows without a wide route): sub-windows of + # the rows it does take, each over the keys up to its own last row, so the history is read + # once per sub-window and never gathered. Row i of the window sees keys up to n - Q + i. + sub = max(1, 32 // gqa) + if not segments_supported(sub, gqa, d): + return None + tail_keys, tail_values, tail_n = segments[-1] + if tail_n < q_len: return None - return sdpa_nax_flash_dsplit_segments(queries=queries, segments=segments, scale=scale) + outs = [] + for r0 in range(0, q_len, sub): + r1 = min(q_len, r0 + sub) + trimmed = segments[:-1] + [(tail_keys, tail_values, tail_n - (q_len - r1))] + part = sdpa_nax_flash_dsplit_segments( + queries=queries[:, :, r0:r1, :], segments=trimmed, scale=scale + ) + if part is None: + return None + outs.append(part) + _count(f"chunked_q{q_len}") + return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=2) # ---- install -------------------------------------------------------------------------------- diff --git a/tests/test_segmented_kv_attention.py b/tests/test_segmented_kv_attention.py index fe0407fee..bc3f58b34 100644 --- a/tests/test_segmented_kv_attention.py +++ b/tests/test_segmented_kv_attention.py @@ -131,3 +131,31 @@ def test_lse_prefill_route_matches_the_gather_route(lane, monkeypatch) -> None: b, _ = _run_turns(lane, SegmentedKVCache(), True) for x, y in zip(a, b): assert _top_ulps(x, y) <= 8 # fp32 reference route vs bf16 SDPA on the same rows + + +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +@pytest.mark.parametrize("window", [6, 9, 17, 32]) +def test_wide_verify_windows_run_on_the_segments_without_a_gather(lane, window) -> None: + """Context-copy windows (6 to 32 rows): sub-windows over the segments, history never gathered.""" + import mtplx.segmented_kv as module + + def run(cache): + outs = [lane(_x(1, 9000), mask="causal", cache=cache)] + if isinstance(cache, SegmentedKVCache): + state = cache.state + cache = SegmentedKVCache() + cache.state = state + outs.append(lane(_x(2, 300), mask="causal", cache=cache)) + for seed in (4, 5): + outs.append(lane(_x(seed, window), mask="causal", cache=cache)) + cache.trim(window - 3) + mx.eval(outs) + return outs + + stock = run(KVCache()) + module.route_counts.clear() + seg = run(SegmentedKVCache()) + assert not any(k.startswith("gather_q") and k[8:].isdigit() for k in module.route_counts), module.route_counts + assert module.route_counts.get(f"chunked_q{window}", 0) + module.route_counts.get(f"fused_q{window}", 0) == 2 + for a, b in zip(stock, seg): + assert _top_ulps(a, b) <= 4 From b05703f62f4fad33b4e11374b0d7ccaa8be61958 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 10:22:02 +0200 Subject: [PATCH 06/37] Segmented decode route also serves single-row steps; seal and merge counters --- mtplx/segmented_kv.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 3ffb844d8..6665e11e9 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -409,10 +409,12 @@ def seal(self) -> bool: if rows == 0: self._tail = None return False + _count("seal") if rows == tail.capacity: keys, values = tail.keys, tail.values mx.eval(keys, values) else: + _count("seal_copy") keys = _own_copy(tail.keys[..., :rows, :]) values = _own_copy(tail.values[..., :rows, :]) self.seal_copies += 1 @@ -496,6 +498,7 @@ def _merge_last_two(self) -> None: ref.segment.release(self) self._sealed[-2:] = [SegRef(segment, int(older.n) + int(newer.n))] self.merges += 1 + _count("merge") self.merge_rows += int(older.n) + int(newer.n) def stats(self) -> dict[str, int]: @@ -671,7 +674,7 @@ def decode_segments_attention( if lane_disabled("nax_flash_dsplit_sdpa"): return None q_len = int(queries.shape[2]) - if not 2 <= q_len <= VERIFY_WINDOW_MAX: + if not 1 <= q_len <= VERIFY_WINDOW_MAX: return None if cache.offset < int(packed_threshold): return None From 554ecc25fbae924340ef806e823c9fc7ca75b557 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 10:22:19 +0200 Subject: [PATCH 07/37] Tests: filter the lse route counters --- tests/test_segmented_kv.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py index 9022dbdb6..5def8c281 100644 --- a/tests/test_segmented_kv.py +++ b/tests/test_segmented_kv.py @@ -415,7 +415,7 @@ def test_stock_mlx_has_no_lse_kernel_and_falls_back_to_the_reference_route(monke module.route_counts.clear() cache, q = _lse_cache() attend_segments_lse(q, cache, scale=0.25) - assert module.route_counts == {"sdpa_lse_reference": 3} + assert {k: v for k, v in module.route_counts.items() if k.startswith("sdpa_lse")} == {"sdpa_lse_reference": 3} def test_lse_kernel_path_plumbing_with_mocked_detection(monkeypatch) -> None: @@ -438,7 +438,7 @@ def _segment_sdpa_lse_reference_for_test(q, k, v, scale, mask): module.route_counts.clear() cache, q = _lse_cache() got = attend_segments_lse(q, cache, scale=0.25) - assert module.route_counts == {"sdpa_lse_kernel": 3} + assert {k: v for k, v in module.route_counts.items() if k.startswith("sdpa_lse")} == {"sdpa_lse_kernel": 3} assert [c[1] for c in calls] == [None, None, "causal"] # only the last segment is causal monkeypatch.setattr(module, "_SDPA_LSE", False) want = attend_segments_lse(q, cache, scale=0.25) From 8bb93dc88dcf6aca2814224673450ae5eb36e34c Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 11:44:48 +0200 Subject: [PATCH 08/37] Segmented KV: a segment made from stock rows owns exact row-contiguous storage (MLX copied the strided view on every kernel launch) --- mtplx/segmented_kv.py | 6 +++++- tests/test_segmented_kv_boundary.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 6665e11e9..1e0cdec89 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -445,7 +445,11 @@ def state(self, value: Any) -> None: refs = list(value.refs) elif isinstance(value, (tuple, list)) and len(value) == 2 and value[0] is not None: # A contiguous pair (an SSD restore, a stock snapshot): one sealed segment. - keys, values = value + # Rows of a stock buffer are a strided view of it (capacity beyond the rows). The + # fused kernel takes row-contiguous inputs and MLX would copy such a view on every + # launch (the whole history, per layer, per verify: +34 ms at 50K), so the segment + # owns exact rows (a no-op when the pair is already row-contiguous). + keys, values = (_own_copy(a) for a in value) segment = KVSegment(keys, values, sealed=True) refs = [SegRef(segment, int(keys.shape[2]))] else: diff --git a/tests/test_segmented_kv_boundary.py b/tests/test_segmented_kv_boundary.py index ab90b223f..c47758165 100644 --- a/tests/test_segmented_kv_boundary.py +++ b/tests/test_segmented_kv_boundary.py @@ -182,3 +182,18 @@ def test_conversion_and_snapshot_do_not_copy_the_history(monkeypatch) -> None: grew = mx.get_active_memory() - before assert isinstance(snap.states[0], SegmentedKVState) assert grew < history // 4 + + +def test_a_converted_stock_buffer_is_stored_as_exact_rows(monkeypatch) -> None: + """A stock buffer has capacity beyond its rows; a segment of its rows must not keep a strided view + (MLX copies a non-contiguous kernel input on every launch).""" + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + cache = _target_cache(1) + cache[0].update_and_fetch(*_rows(1, 300)) # stock capacity: 512 rows + assert cache[0].keys.shape[2] > 300 + configure_tail_owned_attention_kv_cache(cache) + segment = cache[0].attention_segments()[0][0] + assert segment.shape[2] == 300 + flat = mx.contiguous(segment) + mx.eval(flat) + assert mx.array_equal(flat, segment).item() From 738eb8a7bdf5430f81fc10359bc4bfdf3ddeb265 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 7 Oct 2026 16:56:36 +0200 Subject: [PATCH 09/37] 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 --- CHANGELOG.md | 2 +- mtplx/cache_bank/codec.py | 40 +++++- mtplx/segmented_kv.py | 67 ++++++++++ mtplx/session_bank.py | 18 ++- tests/test_segmented_kv_bank.py | 3 +- tests/test_segmented_kv_ssd.py | 226 ++++++++++++++++++++++++++++++++ 6 files changed, 347 insertions(+), 9 deletions(-) create mode 100644 tests/test_segmented_kv_ssd.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a810431c..3b340b343 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ All notable user-facing changes to MTPLX. The format is based on ### 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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Not covered: the compiled verify (segmented caches verify eager), the SSD tier (segmented entries are not persisted), paged and quantized KV. +- **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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. Not covered: the compiled verify (segmented caches verify eager), paged and quantized KV. ### Fixed diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index 52503e5bc..c87d4c0e6 100644 --- a/mtplx/cache_bank/codec.py +++ b/mtplx/cache_bank/codec.py @@ -2,6 +2,7 @@ from __future__ import annotations +import sys import time from dataclasses import dataclass, field from pathlib import Path @@ -150,6 +151,16 @@ def content_fingerprints(value: Any, *, rows: int | None = None) -> list[int] | ] +# Segmented KV (MTPLX_SEGMENTED_KV): blocks gathered per fingerprint unit (16 x 256 rows). +_SEGMENTED_FINGERPRINT_BLOCKS = 16 + + +def _is_segmented_state(value: Any) -> bool: + """True for a SegmentedKVState; never imports the segmented module when it is not loaded.""" + module = sys.modules.get("mtplx.segmented_kv") + return module is not None and isinstance(value, module.SegmentedKVState) + + class TreeCodec: """Flatten JSON-safe trees plus MLX arrays into raw tensor blobs.""" @@ -249,6 +260,20 @@ def encode(self, value: Any) -> Any: return {"kind": "str", "value": str(value)} if isinstance(value, mx.array): return self._encode_tensor(value) + if _is_segmented_state(value): + # MTPLX_SEGMENTED_KV: the same spec and blobs as the stock (keys, values) tuple of + # the same rows; the history is sliced block by block, never gathered. + from mtplx.segmented_kv import segmented_state_tensors + + return { + "kind": "tuple", + "items": [ + self._encode_tensor( + t if t.shape[2] >= self.block_size * 2 else t.rows_slice(0, t.shape[2]) + ) + for t in segmented_state_tensors(value) + ], + } if isinstance(value, tuple): return {"kind": "tuple", "items": [self.encode(item) for item in value]} if isinstance(value, list): @@ -339,7 +364,20 @@ def _encode_tensor_blocks( block_fps: list[int] | None = None if self.reuse is not None: full = (shape[axis] // self.block_size) * self.block_size - if full: + if full and hasattr(value, "rows_slice"): + # Segmented rows: fingerprints per chunk of blocks, so the fingerprint pass + # never builds the whole history either. + block_fps = [] + chunk_rows = self.block_size * _SEGMENTED_FINGERPRINT_BLOCKS + for lo in range(0, full, chunk_rows): + part = self._fingerprints( + value.rows_slice(lo, min(full, lo + chunk_rows)), rows=self.block_size + ) + if part is None: + block_fps = None + break + block_fps.extend(part) + elif full: slices = [slice(None)] * len(shape) slices[axis] = slice(0, full) block_fps = self._fingerprints( diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 1e0cdec89..fc784f111 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -858,6 +858,73 @@ def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = return stats +# ---- SSD tier --------------------------------------------------------------------------------- + + +def segmented_ssd_enabled() -> bool: + """Segmented snapshots may go to the SSD tier (MTPLX_SEGMENTED_KV_SSD, default on with the switch).""" + if not segmented_kv_enabled(): + return False + return os.environ.get("MTPLX_SEGMENTED_KV_SSD", "").strip().lower() not in {"0", "false", "no", "off"} + + +class SegmentedRows: + """One tensor of a segmented snapshot (all K or all V rows) as the codec sees an array. + + The SSD codec asks an array for ``shape``, ``dtype``, ``nbytes`` and slices along axis 2. + A slice inside one segment is a view; a slice across a segment edge is a concatenation of + the two small slices, so the history is never gathered. Encoded with the codec's own + 256-row blocks, the result is the spec and the block blobs of the stock cache holding the + same rows (``docs``: SSD tier of the segmented KV cache). + """ + + def __init__(self, refs: Sequence[SegRef], *, values: bool) -> None: + self._refs = tuple(refs) + self._values = bool(values) + first = self._arrays()[0] + self.dtype = first.dtype + self._rows = sum(int(r.n) for r in self._refs) + self.shape = (int(first.shape[0]), int(first.shape[1]), self._rows, int(first.shape[3])) + self.nbytes = int(first.nbytes) // max(1, int(first.shape[2])) * self._rows + self.ndim = 4 + + def _arrays(self) -> list[mx.array]: + return [r.segment.values if self._values else r.segment.keys for r in self._refs] + + def rows_slice(self, start: int, stop: int) -> mx.array: + """Rows [start, stop) as one array (a view when inside one segment).""" + start, stop = max(0, int(start)), min(self._rows, int(stop)) + parts: list[mx.array] = [] + offset = 0 + for ref, array in zip(self._refs, self._arrays()): + lo, hi = offset, offset + int(ref.n) + if hi > start and lo < stop: + parts.append(array[..., max(start, lo) - lo : min(stop, hi) - lo, :]) + offset = hi + if offset >= stop: + break + if not parts: + return mx.zeros((*self.shape[:2], 0, self.shape[3]), self.dtype) + return parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=2) + + def __getitem__(self, index: Any) -> mx.array: + if not isinstance(index, tuple) or len(index) != 4 or not all( + isinstance(i, slice) and i == slice(None) for k, i in enumerate(index) if k != 2 + ): + raise TypeError("SegmentedRows only slices along axis 2") + rows = index[2] + start, stop, step = rows.indices(self._rows) + if step != 1: + raise TypeError("SegmentedRows slices have step 1") + return self.rows_slice(start, stop) + + +def segmented_state_tensors(state: SegmentedKVState) -> tuple[SegmentedRows, SegmentedRows]: + """(keys, values) of a snapshot for the SSD codec; refs dropped to zero rows are skipped.""" + refs = [r for r in state.refs if int(r.n) > 0] + return SegmentedRows(refs, values=False), SegmentedRows(refs, values=True) + + # ---- accounting ----------------------------------------------------------------------------- diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py index cbc4850d2..68ffdce42 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -35,7 +35,12 @@ from .checkpoint_anchors import AnchorPlan, default_tail_backoff, retain_checkpoints from .runtime import MTPLXRuntime from .runtime_options import block_prefix_restore_enabled -from .segmented_kv import SegmentedKVCache, SegmentedKVState, walk_segmented_states +from .segmented_kv import ( + SegmentedKVCache, + SegmentedKVState, + segmented_ssd_enabled, + walk_segmented_states, +) def _policy_uses_committed_history(policy: str | None) -> bool: @@ -3150,9 +3155,10 @@ def _enqueue_cold_entry( cold["enabled"] = False cold["skip_reason"] = "no_cold_tier" return - if entry_has_segmented_kv(entry): - # The SSD tier encodes contiguous KV arrays; a segmented snapshot is a list of - # references. Not persisted in this phase (MTPLX_SEGMENTED_KV, mtplx/segmented_kv.py). + if entry_has_segmented_kv(entry) and not segmented_ssd_enabled(): + # MTPLX_SEGMENTED_KV_SSD=0: a segmented snapshot is a list of references and is not + # persisted. By default the codec encodes it as the stock tensors of the same rows + # (mtplx/cache_bank/codec.py, mtplx/segmented_kv.py:SegmentedRows). if cold is not None: cold["enabled"] = False cold["skip_reason"] = "segmented_kv" @@ -3322,7 +3328,7 @@ def _schedule_live_ref_spill(self, entry: SessionBankEntry) -> None: cold = self.cold_tier if cold is None or not callable(getattr(cold, "spill_entry", None)): return - if _cache_has_segmented_kv(entry.cache_ref): + if _cache_has_segmented_kv(entry.cache_ref) and not segmented_ssd_enabled(): return if self.cold_enqueue_dispatch is None: self.eviction_log.append( @@ -3386,7 +3392,7 @@ def run_live_ref_spill( spill = getattr(cold, "spill_entry", None) if cold is not None else None if not callable(spill): return False - if _cache_has_segmented_kv(entry.cache_ref): + if _cache_has_segmented_kv(entry.cache_ref) and not segmented_ssd_enabled(): return False try: snapshot = snapshot_cache_lazy_hybrid(entry.cache_ref) diff --git a/tests/test_segmented_kv_bank.py b/tests/test_segmented_kv_bank.py index cc5242da1..ca8c4dcc2 100644 --- a/tests/test_segmented_kv_bank.py +++ b/tests/test_segmented_kv_bank.py @@ -227,7 +227,8 @@ def test_eviction_releases_the_segments(monkeypatch) -> None: assert held - mx.get_active_memory() >= BASE * ROW_BYTES * 0.9 -def test_ssd_tier_skips_segmented_entries_and_leaves_stock_ones_alone(monkeypatch) -> None: +def test_ssd_tier_skips_segmented_entries_with_the_ssd_switch_off(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SEGMENTED_KV_SSD", "0") class Cold: puts = 0 diff --git a/tests/test_segmented_kv_ssd.py b/tests/test_segmented_kv_ssd.py new file mode 100644 index 000000000..3cef9bcaf --- /dev/null +++ b/tests/test_segmented_kv_ssd.py @@ -0,0 +1,226 @@ +"""SSD tier for segmented snapshots: the codec writes the stock format of the same rows.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import mlx.core as mx +import numpy as np +import pytest +from mlx_lm.models.cache import KVCache + +from mtplx.cache_bank.codec import decode_payload, encode_payload +from mtplx.cache_bank.cold_tier import SessionBankColdTier +from mtplx.cache_state import CacheSnapshot, restore_cache +from mtplx.segmented_kv import ( + SegmentedKVCache, + SegmentedKVState, + SegmentedRows, + segmented_kv_enabled, +) + +HEADS, DIM = 2, 16 + + +def _rand(seed, n): + rng = np.random.default_rng(seed) + return mx.array(rng.standard_normal((1, HEADS, n, DIM)).astype(np.float32)).astype(mx.bfloat16) + + +def _segmented_and_stock(lens=(700, 311, 90), layers=2): + """Same rows in a SegmentedKVCache with one sealed segment per turn and in a stock cache.""" + seg, stock = [], [] + for layer in range(layers): + s, c = SegmentedKVCache(), KVCache() + for i, n in enumerate(lens): + k, v = _rand(10 * layer + i, n), _rand(100 + 10 * layer + i, n) + s.update_and_fetch(k, v) + c.update_and_fetch(k, v) + if i < len(lens) - 1: + s.seal() + seg.append(s) + stock.append(c) + return seg, stock + + +def _snap(cache): + return CacheSnapshot( + states=tuple(c.state for c in cache) + ((mx.ones((1, 4, 8), dtype=mx.float32),),), + meta_states=tuple(("",) for _ in cache) + (("g",),), + ) + + +def _encode(snapshot, reuse=None): + return encode_payload( + cache_snapshot=snapshot, logits=mx.zeros((1, 8)), hidden=None, mtp_history_snapshot=None, + block_size=256, reuse=reuse, + ) + + +def test_segmented_snapshot_encodes_to_the_stock_spec_and_bytes() -> None: + seg, stock = _segmented_and_stock() + a, b = _snap(seg), _snap(stock) + assert any(isinstance(x, SegmentedKVState) for x in a.states) + ea, eb = _encode(a), _encode(b) + assert json.dumps(ea.spec, sort_keys=True) == json.dumps(eb.spec, sort_keys=True) + assert ea.tensors == eb.tensors + + +def test_incremental_encode_of_a_segmented_snapshot_matches_the_stock_encode() -> None: + seg, stock = _segmented_and_stock() + first = _encode(_snap(stock), reuse={}) + reuse = {key: {"sha256": name, "nbytes": 1} for name, key in first.fingerprints.items()} + ea, eb = _encode(_snap(seg), reuse=dict(reuse)), _encode(_snap(stock), reuse=dict(reuse)) + assert ea.fingerprints == eb.fingerprints + assert set(ea.reused) == set(eb.reused) and ea.reused + assert ea.tensors == eb.tensors + + +def test_a_segmented_entry_written_to_disk_restores_as_the_stock_rows(tmp_path: Path) -> None: + seg, stock = _segmented_and_stock() + snap = _snap(seg) + decoded = decode_payload(_encode(snap).spec, _encode(snap).tensors.__getitem__) + for layer in range(2): + k, v = decoded.cache_snapshot.states[layer] + ck, cv = stock[layer].state + assert mx.array_equal(k, ck).item() and mx.array_equal(v, cv).item() + + +def test_the_codec_never_gathers_the_history() -> None: + seg, _ = _segmented_and_stock(lens=(2000, 40)) + rows = SegmentedRows(seg[0].state.refs, values=False) + mx.clear_cache() + base = mx.get_active_memory() + mx.eval(rows.rows_slice(256, 512)) # inside the first segment: a view, nothing gathered + assert mx.get_active_memory() - base < 100_000 + straddle = rows.rows_slice(1900, 2040) + assert straddle.shape[2] == 140 + + +def test_rows_slice_matches_the_concatenation() -> None: + seg, stock = _segmented_and_stock() + rows = SegmentedRows(seg[0].state.refs, values=True) + full = stock[0].state[1] + for lo, hi in [(0, 1), (0, 256), (690, 710), (255, 1101), (1000, 1101), (0, 1101)]: + assert mx.array_equal(rows[:, :, lo:hi, :], full[:, :, lo:hi, :]).item() + assert rows.shape == tuple(full.shape) and rows.nbytes == full.nbytes + + +class _Entry: + def __init__(self, cache, tokens, session="s1"): + self.token_ids = tuple(tokens) + self.cache_snapshot = _snap(cache) + self.logits = mx.zeros((1, 8)) + self.hidden = mx.ones((1, 4), dtype=mx.bfloat16) + self.mtp_history_snapshot = None + self.gdn_boundaries = () + self.has_recurrent = True + self.session_id = session + self.model_path = "models/example" + self.mtp_enabled = True + self.snapshot_epoch = len(tokens) + self.mtp_snapshot_epoch = len(tokens) + self.nbytes = 4_000_000 + + +def _tier(path: Path, **kw) -> SessionBankColdTier: + return SessionBankColdTier(base_dir=path, mode="on", min_prefix_tokens=2, block_size=256, **kw) + + +def _blobs(tier): + return {p.name: p.read_bytes() for p in (tier.base_dir / "blobs").rglob("*.bin")} + + +def _entry_json(tier, entry): + eid = tier._metadata_for_entry(entry, capabilities=(), payload_nbytes=0)["entry_id"] + payload = json.loads((tier.base_dir / "entries" / eid[:2] / eid / "payload.json").read_text()) + payload["metadata"].pop("created_at_s", None) + return payload + + +@pytest.mark.parametrize("path", ["put", "spill"]) +def test_tier_files_of_a_segmented_entry_equal_those_of_the_stock_entry(tmp_path, path) -> None: + seg, stock = _segmented_and_stock() + toks = range(1101) + out = {} + for name, cache in (("seg", seg), ("stock", stock)): + tier = _tier(tmp_path / name) + entry = _Entry(cache, toks) + if path == "put": + assert tier.put_entry(entry) and tier.flush(timeout_s=20.0) + else: + assert tier.spill_entry(entry) + out[name] = (_entry_json(tier, entry), _blobs(tier)) + assert out["seg"][0] == out["stock"][0] + assert out["seg"][1] == out["stock"][1] and out["seg"][1] + + +def test_forked_entries_share_blobs(tmp_path) -> None: + """Two snapshots that share their first segments store the shared blocks once.""" + seg, _ = _segmented_and_stock(lens=(1024, 512)) + tier = _tier(tmp_path / "t") + first = _Entry(seg, range(1536), session="a") + assert tier.put_entry(first) and tier.flush(timeout_s=20.0) + one = sum(len(b) for b in _blobs(tier).values()) + fork = [] + for layer in seg: + c = SegmentedKVCache() + c.state = SegmentedKVState([layer.state.refs[0]]) # the first segment only + c.update_and_fetch(_rand(900, 256), _rand(901, 256)) + fork.append(c) + second = _Entry(fork, range(1280), session="b") + assert tier.put_entry(second) and tier.flush(timeout_s=20.0) + two = sum(len(b) for b in _blobs(tier).values()) + assert two - one < 0.6 * one # the 1024 shared rows (4 blocks per tensor) are not stored again + + +def test_flag_off_leaves_the_stock_path_untouched(monkeypatch, tmp_path) -> None: + """With the switch off no segmented object exists; a stock entry writes the same files as before the change.""" + monkeypatch.delenv("MTPLX_SEGMENTED_KV", raising=False) + assert not segmented_kv_enabled() + _, stock = _segmented_and_stock() + tier = _tier(tmp_path / "a") + entry = _Entry(stock, range(1101)) + assert tier.put_entry(entry) and tier.flush(timeout_s=20.0) + payload, blobs = _entry_json(tier, entry), _blobs(tier) + ref = encode_payload( + cache_snapshot=entry.cache_snapshot, logits=entry.logits, hidden=entry.hidden, mtp_history_snapshot=None, + has_recurrent=True, block_size=256, + ) + assert sorted(blobs.values()) == sorted(ref.tensors.values()) + assert json.dumps(payload["payload_spec"], sort_keys=True) == json.dumps(ref.spec, sort_keys=True) + # and the segmented branch of the codec is never entered for stock tensors + import mtplx.segmented_kv as module + + def boom(*a, **k): + raise AssertionError("segmented branch taken") + + monkeypatch.setattr(module, "segmented_state_tensors", boom) + tier2 = _tier(tmp_path / "b") + assert tier2.put_entry(entry) and tier2.flush(timeout_s=20.0) + assert _blobs(tier2) == blobs + assert not module.segmented_ssd_enabled() + + +def test_a_stock_entry_on_disk_restores_with_the_switch_on_as_segments(monkeypatch, tmp_path) -> None: + _, stock = _segmented_and_stock() + tier = _tier(tmp_path / "t") + entry = _Entry(stock, range(1101)) + assert tier.put_entry(entry) and tier.flush(timeout_s=20.0) + eid = tier._metadata_for_entry(entry, capabilities=(), payload_nbytes=0)["entry_id"] + payload = json.loads((tier.base_dir / "entries" / eid[:2] / eid / "payload.json").read_text()) + names = payload["tensor_blobs"] + blobs = {n: (tier.base_dir / "blobs" / b["sha256"][:2] / (b["sha256"] + ".bin")).read_bytes() for n, b in names.items()} \ + if (tier.base_dir / "blobs" / next(iter(names.values()))["sha256"][:2]).exists() else None + decoded = decode_payload(payload["payload_spec"], (lambda n: blobs[n]) if blobs else None) + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") + from mtplx.cache_state import configure_tail_owned_attention_kv_cache + cache = [KVCache(), KVCache(), KVCache()] + cache = cache[:2] + configure_tail_owned_attention_kv_cache(cache) + snap = CacheSnapshot(states=decoded.cache_snapshot.states[:2], meta_states=decoded.cache_snapshot.meta_states[:2]) + restore_cache(cache, snap) + assert all(isinstance(c, SegmentedKVCache) and c.offset == 1101 for c in cache) + assert mx.array_equal(cache[0].gather()[0], stock[0].state[0]).item() From 2e6edbebe69e0af8b346530b1c4f6a20a0589d64 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 08:40:58 +0200 Subject: [PATCH 10/37] Segmented KV: spill hashes each sealed segment block once (digest memo on the segment), /health block, SSD counters --- mtplx/cache_bank/codec.py | 37 +++++++++++++- mtplx/cache_bank/cold_tier.py | 21 ++++++++ mtplx/segmented_kv.py | 86 ++++++++++++++++++++++++++++++++- mtplx/session_bank.py | 8 +++ tests/test_segmented_kv_bank.py | 21 ++++++++ tests/test_segmented_kv_ssd.py | 72 +++++++++++++++++++++++++++ 6 files changed, 243 insertions(+), 2 deletions(-) diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index c87d4c0e6..5f7fe639b 100644 --- a/mtplx/cache_bank/codec.py +++ b/mtplx/cache_bank/codec.py @@ -155,6 +155,12 @@ def content_fingerprints(value: Any, *, rows: int | None = None) -> list[int] | _SEGMENTED_FINGERPRINT_BLOCKS = 16 +def _count_ssd(route: str) -> None: + from mtplx.segmented_kv import _count + + _count(route) + + def _is_segmented_state(value: Any) -> bool: """True for a SegmentedKVState; never imports the segmented module when it is not loaded.""" module = sys.modules.get("mtplx.segmented_kv") @@ -182,6 +188,12 @@ def __init__( self.reuse = reuse self.reused: dict[str, dict[str, Any]] = {} self.fingerprints: dict[str, tuple] = {} + # Streaming writer hooks for segmented rows (MTPLX_SEGMENTED_KV, spill_entry): a block + # inside one sealed segment that an earlier spill already hashed is referenced + # (known_blob confirms the blob is still on disk) instead of read back and hashed again; + # on_block records the digest of every block that was hashed. + self.known_blob: Callable[[dict[str, Any]], dict[str, Any] | None] | None = None + self.on_block: Callable[[Any, int, int, str], None] | None = None self.tensors: dict[str, bytes] = {} self.block_size = max(1, int(block_size)) self.should_abort = should_abort @@ -263,7 +275,9 @@ def encode(self, value: Any) -> Any: if _is_segmented_state(value): # MTPLX_SEGMENTED_KV: the same spec and blobs as the stock (keys, values) tuple of # the same rows; the history is sliced block by block, never gathered. - from mtplx.segmented_kv import segmented_state_tensors + from mtplx.segmented_kv import _count, segmented_state_tensors + + _count("ssd_layer_states_encoded") return { "kind": "tuple", @@ -412,8 +426,29 @@ def _encode_tensor_blocks( } ) continue + seg_rows = value if hasattr(value, "known_digest") and self.known_blob is not None else None + if seg_rows is not None: + known = seg_rows.known_digest(start, end) + known = self.known_blob(known) if known is not None else None + if known is not None: + _count_ssd("ssd_blocks_reused") + self.reused[name] = dict(known) + total += int(known["nbytes"]) + blocks.append( + { + "name": name, + "start": int(start), + "end": int(end), + "shape": [int(dim) for dim in chunk.shape], + "nbytes": int(known["nbytes"]), + } + ) + continue raw = self._materialize("block", chunk) self.tensors[name] = raw + if seg_rows is not None and self.on_block is not None: + _count_ssd("ssd_blocks_hashed") + self.on_block(seg_rows, int(start), int(end), name) total += len(raw) blocks.append( { diff --git a/mtplx/cache_bank/cold_tier.py b/mtplx/cache_bank/cold_tier.py index d11145079..8e82c060e 100644 --- a/mtplx/cache_bank/cold_tier.py +++ b/mtplx/cache_bank/cold_tier.py @@ -984,6 +984,25 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 on_unit=self._observe_encode_unit, ) codec.tensors = _WriteThroughTensors() + + def _known_blob(blob: dict[str, Any]) -> dict[str, Any] | None: + # MTPLX_SEGMENTED_KV: a block of an immutable segment hashed by an earlier spill. + # Claim first (so the orphan cleanup keeps it), then require it on disk. + digest = str(blob["sha256"]) + if digest not in claimed_digests: + claimed_digests.add(digest) + tier._claim_inflight(digests={digest}) + if not tier._blob_path(digest).exists(): + return None + written_state["logical"] += int(blob["nbytes"]) + written_state["deduped_hits"] += 1 + return blob + + def _on_block(rows: Any, start: int, end: int, name: str) -> None: + rows.remember_digest(start, end, tensor_blobs[name]) + + codec.known_blob = _known_blob + codec.on_block = _on_block boundaries = tuple( (int(r[0]), r[1], r[2] if len(r) > 2 else None) for r in (getattr(entry, "gdn_boundaries", None) or []) @@ -1015,6 +1034,8 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 "SessionBank SSD spill skipped: %s: %s", type(exc).__name__, exc ) return False + for name, blob in codec.reused.items(): + tensor_blobs[name] = dict(blob) if should_abort is not None: # Same fencing rationale as put_entry: the per-tensor evals # above schedule GPU work whose command buffers must drain in diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index fc784f111..b329adddf 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -119,7 +119,7 @@ class KVSegment: The tail of a cache is an unsealed segment: only its owner writes it. """ - __slots__ = ("__weakref__", "_holders", "id", "keys", "sealed", "values") + __slots__ = ("__weakref__", "_holders", "block_digests", "id", "keys", "sealed", "values") def __init__(self, keys: mx.array, values: mx.array, *, sealed: bool = False) -> None: self.id = next(_SEGMENT_IDS) @@ -127,6 +127,9 @@ def __init__(self, keys: mx.array, values: mx.array, *, sealed: bool = False) -> self.values = values self.sealed = bool(sealed) self._holders: weakref.WeakSet[Any] = weakref.WeakSet() + # SSD tier: sha256 and size of the 256-row blocks of this (immutable) segment that a + # spill already hashed, keyed (is_values, first_row, last_row) in segment rows. + self.block_digests: dict[tuple[bool, int, int], dict[str, Any]] = {} @property def capacity(self) -> int: @@ -798,6 +801,8 @@ def adapt_layer_for_restore(cache: list[Any], idx: int, state: Any) -> tuple[Any floor = segmented_kv_min_tokens() want = rows > floor if floor > 0 else isinstance(entry, SegmentedKVCache) if want and not isinstance(entry, SegmentedKVCache): + if not isinstance(state, SegmentedKVState): + _count("restored_stock_as_segment") # a stock snapshot (SSD restore, stock bank entry) entry = SegmentedKVCache(step=int(getattr(entry, "step", 256))) entry._mtplx_segmentable = True cache[idx] = entry @@ -907,6 +912,32 @@ def rows_slice(self, start: int, stop: int) -> mx.array: return mx.zeros((*self.shape[:2], 0, self.shape[3]), self.dtype) return parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=2) + def _block_home(self, start: int, stop: int) -> tuple[KVSegment, int, int] | None: + """(segment, first, last row in segment rows) when the block lies in one sealed segment.""" + offset = 0 + for ref in self._refs: + lo, hi = offset, offset + int(ref.n) + if lo <= start and stop <= hi: + seg = ref.segment + return (seg, start - lo, stop - lo) if seg.sealed else None + offset = hi + return None + + def known_digest(self, start: int, stop: int) -> dict[str, Any] | None: + """The blob a spill already wrote for this block, when it lies inside one sealed segment.""" + home = self._block_home(start, stop) + if home is None: + return None + seg, lo, hi = home + hit = seg.block_digests.get((self._values, lo, hi)) + return dict(hit) if hit is not None else None + + def remember_digest(self, start: int, stop: int, blob: dict[str, Any]) -> None: + home = self._block_home(start, stop) + if home is not None: + seg, lo, hi = home + seg.block_digests[(self._values, lo, hi)] = {"sha256": str(blob["sha256"]), "nbytes": int(blob["nbytes"])} + def __getitem__(self, index: Any) -> mx.array: if not isinstance(index, tuple) or len(index) != 4 or not all( isinstance(i, slice) and i == slice(None) for k, i in enumerate(index) if k != 2 @@ -925,6 +956,59 @@ def segmented_state_tensors(state: SegmentedKVState) -> tuple[SegmentedRows, Seg return SegmentedRows(refs, values=False), SegmentedRows(refs, values=True) +# ---- /health ------------------------------------------------------------------------------ + + +def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: + """The segmented-KV block of /health, None when the switch is off (the key is then absent). + + ``entries`` are the session bank's entries. Per entry: the segment count of its snapshot + (largest over the layers) or of its live cache, and the bytes of the segments it references. + """ + if not segmented_kv_enabled(): + return None + per_entry: list[dict[str, Any]] = [] + unique: dict[int, int] = {} + for entry in entries: + counts: list[int] = [] + sealed_bytes = 0 + snapshot = getattr(entry, "cache_snapshot", None) + for state in walk_segmented_states(snapshot.states) if snapshot is not None else (): + counts.append(len(state.refs)) + sealed_bytes += int(state.nbytes) + for ref in state.refs: + unique.setdefault(ref.segment.id, int(ref.segment.nbytes)) + live = getattr(entry, "cache_ref", None) + for layer in live or (): + if isinstance(layer, SegmentedKVCache): + counts.append(layer.segment_count) + if counts: + per_entry.append( + { + "session_id": getattr(entry, "session_id", None), + "prefix_len": getattr(entry, "prefix_len", None), + "segments": max(counts), + "sealed_bytes": sealed_bytes, + "live": bool(live), + } + ) + counts_now = dict(route_counts) + return { + "enabled": True, + "min_tokens": segmented_kv_min_tokens(), + "ssd": segmented_ssd_enabled(), + "prefill_route": prefill_route(), + "entries": per_entry, + "max_segments": max((e["segments"] for e in per_entry), default=0), + "sealed_bytes_unique": sum(unique.values()), + "merges": counts_now.get("merge", 0), + "seals": counts_now.get("seal", 0), + "seal_copies": counts_now.get("seal_copy", 0), + "route_counts": counts_now, + "ssd_counts": {k: v for k, v in counts_now.items() if k.startswith("ssd_")}, + } + + # ---- accounting ----------------------------------------------------------------------------- diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py index 68ffdce42..7d99e5c36 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -38,6 +38,7 @@ from .segmented_kv import ( SegmentedKVCache, SegmentedKVState, + segmented_kv_health, segmented_ssd_enabled, walk_segmented_states, ) @@ -2953,6 +2954,13 @@ def flush_cold_tier(self, *, timeout_s: float = 30.0) -> bool: return bool(flush(timeout_s=timeout_s)) def to_dict(self) -> dict[str, Any]: + payload = self._to_dict() + segmented = segmented_kv_health(list(self._entries.values())) + if segmented is not None: + payload["segmented_kv"] = segmented # MTPLX_SEGMENTED_KV only; absent otherwise + return payload + + def _to_dict(self) -> dict[str, Any]: return { "max_entries": self.max_entries, "max_bytes": self.max_bytes, diff --git a/tests/test_segmented_kv_bank.py b/tests/test_segmented_kv_bank.py index ca8c4dcc2..a845cf7e9 100644 --- a/tests/test_segmented_kv_bank.py +++ b/tests/test_segmented_kv_bank.py @@ -262,3 +262,24 @@ def test_admission_prices_no_history_copy_for_a_segmented_entry(monkeypatch) -> assert _admission_restore_copies_prefix(seg_entry, "clone", "x") is False assert _admission_restore_copies_prefix(seg_entry, "reference", "x") is False assert _admission_restore_copies_prefix(stock_entry, "clone", "y") is True + + +def test_health_block_is_absent_with_the_switch_off_and_reports_segments_with_it_on(monkeypatch) -> None: + monkeypatch.delenv("MTPLX_SEGMENTED_KV", raising=False) + bank = _bank(monkeypatch) + runtime = Runtime("segmented") + tokens = list(range(BASE)) + _, entry = _turn(bank, runtime, tokens, BASE, 1, "s") + tokens2 = tokens + list(range(BASE * 10, BASE * 10 + TURN)) + _turn(bank, runtime, tokens2, TURN, 2, "s", prev=tokens) + assert "segmented_kv" not in bank.to_dict() + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + block = bank.to_dict()["segmented_kv"] + assert block["enabled"] is True and block["entries"] + assert block["max_segments"] >= 2 + assert block["sealed_bytes_unique"] > 0 + assert {"route_counts", "ssd_counts", "merges", "seals", "min_tokens"} <= set(block) + assert all(e["segments"] >= 1 and e["sealed_bytes"] > 0 for e in block["entries"]) + import json + + json.dumps(block) # /health serialises it diff --git a/tests/test_segmented_kv_ssd.py b/tests/test_segmented_kv_ssd.py index 3cef9bcaf..e4e886095 100644 --- a/tests/test_segmented_kv_ssd.py +++ b/tests/test_segmented_kv_ssd.py @@ -224,3 +224,75 @@ def test_a_stock_entry_on_disk_restores_with_the_switch_on_as_segments(monkeypat restore_cache(cache, snap) assert all(isinstance(c, SegmentedKVCache) and c.offset == 1101 for c in cache) assert mx.array_equal(cache[0].gather()[0], stock[0].state[0]).item() + + +def _count_block_units(monkeypatch): + from mtplx.cache_bank import codec as codec_module + + calls = [] + original = codec_module.TreeCodec._materialize + + def counting(self, kind, array): + calls.append(kind) + return original(self, kind, array) + + monkeypatch.setattr(codec_module.TreeCodec, "_materialize", counting) + return calls + + +def _grown(seg, extra_rows, seed): + """The same segments plus one more sealed segment (the next turn).""" + out = [] + for layer in seg: + c = SegmentedKVCache() + c.state = layer.state + c.update_and_fetch(_rand(seed, extra_rows), _rand(seed + 1, extra_rows)) + c.seal() + out.append(c) + return out + + +def test_spill_does_not_rehash_sealed_segments_it_already_hashed(monkeypatch, tmp_path) -> None: + seg, _ = _segmented_and_stock(lens=(4000, 700, 90)) + for layer in seg: + layer.seal() + calls = _count_block_units(monkeypatch) + tier = _tier(tmp_path / "seg") + assert tier.spill_entry(_Entry(seg, range(4790))) + first = calls.count("block") + grown = _grown(seg, 300, 500) + calls.clear() + assert tier.spill_entry(_Entry(grown, range(5090))) + second = calls.count("block") + assert first > 40 and second < 0.45 * first, (first, second) + # and the result is what the stock cache of the same rows writes + stock = [] + for layer in grown: + k, v = layer.gather() + c = KVCache() + c.keys, c.values, c.offset = k, v, k.shape[2] + stock.append(c) + ref = _tier(tmp_path / "stock") + assert ref.spill_entry(_Entry(stock, range(5090))) + assert _blobs(tier).keys() >= _blobs(ref).keys() + eid = lambda t, e: t._metadata_for_entry(e, capabilities=(), payload_nbytes=0)["entry_id"] + a = _entry_json(tier, _Entry(grown, range(5090))) + b = _entry_json(ref, _Entry(stock, range(5090))) + assert a["payload_spec"] == b["payload_spec"] and a["tensor_blobs"] == b["tensor_blobs"] + assert eid(tier, _Entry(grown, range(5090))) == eid(ref, _Entry(stock, range(5090))) + + +def test_spill_rehashes_a_block_whose_blob_is_gone(monkeypatch, tmp_path) -> None: + seg, _ = _segmented_and_stock(lens=(1500, 700)) + for layer in seg: + layer.seal() + tier = _tier(tmp_path / "seg") + assert tier.spill_entry(_Entry(seg, range(2200))) + full = max(p.stat().st_size for p in (tier.base_dir / "blobs").rglob("*.bin")) + victims = [p for p in sorted((tier.base_dir / "blobs").rglob("*.bin")) if p.stat().st_size == full][:5] + for blob in victims: + blob.unlink() + grown = _grown(seg, 256, 700) + assert tier.spill_entry(_Entry(grown, range(2456))) + for blob in victims: + assert blob.exists() # rewritten from the live rows From d80590b3649d2d5ad3b8b52f1325e8a3d224994f Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 11:45:18 +0200 Subject: [PATCH 11/37] Segmented KV: the tiered merge also merges segments that bank snapshots hold (nothing merged before: every turn's snapshot holds its segments) --- mtplx/segmented_kv.py | 24 +++++++++++++---------- tests/test_segmented_kv.py | 40 ++++++++++++++++++++++++++++++++++---- 2 files changed, 50 insertions(+), 14 deletions(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index b329adddf..f8ae585da 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -466,13 +466,18 @@ def state(self, value: Any) -> None: def compact(self, *, max_segments: int | None = None) -> int: """Merge small recent sealed segments among themselves (geometric tiers). - Two adjacent sealed segments merge when no snapshot holds either (reference count 1: - this cache), the older is at most ``MTPLX_SEGMENTED_KV_TIER`` times (default 2) the - newer one and the result stays within ``MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS`` (default - 16384) rows. The big history segment is never merged here: a full copy would bring the - duplicate back while a snapshot, a subagent or a lease references it. Copying two small - segments costs 0.8 ms (512 + 512 rows) to 4.1 ms (4096 + 4096 rows) for 16 layers (phase 1 - measurement). Returns the number of merges. + Two adjacent sealed segments merge when the older is at most ``MTPLX_SEGMENTED_KV_TIER`` + times (default 2) the newer one and the result stays within + ``MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS`` (default 16384) rows. The big history segment is + never merged here: a full copy would bring the duplicate back while a snapshot, a + subagent or a lease references it. A segment that bank snapshots also hold (reference + count above 1) is merged too: the merged copy is new memory (at most the merged rows, + released when the snapshots holding the originals go), and without it nothing would + ever merge, because every turn's snapshot holds that turn's segments. The segment count + is thereby bounded by the tier structure (logarithmic in the rows added) instead of + growing by two per turn. Copying two small segments costs 0.8 ms (512 + 512 rows) to + 4.1 ms (4096 + 4096 rows) for 16 layers (phase 1 measurement). Returns the number of + merges. """ tier = max(1, _int_env("MTPLX_SEGMENTED_KV_TIER", 2)) cap_rows = _int_env("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", 16384) @@ -480,11 +485,10 @@ def compact(self, *, max_segments: int | None = None) -> int: merged = 0 while len(self._sealed) >= 2: older, newer = self._sealed[-2], self._sealed[-1] - alone = older.segment.refcount <= 1 and newer.segment.refcount <= 1 small = int(older.n) + int(newer.n) <= cap_rows geometric = int(older.n) <= tier * int(newer.n) - over = len(self._sealed) > limit and alone - if not (alone and ((small and geometric) or over)): + over = len(self._sealed) > limit and small + if not (small and (geometric or over)): break self._merge_last_two() merged += 1 diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py index 5def8c281..584a1c046 100644 --- a/tests/test_segmented_kv.py +++ b/tests/test_segmented_kv.py @@ -293,17 +293,49 @@ def test_tiered_merge_only_merges_small_recent_unshared_segments(monkeypatch) -> assert _same(cache, control) -def test_merge_leaves_segments_that_a_snapshot_holds(monkeypatch) -> None: +def test_merge_also_merges_segments_a_snapshot_holds_and_leaves_the_snapshot_intact(monkeypatch) -> None: + """Every turn's bank snapshot holds that turn's segments; without merging held segments nothing + would ever merge and the count would grow by two per turn.""" monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "16384") cache = SegmentedKVCache() for seed in (1, 2): _fill(cache, [(seed, 300)]) cache.seal() held = cache.state # a bank entry references both segments + held_segments = [ref.segment for ref in held.refs] before = [ref.segment for ref in cache._sealed] - _fill(cache, [(3, 300)]) # new tail: compact() must not merge held segments - assert [ref.segment for ref in cache._sealed] == before - assert held.rows == 600 + assert before == held_segments + _fill(cache, [(3, 300)]) # new tail: compact() merges the two equal-sized held segments + assert len(cache._sealed) == 1 and cache._sealed[0].n == 600 + assert cache.merges == 1 and cache.merge_rows == 600 + assert held.rows == 600 and [ref.segment for ref in held.refs] == held_segments # the snapshot is untouched + control = KVCache() + for seed in (1, 2): + control.update_and_fetch(*_rows(seed, 300)) + k, v = held.to_arrays() + assert mx.array_equal(k, control.state[0]).item() and mx.array_equal(v, control.state[1]).item() + merged_k = cache._sealed[0].segment.keys[..., :600, :] + assert mx.array_equal(merged_k, control.state[0]).item() + + +def test_segment_count_stays_bounded_over_many_turns_with_snapshots(monkeypatch) -> None: + """Two seals per turn (a prompt snapshot and the final one), every snapshot kept alive.""" + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "16384") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_TIER", "2") + cache = SegmentedKVCache() + _fill(cache, [(1, 20000)]) + cache.seal() + snapshots = [cache.state] + peak = 0 + for turn in range(2, 26): + _fill(cache, [(turn, 700)]) + snapshots.append(cache.state) # prompt snapshot + _fill(cache, [(100 + turn, 150)]) + snapshots.append(cache.state) # final snapshot + peak = max(peak, cache.segment_count) + assert peak <= 12, peak + assert cache.segment_count <= 12 + assert cache.merges > 0 def test_nbytes_counts_each_shared_segment_once() -> None: From cac2908ea0736a855b5332da133df9a8af97d3c0 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 12:38:22 +0200 Subject: [PATCH 12/37] Changelog: spill memo, merge of held segments, /health block --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b340b343..bbd9e756d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ All notable user-facing changes to MTPLX. The format is based on ### 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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. Not covered: the compiled verify (segmented caches verify eager), paged and quantized KV. +- **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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. The streaming spill hashes each sealed segment's blocks once (the digest is kept on the segment, and the blob is confirmed on disk before it is referenced), so a follow-up turn no longer rereads and rehashes the whole history. The tiered merge of small recent segments also merges segments that bank snapshots hold (before, nothing ever merged because every turn's snapshot holds its segments), which keeps a conversation at a handful of segments (3 to 5 over 12 follow-up turns at 50K) instead of two more per turn, below the 12-segment fused launch. `/health` gets a `segmented_kv` block (segments and sealed bytes per entry, merges, route and SSD counters), absent when the switch is off. Not covered: the compiled verify (segmented caches verify eager), paged and quantized KV. ### Fixed From 2966d66de7fd5fa2f42c53fc07772de8cb158dc3 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 16:12:43 +0200 Subject: [PATCH 13/37] Segmented KV test: an SSD restore builds the request's cache layout with segments too (the fix itself is the base commit) --- tests/test_segmented_kv_ssd.py | 47 ++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/test_segmented_kv_ssd.py b/tests/test_segmented_kv_ssd.py index e4e886095..0ab6e9be1 100644 --- a/tests/test_segmented_kv_ssd.py +++ b/tests/test_segmented_kv_ssd.py @@ -296,3 +296,50 @@ def test_spill_rehashes_a_block_whose_blob_is_gone(monkeypatch, tmp_path) -> Non assert tier.spill_entry(_Entry(grown, range(2456))) for blob in victims: assert blob.exists() # rewritten from the live rows + + +class _FakeRuntime: + from pathlib import Path as _P + + model_path = _P("models/example") + mtp_enabled = False + + def __init__(self): + self.made = 0 + + def make_cache(self): + self.made += 1 + return [KVCache() for _ in range(2)] + + +@pytest.mark.parametrize("flag", ["1", ""]) +def test_ssd_restore_builds_the_requests_cache_layout_with_and_without_the_switch(monkeypatch, tmp_path, flag) -> None: + """The SSD restore builds the request's cache layout, with or without segments (the fix itself and + its own tests are in tests/test_ssd_restore_cache_layout.py).""" + from mtplx.session_bank import SessionBank + + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100000") # stock layers stay stock + if flag: + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + else: + monkeypatch.delenv("MTPLX_SEGMENTED_KV", raising=False) + runtime = _FakeRuntime() + tier = _tier(tmp_path / "t") + bank = SessionBank(cold_tier=tier) + cache = [KVCache() for _ in range(2)] + for c in cache: + c.update_and_fetch(_rand(1, 700), _rand(2, 700)) + tokens = list(range(700)) + bank.put(runtime=runtime, token_ids=tokens, cache=cache, logits=mx.zeros((1, 8)), hidden=None, session_id="s") + assert tier.flush(timeout_s=30.0) + fresh = SessionBank(cold_tier=tier) + made_by_factory = [] + + def factory(): + made_by_factory.append(1) + return [KVCache() for _ in range(2)] + + restored = fresh.restore(runtime, tokens + [701, 702], session_id="s", cache_factory=factory) + assert restored is not None and restored.restore_mode == "ssd_clone" + assert made_by_factory + assert restored.cache[0].offset == 700 From caaed25116751357676008677c0b38be90d28ad6 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 12:52:18 +0200 Subject: [PATCH 14/37] Changelog: SSD restore layout --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bbd9e756d..0b72cc46e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ All notable user-facing changes to MTPLX. The format is based on ### 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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. The streaming spill hashes each sealed segment's blocks once (the digest is kept on the segment, and the blob is confirmed on disk before it is referenced), so a follow-up turn no longer rereads and rehashes the whole history. The tiered merge of small recent segments also merges segments that bank snapshots hold (before, nothing ever merged because every turn's snapshot holds its segments), which keeps a conversation at a handful of segments (3 to 5 over 12 follow-up turns at 50K) instead of two more per turn, below the 12-segment fused launch. `/health` gets a `segmented_kv` block (segments and sealed bytes per entry, merges, route and SSD counters), absent when the switch is off. Not covered: the compiled verify (segmented caches verify eager), paged and quantized KV. +- **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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. The streaming spill hashes each sealed segment's blocks once (the digest is kept on the segment, and the blob is confirmed on disk before it is referenced), so a follow-up turn no longer rereads and rehashes the whole history. The tiered merge of small recent segments also merges segments that bank snapshots hold (before, nothing ever merged because every turn's snapshot holds its segments), which keeps a conversation at a handful of segments (3 to 5 over 12 follow-up turns at 50K) instead of two more per turn, below the 12-segment fused launch. An SSD restore builds the request's dense cache layout instead of the paged default (with the switch on). `/health` gets a `segmented_kv` block (segments and sealed bytes per entry, merges, route and SSD counters), absent when the switch is off. Not covered: the compiled verify (segmented caches verify eager), paged and quantized KV. ### Fixed From b9e424fced6cd7758178002d5b7fff8e786908e3 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 14:08:50 +0200 Subject: [PATCH 15/37] 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 --- mtplx/attention_split.py | 13 ++- mtplx/kernels/sdpa_segmented.py | 3 + mtplx/segmented_kv.py | 96 +++++++++++++++++++- tests/test_segmented_kv_support.py | 140 +++++++++++++++++++++++++++++ 4 files changed, 248 insertions(+), 4 deletions(-) create mode 100644 tests/test_segmented_kv_support.py diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index e8c2f248f..b77de93cf 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -933,6 +933,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 @@ -944,7 +945,10 @@ 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))) + stats["installed"] += int(newly_hooked) attn._mtplx_split_full_attention_enabled = bool( active or sdpa_2pass or vllm_metal_paged or gqa_packed ) @@ -966,4 +970,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 _env_enabled("MTPLX_SEGMENTED_KV"): + # 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 diff --git a/mtplx/kernels/sdpa_segmented.py b/mtplx/kernels/sdpa_segmented.py index accdfc71a..d317b6ed0 100644 --- a/mtplx/kernels/sdpa_segmented.py +++ b/mtplx/kernels/sdpa_segmented.py @@ -44,6 +44,9 @@ from .sdpa_2pass import unnormalized_partials_dtype from .sdpa_nax_flash import _HEADER +#: Head dimensions the segment kernel serves (the head-dim-split TensorOps kernel). +SUPPORTED_HEAD_DIMS = (256,) + #: Metal allows 31 buffer bindings: queries + 2 per segment + meta + scale + 3 outputs. MAX_FUSED_SEGMENTS = 12 #: The reduce kernel binds one partial buffer per launch group plus sums, maxs, starts, out. diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index f8ae585da..d6eee64cb 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -50,11 +50,96 @@ _SEGMENT_IDS = itertools.count(1) -def segmented_kv_enabled() -> bool: - """MTPLX_SEGMENTED_KV (default off), read on every call.""" +def segmented_kv_requested() -> bool: + """MTPLX_SEGMENTED_KV (default off), read on every call: what the operator asked for.""" return os.environ.get("MTPLX_SEGMENTED_KV", "").strip().lower() in {"1", "true", "yes", "on"} +#: Verdict on the loaded model (``evaluate_model_support``); None until a model was checked. +_MODEL_SUPPORT: dict[str, Any] | None = None +_UNSUPPORTED_LOGGED = False + + +def segmented_kv_enabled() -> bool: + """MTPLX_SEGMENTED_KV, unless the loaded model cannot use segments (then everything behaves + as with the switch off: stock caches, stock bank, stock SSD path).""" + if not segmented_kv_requested(): + return False + return _MODEL_SUPPORT is None or bool(_MODEL_SUPPORT.get("supported", True)) + + +def model_support() -> dict[str, Any] | None: + return None if _MODEL_SUPPORT is None else dict(_MODEL_SUPPORT) + + +def reset_model_support() -> None: + global _MODEL_SUPPORT, _UNSUPPORTED_LOGGED + _MODEL_SUPPORT = None + _UNSUPPORTED_LOGGED = False + + +def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | None = None) -> dict[str, Any]: + """Whether the segmented cache can serve this model without ever gathering the history. + + A call that the segment kernel does not take runs on one gathered copy of the history + (correct, but a full copy per layer per verify), so a model whose full-attention layers the + kernel cannot serve keeps the stock cache instead. ``layers`` are the full-attention modules + (``head_dim``, ``num_attention_heads``, ``num_key_value_heads``; ``hooked`` says which of + them run the split-attention hook, the only code that reads segments). The verdict is + stored; it switches the feature off (``segmented_kv_enabled``) and is reported on /health. + """ + from .kernels.sdpa_segmented import SUPPORTED_HEAD_DIMS, segments_supported + from .nax_verify import nax_available + + global _MODEL_SUPPORT, _UNSUPPORTED_LOGGED + layers = list(layers) + flags = list(hooked) if hooked is not None else [True] * len(layers) + reasons: list[str] = [] + shapes: set[tuple[int, int, int]] = set() + if not layers: + reasons.append("no_full_attention_layers") + if not all(flags): + reasons.append("attention_not_hooked") + for attn in layers: + head_dim = getattr(attn, "head_dim", None) + q_heads = getattr(attn, "num_attention_heads", None) or getattr(attn, "n_heads", None) + kv_heads = getattr(attn, "num_key_value_heads", None) or getattr(attn, "n_kv_heads", None) + if not (head_dim and q_heads and kv_heads): + reasons.append("attention_shape_unknown") + continue + shapes.add((int(head_dim), int(q_heads), int(kv_heads))) + for head_dim, q_heads, kv_heads in sorted(shapes): + if head_dim not in SUPPORTED_HEAD_DIMS: + reasons.append(f"head_dim_{head_dim}_unsupported") + elif q_heads % kv_heads or not segments_supported(1, q_heads // kv_heads, head_dim): + reasons.append(f"gqa_{q_heads}_{kv_heads}_unsupported") + if not all(getattr(a, "_mtplx_gqa_packed_sdpa_enabled", False) for a in layers): + reasons.append("gqa_packed_sdpa_off") + if os.environ.get("MTPLX_NAX_FLASH_ROUTE", "").strip().lower() not in {"1", "true", "yes", "on"}: + reasons.append("nax_flash_route_off") + if not nax_available(): + reasons.append("nax_unavailable") + reasons = list(dict.fromkeys(reasons)) + verdict = { + "supported": not reasons, + "reasons": reasons, + "head_dims": sorted({s[0] for s in shapes}), + "gqa": sorted({s[1] // s[2] for s in shapes if s[2] and s[1] % s[2] == 0}), + "supported_head_dims": list(SUPPORTED_HEAD_DIMS), + "full_attention_layers": len(layers), + } + _MODEL_SUPPORT = verdict + if reasons and segmented_kv_requested() and not _UNSUPPORTED_LOGGED: + _UNSUPPORTED_LOGGED = True + print( + "[mtplx] MTPLX_SEGMENTED_KV is on but this model cannot use it (" + + ", ".join(reasons) + + "): keeping the stock KV cache", + flush=True, + ) + return verdict + + def _int_env(name: str, default: int) -> int: raw = os.environ.get(name, "").strip() try: @@ -969,8 +1054,11 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: ``entries`` are the session bank's entries. Per entry: the segment count of its snapshot (largest over the layers) or of its live cache, and the bytes of the segments it references. """ - if not segmented_kv_enabled(): + if not segmented_kv_requested(): return None + if not segmented_kv_enabled(): + # Asked for, but the loaded model cannot use it: say so (and why) instead of staying silent. + return {"enabled": False, "requested": True, "model_support": model_support()} per_entry: list[dict[str, Any]] = [] unique: dict[int, int] = {} for entry in entries: @@ -999,6 +1087,8 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: counts_now = dict(route_counts) return { "enabled": True, + "requested": True, + "model_support": model_support(), "min_tokens": segmented_kv_min_tokens(), "ssd": segmented_ssd_enabled(), "prefill_route": prefill_route(), diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py new file mode 100644 index 000000000..86d6d5d44 --- /dev/null +++ b/tests/test_segmented_kv_support.py @@ -0,0 +1,140 @@ +"""Models the segment kernel cannot serve keep the stock cache (and say why), never the gather route.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import mlx.core as mx +import pytest +from mlx_lm.models.cache import KVCache + +import mtplx.segmented_kv as module +from mtplx.cache_state import configure_tail_owned_attention_kv_cache +from mtplx.segmented_kv import ( + evaluate_model_support, + model_support, + reset_model_support, + segmented_kv_enabled, + segmented_kv_health, + segmented_ssd_enabled, +) + + +@pytest.fixture(autouse=True) +def _cpu_and_clean(monkeypatch): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + reset_model_support() + monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + monkeypatch.delenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", raising=False) + monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: True) + yield + reset_model_support() + mx.set_default_device(previous) + + +def _attn(head_dim=256, q=24, kv=4, packed=True): + return SimpleNamespace( + head_dim=head_dim, num_attention_heads=q, num_key_value_heads=kv, _mtplx_gqa_packed_sdpa_enabled=packed + ) + + +def test_a_supported_model_stays_enabled() -> None: + verdict = evaluate_model_support([_attn(), _attn()]) + assert verdict["supported"] and verdict["reasons"] == [] + assert verdict["head_dims"] == [256] and verdict["gqa"] == [6] + assert segmented_kv_enabled() is True + + +@pytest.mark.parametrize( + "layers, hooked, reason", + [ + ([_attn(head_dim=64, q=16, kv=4)], None, "head_dim_64_unsupported"), + ([_attn(head_dim=256, q=40, kv=1)], None, "gqa_40_1_unsupported"), + ([_attn()], [False], "attention_not_hooked"), + ([_attn(packed=False)], None, "gqa_packed_sdpa_off"), + ([], None, "no_full_attention_layers"), + ([SimpleNamespace(head_dim=None)], None, "attention_shape_unknown"), + ], +) +def test_an_unsupported_model_switches_the_feature_off_with_a_reason(layers, hooked, reason, capsys) -> None: + verdict = evaluate_model_support(layers, hooked=hooked) + assert not verdict["supported"] and reason in verdict["reasons"] + assert segmented_kv_enabled() is False + assert segmented_ssd_enabled() is False + out = capsys.readouterr().out + assert out.count("MTPLX_SEGMENTED_KV is on but this model cannot use it") == 1 and reason in out + + +def test_the_log_line_appears_once(capsys) -> None: + evaluate_model_support([_attn(head_dim=64)]) + evaluate_model_support([_attn(head_dim=64)]) + assert capsys.readouterr().out.count("keeping the stock KV cache") == 1 + + +def test_missing_flash_route_or_nax_is_a_reason(monkeypatch) -> None: + monkeypatch.delenv("MTPLX_NAX_FLASH_ROUTE") + assert "nax_flash_route_off" in evaluate_model_support([_attn()])["reasons"] + monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") + monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: False) + assert "nax_unavailable" in evaluate_model_support([_attn()])["reasons"] + + +def test_an_unsupported_model_keeps_stock_caches_and_never_builds_a_segmented_one() -> None: + evaluate_model_support([_attn(head_dim=64)]) + cache = [KVCache(), KVCache()] + configure_tail_owned_attention_kv_cache(cache) + assert all(type(c) is KVCache for c in cache) + assert not any(getattr(c, "_mtplx_segmentable", False) for c in cache) + for c in cache: + c.update_and_fetch(mx.zeros((1, 2, 9000, 8)), mx.zeros((1, 2, 9000, 8))) + configure_tail_owned_attention_kv_cache(cache) # what the repage after a prefill runs + assert all(type(c) is KVCache for c in cache) + + +def test_health_reports_the_reason_and_is_absent_when_not_requested(monkeypatch) -> None: + evaluate_model_support([_attn(head_dim=64)]) + block = segmented_kv_health([]) + assert block["enabled"] is False and block["requested"] is True + assert "head_dim_64_unsupported" in block["model_support"]["reasons"] + monkeypatch.delenv("MTPLX_SEGMENTED_KV") + assert segmented_kv_health([]) is None + + +def test_health_of_a_supported_model_carries_the_verdict() -> None: + evaluate_model_support([_attn()]) + block = segmented_kv_health([]) + assert block["enabled"] is True and block["model_support"]["supported"] is True + + +def test_nothing_is_decided_before_a_model_is_checked() -> None: + assert model_support() is None and segmented_kv_enabled() is True + assert module._MODEL_SUPPORT is None + + +def test_configure_split_full_attention_decides_from_the_hooked_layers(monkeypatch) -> None: + from mlx_lm.models.qwen3_next import Qwen3NextAttention + + from mtplx.attention_split import configure_split_full_attention + + monkeypatch.setenv("MTPLX_GQA_PACKED_SDPA", "1") + + def model(head_dim): + args = SimpleNamespace( + hidden_size=64, num_attention_heads=4, num_key_value_heads=2, head_dim=head_dim, attention_bias=False, + rms_norm_eps=1e-6, partial_rotary_factor=0.25, rope_theta=10000.0, rope_scaling=None, + max_position_embeddings=1024, + ) + attn = Qwen3NextAttention(args) + layer = SimpleNamespace(is_linear=False, self_attn=attn) + return SimpleNamespace(model=SimpleNamespace(layers=[layer])) + + stats = configure_split_full_attention(model(32)) + assert stats["segmented_kv_supported"] is False + assert "head_dim_32_unsupported" in model_support()["reasons"] + assert segmented_kv_enabled() is False + reset_model_support() + stats = configure_split_full_attention(model(256)) + assert stats["segmented_kv_supported"] is True + assert segmented_kv_enabled() is True From 377aa86bdc8ce3486cd99c42f057f1a6bf44afc7 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 14:22:46 +0200 Subject: [PATCH 16/37] Segmented verify kernel at head_dim 128 (D is a template parameter), defensive lse fallback, tests (kernel tests pending a power-connected run) --- mtplx/kernels/sdpa_segmented.py | 17 ++- mtplx/segmented_kv.py | 11 +- tests/test_sdpa_segmented_hd128.py | 180 +++++++++++++++++++++++++++++ 3 files changed, 200 insertions(+), 8 deletions(-) create mode 100644 tests/test_sdpa_segmented_hd128.py diff --git a/mtplx/kernels/sdpa_segmented.py b/mtplx/kernels/sdpa_segmented.py index d317b6ed0..7273a6fc1 100644 --- a/mtplx/kernels/sdpa_segmented.py +++ b/mtplx/kernels/sdpa_segmented.py @@ -44,8 +44,11 @@ from .sdpa_2pass import unnormalized_partials_dtype from .sdpa_nax_flash import _HEADER -#: Head dimensions the segment kernel serves (the head-dim-split TensorOps kernel). -SUPPORTED_HEAD_DIMS = (256,) +#: Head dimensions the segment kernel serves (the head-dim-split TensorOps kernel; ``D`` is a +#: template parameter, 128 gives 64 dims per simdgroup at the halves split). The contiguous +#: dispatcher in sdpa_nax_flash_dsplit stays 256 only: a model with head_dim 128 that does not +#: use segments keeps exactly the routes it has today. +SUPPORTED_HEAD_DIMS = (128, 256) #: Metal allows 31 buffer bindings: queries + 2 per segment + meta + scale + 3 outputs. MAX_FUSED_SEGMENTS = 12 @@ -210,7 +213,7 @@ def sdpa_nax_flash_dsplit_segments( ) -> mx.array | None: """Exact attention of ``queries`` over ``segments`` = [(keys, values, n), ...], oldest first. - ``keys``/``values`` are (1, Hk, capacity, 256) buffers of which the first ``n`` rows are + ``keys``/``values`` are (1, Hk, capacity, D) buffers, D in SUPPORTED_HEAD_DIMS, of which the first ``n`` rows are live. Every segment but the last is fully visible; the last holds the rows being verified as its final ``q_len`` rows and keeps the tail-causal mask. None when the contract is not met (the caller falls back to ``segmented_kv.gather_attention``). @@ -224,7 +227,7 @@ def sdpa_nax_flash_dsplit_segments( if queries.ndim != 4: return _bail("ndim") bsz, hq, q_len, d = (int(x) for x in queries.shape) - if bsz != 1 or d != 256: + if bsz != 1 or d not in SUPPORTED_HEAD_DIMS: return _bail("shape_gate") if not segments: return _bail("no_segments") @@ -253,7 +256,9 @@ def sdpa_nax_flash_dsplit_segments( if int(segments[-1][2]) < q_len: return _bail("tail_shorter_than_q") - if len(segments) == 1: + if len(segments) == 1 and d == 256: + # One segment is the contiguous kernel's input (bit-identical). The contiguous + # dispatcher is 256 only; a head_dim-128 segment goes through the fused path below. keys, values, n = segments[0] return _dsplit.sdpa_nax_flash_dsplit( queries=queries, keys=keys, values=values, offset=int(n), scale=scale @@ -336,5 +341,5 @@ def sdpa_nax_flash_dsplit_segments( def segments_supported(q_len: int, gqa_factor: int, head_dim: int) -> bool: """Static part of the contract (no device work), for the route gate in attention_split.""" - return head_dim == 256 and 1 <= q_len <= 10 and gqa_factor * q_len <= 32 + return head_dim in SUPPORTED_HEAD_DIMS and 1 <= q_len <= 10 and gqa_factor * q_len <= 32 diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index d6eee64cb..e166f20f6 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -702,8 +702,15 @@ def _segment_sdpa_lse_fast(queries, keys, values, n, *, scale, causal): def segment_sdpa_lse(queries, keys, values, n, *, scale, causal): """Per segment (out, lse): the MLX kernel when present, else the exact reference route.""" if sdpa_lse_available(): - _count("sdpa_lse_kernel") - return _segment_sdpa_lse_fast(queries, keys, values, n, scale=scale, causal=causal) + try: + out = _segment_sdpa_lse_fast(queries, keys, values, n, scale=scale, causal=causal) + except (ValueError, RuntimeError): + # The fused kernel does not take this shape (the detection probe uses a tiny head + # dim): the exact reference route serves the call, once counted. + _count("sdpa_lse_kernel_unsupported") + else: + _count("sdpa_lse_kernel") + return out _count("sdpa_lse_reference") return _segment_sdpa_lse_reference(queries, keys, values, n, scale=scale, causal=causal) diff --git a/tests/test_sdpa_segmented_hd128.py b/tests/test_sdpa_segmented_hd128.py new file mode 100644 index 000000000..e7ec614bc --- /dev/null +++ b/tests/test_sdpa_segmented_hd128.py @@ -0,0 +1,180 @@ +"""Segmented verify kernel at head_dim 128 (Qwen3-8B, Llama-like shapes: 32 query / 8 KV heads, GQA 4). + +The kernel tests need a TensorOps GPU (pending on a machine on power); the support-contract and the +lse-fallback tests are CPU-safe. +""" + +from __future__ import annotations + +import math + +import mlx.core as mx +import pytest + +import mtplx.segmented_kv as module +from mtplx.kernels import sdpa_segmented as S +from mtplx.nax_verify import nax_available + +HQ, HK, D = 32, 8, 128 +SCALE = 1.0 / math.sqrt(D) + +needs_nax = pytest.mark.skipif( + not mx.metal.is_available() or not nax_available(), + reason="TensorOps kernel unavailable on this machine", +) + + +def _data(total, q_len, peaked=False, seed=3): + mx.random.seed(seed) + q = (mx.random.normal((1, HQ, q_len, D)) * (2.5 if peaked else 0.5)).astype(mx.bfloat16) + k = (mx.random.normal((1, HK, total + 64, D)) * 0.5).astype(mx.bfloat16) + v = (mx.random.normal((1, HK, total + 64, D)) * 0.5).astype(mx.bfloat16) + mx.eval(q, k, v) + return q, k, v + + +def _split(k, v, lens, spare=0): + segs, a = [], 0 + for n in lens: + kk, vv = mx.contiguous(k[:, :, a : a + n + spare]), mx.contiguous(v[:, :, a : a + n + spare]) + mx.eval(kk, vv) + segs.append((kk, vv, n)) + a += n + return segs + + +def _lens(total, nseg): + lens = [total // nseg] * nseg + lens[-1] += total - sum(lens) + return lens + + +def _ref(q, k, v, total): + q_len = q.shape[2] + kf = mx.repeat(k[:, :, :total].astype(mx.float32), HQ // HK, axis=1) + vf = mx.repeat(v[:, :, :total].astype(mx.float32), HQ // HK, axis=1) + s = (q.astype(mx.float32) @ kf.transpose(0, 1, 3, 2)) * SCALE + vis = mx.arange(total)[None, :] <= (total - q_len + mx.arange(q_len))[:, None] + return mx.softmax(mx.where(vis[None, None], s, -mx.inf), axis=-1) @ vf + + +def test_the_support_contract_names_head_dim_128_and_nothing_below() -> None: + assert S.segments_supported(4, 4, 128) and S.segments_supported(1, 4, 256) + assert not S.segments_supported(4, 4, 64) and not S.segments_supported(4, 4, 96) + assert S.SUPPORTED_HEAD_DIMS == (128, 256) + + +@needs_nax +@pytest.mark.parametrize("nseg", [1, 2, 6, 13]) +@pytest.mark.parametrize("q_len", [1, 4, 5]) +@pytest.mark.parametrize("peaked", [False, True]) +def test_segments_match_the_fp32_reference_as_well_as_the_stock_sdpa_does(nseg, q_len, peaked) -> None: + total = 12000 + q, k, v = _data(total, q_len, peaked) + out = S.sdpa_nax_flash_dsplit_segments(queries=q, segments=_split(k, v, _lens(total, nseg)), scale=SCALE) + assert out is not None and bool(mx.all(mx.isfinite(out)).item()) + ref = _ref(q, k, v, total) + err = float(mx.abs(out.astype(mx.float32) - ref).max()) + mask = (mx.arange(total)[None, :] <= (total - q_len + mx.arange(q_len))[:, None]) + stock = mx.fast.scaled_dot_product_attention( + q, k[:, :, :total], v[:, :, :total], scale=SCALE, mask=mask[None, None] + ) + err_stock = float(mx.abs(stock.astype(mx.float32) - ref).max()) + assert err <= max(2.5 * err_stock, 1e-3), (err, err_stock) + + +@needs_nax +def test_head_dim_128_one_segment_does_not_go_through_the_contiguous_dispatcher() -> None: + q, k, v = _data(9000, 4) + before = dict(S._dsplit.nax_flash_dsplit_dispatch_counts) + out = S.sdpa_nax_flash_dsplit_segments(queries=q, segments=[(k, v, 9000)], scale=SCALE) + assert out is not None + assert dict(S._dsplit.nax_flash_dsplit_dispatch_counts) == before # the contiguous kernel stays 256 only + assert float(mx.abs(out.astype(mx.float32) - _ref(q, k, v, 9000)).max()) < 5e-3 + + +@needs_nax +def test_a_reference_shorter_than_its_buffer_reads_only_its_rows() -> None: + total, q_len = 9000, 4 + q, k, v = _data(total, q_len) + lens = _lens(total, 3) + segs = _split(k, v, lens) + junk = mx.ones((1, HK, 300, D), mx.bfloat16) * 9 + segs[0] = (mx.concatenate([segs[0][0], junk], axis=2), mx.concatenate([segs[0][1], junk], axis=2), lens[0]) + out = S.sdpa_nax_flash_dsplit_segments(queries=q, segments=segs, scale=SCALE) + assert float(mx.abs(out.astype(mx.float32) - _ref(q, k, v, total)).max()) < 5e-3 + + +def test_prefill_lse_falls_back_to_the_reference_route_when_the_kernel_refuses_a_shape(monkeypatch) -> None: + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + monkeypatch.setattr(module, "_SDPA_LSE", True) + module.route_counts.clear() + + def refuses(*a, **k): + raise ValueError("unsupported head dim for return_lse") + + monkeypatch.setattr(mx.fast, "scaled_dot_product_attention", refuses) + q = mx.random.normal((1, 4, 6, 8)).astype(mx.bfloat16) + keys = mx.random.normal((1, 2, 40, 8)).astype(mx.bfloat16) + out, lse = module.segment_sdpa_lse(q, keys, keys, 40, scale=0.35, causal=False) + assert out.shape == (1, 4, 6, 8) and lse.shape == (1, 4, 6) + assert module.route_counts.get("sdpa_lse_kernel_unsupported") == 1 + assert module.route_counts.get("sdpa_lse_reference") == 1 + finally: + mx.set_default_device(previous) + + +@needs_nax +def test_attention_layer_at_head_dim_128_matches_the_stock_cache_through_turns_and_rollback(monkeypatch) -> None: + """Qwen3-style attention (32 query / 8 KV heads, head_dim 128) through the attention_split ladder.""" + from types import SimpleNamespace + + from mlx.utils import tree_map + from mlx_lm.models.cache import KVCache + from mlx_lm.models.qwen3_next import Qwen3NextAttention + + from mtplx.attention_split import configure_split_full_attention + from mtplx.segmented_kv import SegmentedKVCache + + for name in ("MTPLX_GQA_PACKED_SDPA", "MTPLX_NAX_FLASH_ROUTE"): + monkeypatch.setenv(name, "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + hidden = 128 + args = SimpleNamespace( + hidden_size=hidden, num_attention_heads=HQ, num_key_value_heads=HK, head_dim=D, attention_bias=False, + rms_norm_eps=1e-6, partial_rotary_factor=0.25, rope_theta=10000.0, rope_scaling=None, + max_position_embeddings=65536, + ) + mx.random.seed(0) + attn = Qwen3NextAttention(args) + attn.update(tree_map(lambda p: p.astype(mx.bfloat16), attn.parameters())) + mx.eval(attn.parameters()) + configure_split_full_attention(SimpleNamespace(model=SimpleNamespace(layers=[SimpleNamespace(is_linear=False, self_attn=attn)]))) + + def x(seed, n): + return (mx.random.normal((1, n, hidden), key=mx.random.key(seed)) * 0.5).astype(mx.bfloat16) + + def run(cache): + outs = [attn(x(1, 9000), mask="causal", cache=cache)] + if isinstance(cache, SegmentedKVCache): + state = cache.state + cache = SegmentedKVCache() + cache.state = state + outs.append(attn(x(2, 300), mask="causal", cache=cache)) + for seed, window in ((4, 4), (5, 4), (6, 9), (7, 1)): + outs.append(attn(x(seed, window), mask="causal" if window > 1 else None, cache=cache)) + cache.trim(max(1, window - 2)) + mx.eval(outs) + return outs, cache + + stock, _ = run(KVCache()) + module.route_counts.clear() + seg, cache = run(SegmentedKVCache()) + assert cache.segment_count >= 2 + assert not any(k.startswith("gather_q") and k[8:].isdigit() for k in module.route_counts), module.route_counts + for a, b in zip(stock, seg): + top = float(mx.maximum(mx.abs(a).max(), mx.abs(b).max()).item()) + ulp = 2.0 ** (math.floor(math.log2(top)) - 7) + assert float(mx.abs(a.astype(mx.float32) - b.astype(mx.float32)).max().item()) / ulp <= 4 From b7144b27c5b8f7f243c0639a911d4380fd662a66 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 17:08:14 +0200 Subject: [PATCH 17/37] Segmented KV support check: read head_dim from the attention scale (mlx-lm Qwen3/Llama attention has no head_dim attribute) --- mtplx/segmented_kv.py | 2 ++ tests/test_segmented_kv_support.py | 8 ++++++++ 2 files changed, 10 insertions(+) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index e166f20f6..bf9d4b432 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -102,6 +102,8 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No reasons.append("attention_not_hooked") for attn in layers: head_dim = getattr(attn, "head_dim", None) + if not head_dim and getattr(attn, "scale", None): + head_dim = round(float(attn.scale) ** -2) # mlx-lm Qwen3/Llama attention: scale = head_dim ** -0.5 q_heads = getattr(attn, "num_attention_heads", None) or getattr(attn, "n_heads", None) kv_heads = getattr(attn, "num_key_value_heads", None) or getattr(attn, "n_kv_heads", None) if not (head_dim and q_heads and kv_heads): diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index 86d6d5d44..cbb24cb12 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -138,3 +138,11 @@ def model(head_dim): stats = configure_split_full_attention(model(256)) assert stats["segmented_kv_supported"] is True assert segmented_kv_enabled() is True + + +def test_head_dim_is_read_from_the_scale_when_the_module_has_no_attribute() -> None: + """mlx-lm's Qwen3 attention (head_dim 128) keeps n_heads, n_kv_heads and scale only.""" + attn = SimpleNamespace(n_heads=32, n_kv_heads=8, scale=128**-0.5, _mtplx_gqa_packed_sdpa_enabled=True) + verdict = evaluate_model_support([attn]) + assert verdict["supported"], verdict + assert verdict["head_dims"] == [128] and verdict["gqa"] == [4] From 07571bd62d38a48292f326cd28ac9e16f5d28bfe Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 17:11:29 +0200 Subject: [PATCH 18/37] 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 reachable by that body --- mtplx/attention_split.py | 48 ++++++++++++++++++----- tests/test_segmented_kv_support.py | 63 ++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 10 deletions(-) diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index b77de93cf..d8bcb2469 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -345,6 +345,18 @@ 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 _install_split_attention_hook(attn: Any) -> bool: cls = type(attn) if getattr(cls, "_mtplx_split_full_attention_installed", False): @@ -380,27 +392,38 @@ 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): + gated_q_proj = _attention_has_gated_q_proj(self) + if not gated_q_proj and not ( + isinstance(cache, SegmentedKVCache) and _attention_has_plain_q_proj(self) + ): + # Only the gated (Qwen3-Next style) attention goes through this body, and, for + # MTPLX_SEGMENTED_KV, a plain q/k-norm attention (Qwen3, head_dim 128) 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) + keys.reshape(B, L, n_kv_heads, -1) ).transpose(0, 2, 1, 3) - values = values.reshape(B, L, self.num_key_value_heads, -1).transpose( + values = values.reshape(B, L, n_kv_heads, -1).transpose( 0, 2, 1, @@ -849,6 +872,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 @@ -947,7 +972,10 @@ def configure_split_full_attention( ) 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))) + 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_plain_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 diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index cbb24cb12..bf3246028 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -146,3 +146,66 @@ def test_head_dim_is_read_from_the_scale_when_the_module_has_no_attribute() -> N verdict = evaluate_model_support([attn]) assert verdict["supported"], verdict assert verdict["head_dims"] == [128] and verdict["gqa"] == [4] + + +def _qwen3_attention(head_dim=32, q=4, kv=2, hidden=64): + from mlx_lm.models.qwen3 import Attention, ModelArgs + + args = ModelArgs( + model_type="qwen3", hidden_size=hidden, num_hidden_layers=1, intermediate_size=64, num_attention_heads=q, + rms_norm_eps=1e-6, vocab_size=64, num_key_value_heads=kv, max_position_embeddings=4096, rope_theta=10000.0, + head_dim=head_dim, tie_word_embeddings=False, + ) + mx.random.seed(0) + return Attention(args) + + +def test_plain_qwen3_attention_is_hooked_for_segmented_caches_only_and_matches_the_stock_forward() -> None: + """A plain q/k-norm attention (no gate) goes through the hook body only on a segmented cache; with a + stock cache it is the unchanged mlx-lm forward, so models without the switch keep their code.""" + from mlx_lm.models.cache import KVCache + + from mtplx.attention_split import _attention_has_gated_q_proj, _attention_has_plain_q_proj, _install_split_attention_hook + from mtplx.segmented_kv import SegmentedKVCache + + attn = _qwen3_attention() + assert not _attention_has_gated_q_proj(attn) and _attention_has_plain_q_proj(attn) + + def run(cache): + outs = [] + for seed, n in ((1, 40), (2, 9), (3, 1)): + x = mx.random.normal((1, n, 64), key=mx.random.key(seed)) + outs.append(attn(x, mask="causal" if n > 1 else None, cache=cache)) + mx.eval(outs) + return outs + + before = run(KVCache()) + _install_split_attention_hook(attn) + attn._mtplx_split_full_attention_enabled = True + after_stock = run(KVCache()) + for a, b in zip(before, after_stock): + assert mx.array_equal(a, b).item() # stock cache: the unchanged forward + cache = SegmentedKVCache() + seg = run(cache) + cache.seal() + for a, b in zip(before, seg): + assert float(mx.abs(a - b).max().item()) < 1e-4 # same math on the gathered rows (CPU: no kernel) + assert cache.offset == 50 + + +def test_the_hook_covers_plain_attention_in_the_support_check() -> None: + from types import SimpleNamespace + + from mtplx.attention_split import configure_split_full_attention + + attn = _qwen3_attention(head_dim=128, q=8, kv=2, hidden=128) + model = SimpleNamespace(model=SimpleNamespace(layers=[SimpleNamespace(is_linear=False, self_attn=attn)])) + import os + + os.environ["MTPLX_GQA_PACKED_SDPA"] = "1" + try: + stats = configure_split_full_attention(model) + finally: + os.environ.pop("MTPLX_GQA_PACKED_SDPA") + assert stats["segmented_kv_supported"] is True + assert model_support()["head_dims"] == [128] and model_support()["gqa"] == [4] From 6da2451f9773146daec72ee25237a56bb31573e5 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 23:28:35 +0200 Subject: [PATCH 19/37] 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 gathers; several segments take the kernels at any length --- CHANGELOG.md | 2 +- mtplx/attention_split.py | 6 +- mtplx/cache_state.py | 4 +- mtplx/segmented_kv.py | 145 +++++-------------- tests/test_segmented_kv_attention.py | 35 +++++ tests/test_segmented_kv_bank.py | 2 +- tests/test_segmented_kv_boundary.py | 199 --------------------------- tests/test_segmented_kv_one_path.py | 153 ++++++++++++++++++++ tests/test_segmented_kv_ssd.py | 2 - tests/test_segmented_kv_support.py | 1 - 10 files changed, 234 insertions(+), 315 deletions(-) delete mode 100644 tests/test_segmented_kv_boundary.py create mode 100644 tests/test_segmented_kv_one_path.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b72cc46e..c19cb0475 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ All notable user-facing changes to MTPLX. The format is based on ### 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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); other calls gather the rows and run the stock route. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. Segments only apply above a context boundary (`MTPLX_SEGMENTED_KV_MIN_TOKENS`, default the compiled-verify ceiling `MTPLX_COMPILED_VERIFY_MAX_CONTEXT`, 0 when the compiled verify is off): below it the stock cache stays so the compiled verify keeps running; a layer converts to one sealed segment after a prefill above the boundary (no copy) or at the next restore, and back to contiguous rows when restored below it. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. The streaming spill hashes each sealed segment's blocks once (the digest is kept on the segment, and the blob is confirmed on disk before it is referenced), so a follow-up turn no longer rereads and rehashes the whole history. The tiered merge of small recent segments also merges segments that bank snapshots hold (before, nothing ever merged because every turn's snapshot holds its segments), which keeps a conversation at a handful of segments (3 to 5 over 12 follow-up turns at 50K) instead of two more per turn, below the 12-segment fused launch. An SSD restore builds the request's dense cache layout instead of the paged default (with the switch on). `/health` gets a `segmented_kv` block (segments and sealed bytes per entry, merges, route and SSD counters), absent when the switch is off. Not covered: the compiled verify (segmented caches verify eager), paged and quantized KV. +- **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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); a call over several segments that no kernel takes gathers the rows and runs the stock route (counter `gather_q*`; none fires in a normal turn). Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. With the switch on the full-attention layers are segment caches from the first token: one path, no context fence (an earlier version kept the stock cache below `MTPLX_COMPILED_VERIFY_MAX_CONTEXT` so the compiled verify could run; measured on Qwen3.8 27B at 8K to 32K, segments with the eager verify are 3 to 12 % faster per verify and lower the peak memory by 1 to 4.5 GiB at every context, so the fence and its `MTPLX_SEGMENTED_KV_MIN_TOKENS` setting are gone). A cache with one segment (a cold prompt) runs the stock route on that segment's rows, a zero-copy view; calls over several segments take the kernels at any length. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. The streaming spill hashes each sealed segment's blocks once (the digest is kept on the segment, and the blob is confirmed on disk before it is referenced), so a follow-up turn no longer rereads and rehashes the whole history. The tiered merge of small recent segments also merges segments that bank snapshots hold (before, nothing ever merged because every turn's snapshot holds its segments), which keeps a conversation at a handful of segments (3 to 5 over 12 follow-up turns at 50K) instead of two more per turn, below the 12-segment fused launch. An SSD restore builds the request's dense cache layout instead of the paged default (with the switch on). `/health` gets a `segmented_kv` block (segments and sealed bytes per entry, merges, route and SSD counters), absent when the switch is off. Not covered: the compiled verify (segmented caches verify eager; keep `MTPLX_VERIFY_ASYNC_CHUNK_LAYERS=8`, it still helps by 4 to 19 % with segments), paged and quantized KV. ### Fixed diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index d8bcb2469..3126e9b1c 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -495,7 +495,11 @@ def split_call( ) segmented_route = "nax_flash_segments" if segmented_output is None: - segmented_count(f"gather_q{int(queries.shape[2]) if int(queries.shape[2]) <= 32 else 'prefill'}") + q_rows = int(queries.shape[2]) + segmented_count( + f"{'single' if cache.segment_count == 1 else 'gather'}_q" + f"{q_rows if q_rows <= 32 else 'prefill'}" + ) view = gathered_view(cache) keys, values = view.keys, view.values cache = view diff --git a/mtplx/cache_state.py b/mtplx/cache_state.py index 4454237c5..551b73ff4 100644 --- a/mtplx/cache_state.py +++ b/mtplx/cache_state.py @@ -5026,8 +5026,8 @@ def restore_cache( from .segmented_kv import adapt_layer_for_restore, segmented_kv_enabled if segmented_kv_enabled(): - # MTPLX_SEGMENTED_KV: the layer takes the layout its rows call for (above the - # compiled-verify boundary segments, below it the stock cache). + # MTPLX_SEGMENTED_KV: a segment cache takes either snapshot kind; an + # untagged stock layer gets a segmented snapshot as rows. entry, state = adapt_layer_for_restore(cache, idx, state) if state is not None: install_view = not clone_states and _is_trimmable(entry) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index bf9d4b432..65062fcbf 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -22,9 +22,12 @@ 2026-10-07 phase 1 measurement) and a branch from the same history costs 0. Attention over the segments: ``mtplx/kernels/sdpa_segmented.py`` (exact split-softmax merge, -one launch for up to 12 segments) for the verify/decode window; everything else gathers the -rows into one contiguous array per layer and call (``gather``, correct and slower for long -histories) and runs the stock route on it. The prefill over existing segments has its own +one launch for up to 12 segments) for the verify/decode window. A cache with a single segment +(a cold prompt, or a conversation not yet sealed) runs the stock route on that segment's rows, +a zero-copy view of the buffer the stock cache would have (counter ``single_q``). Only +a call over several segments that no kernel takes (a masked or oversized window, a missing +logsumexp kernel for prefill) gathers the rows into one contiguous array (``gather_q``, +correct and slower for long histories) and runs the stock route on it. The prefill over existing segments has its own interface, ``attend_segments_lse`` (per segment ``(out, lse)`` plus one merge), so that an MLX kernel which returns the logsumexp can replace ``segment_sdpa_lse``. @@ -763,7 +766,7 @@ def decode_segments_attention( Engages under the same conditions as the packed route over a dense cache, with the total length where that route reads the buffer capacity: MTPLX_NAX_FLASH_ROUTE on, a verify - window of 2 to 8 rows, total length at least ``packed_threshold``, a causal or no mask. + window of 1 to 8 rows, a causal or no mask, and total length at least ``packed_threshold`` when the cache is one segment (several segments take the kernel at any length). The head-dim-split kernel serves windows of up to 5 rows at GQA 6 (wide and dsplit4 are not in this tree); a longer window, an unsupported shape or a kernel bail returns None and the call runs on the gathered rows through the stock ladder. @@ -781,7 +784,10 @@ def decode_segments_attention( q_len = int(queries.shape[2]) if not 1 <= q_len <= VERIFY_WINDOW_MAX: return None - if cache.offset < int(packed_threshold): + if cache.segment_count == 1 and cache.offset < int(packed_threshold): + # One segment below the packed threshold: the stock route over a zero-copy view, + # exactly what the stock cache does there (several segments always take the kernel, + # the alternative would be a gather). return None if not (mask is None or (isinstance(mask, str) and mask == "causal")): return None @@ -828,108 +834,40 @@ def is_segmented(entry: Any) -> bool: return isinstance(entry, SegmentedKVCache) -def segmented_kv_min_tokens() -> int: - """Context length above which a layer is kept as segments (MTPLX_SEGMENTED_KV). - - The segment cache cannot be promoted to the compiled verify (graphbank: - ``promotion_failure:segmented_kv_cache``), so below the compiled-verify context ceiling the - stock cache stays: production runs compiled there and the switch would only slow it down. - ``MTPLX_SEGMENTED_KV_MIN_TOKENS`` sets the boundary explicitly (0 = always segmented). - Otherwise it is the compiled-verify ceiling itself (``MTPLX_COMPILED_VERIFY_MAX_CONTEXT``, - 32768 in the turbo profile, 6144 without a profile), read per call; with the compiled verify - off the boundary is 0, and with the compiled verify unlimited (ceiling 0) it is infinite, - because a compiled verify that runs at every length leaves the segments nothing to win. - A layer is segmented when its rows are strictly above the boundary, the same comparison the - compiled verify uses to fall back to eager. - """ - raw = os.environ.get("MTPLX_SEGMENTED_KV_MIN_TOKENS", "").strip() - if raw: - try: - return max(0, int(raw)) - except ValueError: - pass - from .graphbank import _compiled_verify_max_context, compiled_verify_mode - - if compiled_verify_mode() == "off": - return 0 - ceiling = _compiled_verify_max_context() - return ceiling if ceiling > 0 else (1 << 62) - - def _is_stock_plain_kv(entry: Any) -> bool: return type(entry).__name__ == "KVCache" and getattr(entry, "_idx", None) is None -def _segmentable(entry: Any) -> bool: - """A stock full-attention layer of the target cache (tagged at install; MTP caches never are).""" - return bool(getattr(entry, "_mtplx_segmentable", False)) and _is_stock_plain_kv(entry) - - -def _segmented_from_stock(entry: Any) -> SegmentedKVCache: - """The stock cache's rows as one sealed segment, without a copy. - - The segment views the stock buffers and is never written again; the stock object is dropped - by the caller, so nothing writes into those buffers any more. - """ - new = SegmentedKVCache(step=int(getattr(entry, "step", 256))) - rows = int(getattr(entry, "offset", 0) or 0) - if rows and entry.keys is not None: - new.state = (entry.keys[..., :rows, :], entry.values[..., :rows, :]) - return new - - def adapt_layer_for_restore(cache: list[Any], idx: int, state: Any) -> tuple[Any, Any]: - """Pick the layout of layer ``idx`` for a restore of ``state``; returns (entry, state). + """Prepare layer ``idx`` for a restore of ``state``; returns (entry, state). - Above ``segmented_kv_min_tokens()`` the layer is a segment cache (a stock snapshot becomes - one sealed segment, no copy); at or below it the layer is a stock cache (a segmented - snapshot is gathered into contiguous rows), so a conversation that crosses the boundary - changes layout at its next restore. Layers that were never tagged (MTP caches, recurrent - layers) are left alone. The entry is replaced in ``cache`` when its class changes. + A segment cache takes either snapshot kind (a stock snapshot, from the SSD tier or a stock + bank entry, becomes its rows). A stock layer that cannot hold segments (a layer that was + never tagged: MTP caches, recurrent layers, or one that already had rows at install) gets a + segmented snapshot as contiguous rows. The entry is never replaced: with the switch on, the + layout is decided once, when the cache is built (``install_segmented_attention_kv_cache``). """ entry = cache[idx] - if not (isinstance(entry, SegmentedKVCache) or _segmentable(entry)): + if isinstance(entry, SegmentedKVCache): + if not isinstance(state, SegmentedKVState) and isinstance(state, (tuple, list)) and len(state) == 2 and state[0] is not None: + _count("restored_stock_as_segment") return entry, state if isinstance(state, SegmentedKVState): - rows = state.rows - elif isinstance(state, (tuple, list)) and len(state) == 2 and state[0] is not None: - rows = int(state[0].shape[2]) - else: - return entry, state - floor = segmented_kv_min_tokens() - want = rows > floor if floor > 0 else isinstance(entry, SegmentedKVCache) - if want and not isinstance(entry, SegmentedKVCache): - if not isinstance(state, SegmentedKVState): - _count("restored_stock_as_segment") # a stock snapshot (SSD restore, stock bank entry) - entry = SegmentedKVCache(step=int(getattr(entry, "step", 256))) - entry._mtplx_segmentable = True - cache[idx] = entry - elif not want and isinstance(entry, SegmentedKVCache): - from mlx_lm.models.cache import KVCache - - entry = KVCache() - entry._mtplx_segmentable = True - cache[idx] = entry - if isinstance(state, SegmentedKVState) and not isinstance(entry, SegmentedKVCache): state = state.to_arrays() return entry, state def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = None) -> dict[str, int | str]: - """Choose the layout of the full-attention KV layers of a target cache list. - - Plain dense ``KVCache`` layers are tagged as segmentable. A layer without rows becomes a - segment cache straight away when the boundary is 0 (``segmented_kv_min_tokens``); otherwise - it stays stock (compiled verify below the boundary) and is converted after its prefill when - its rows are above the boundary (this function runs again from ``repage_target_prefill_cache``) - or at the next restore (``adapt_layer_for_restore``). A stock layer with rows above the - boundary becomes one sealed segment (no copy). Recurrent layers, rotating/indexed caches and + """Make the empty full-attention KV layers of a target cache list segment caches. + + Plain dense ``KVCache`` layers without rows become a ``SegmentedKVCache`` from the first + token (one path with the switch on; there is no context boundary below which the stock + cache stays). Recurrent layers, rotating/indexed caches, layers that already hold rows and caches of other lists (MTP) are left alone. """ - stats: dict[str, int | str] = {"enabled": 1, "mode": "segmented", "entries": 0, "skipped": 0, "converted": 0} + stats: dict[str, int | str] = {"enabled": 1, "mode": "segmented", "entries": 0, "skipped": 0} from .cache_state import _is_trimmable - floor = segmented_kv_min_tokens() for idx, entry in enumerate(cache or []): if entry is None: stats["skipped"] = int(stats["skipped"]) + 1 @@ -937,27 +875,19 @@ def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = if isinstance(entry, SegmentedKVCache): stats["entries"] = int(stats["entries"]) + 1 continue - if not _is_trimmable(entry) or not _is_stock_plain_kv(entry): - stats["skipped"] = int(stats["skipped"]) + 1 - continue rows = int(getattr(entry, "offset", 0) or 0) - if rows == 0 and getattr(entry, "keys", None) is None: - if floor == 0: - new = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) - new._mtplx_segmentable = True - cache[idx] = new - stats["entries"] = int(stats["entries"]) + 1 - else: - entry._mtplx_segmentable = True - stats["skipped"] = int(stats["skipped"]) + 1 - continue - if getattr(entry, "_mtplx_segmentable", False) and rows > floor > 0: - cache[idx] = _segmented_from_stock(entry) - cache[idx]._mtplx_segmentable = True + if ( + _is_trimmable(entry) + and _is_stock_plain_kv(entry) + and rows == 0 + and getattr(entry, "keys", None) is None + ): + new = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) + new._mtplx_segmentable = True + cache[idx] = new stats["entries"] = int(stats["entries"]) + 1 - stats["converted"] = int(stats["converted"]) + 1 - continue - stats["skipped"] = int(stats["skipped"]) + 1 + else: + stats["skipped"] = int(stats["skipped"]) + 1 return stats @@ -1098,7 +1028,6 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: "enabled": True, "requested": True, "model_support": model_support(), - "min_tokens": segmented_kv_min_tokens(), "ssd": segmented_ssd_enabled(), "prefill_route": prefill_route(), "entries": per_entry, diff --git a/tests/test_segmented_kv_attention.py b/tests/test_segmented_kv_attention.py index bc3f58b34..f962c2307 100644 --- a/tests/test_segmented_kv_attention.py +++ b/tests/test_segmented_kv_attention.py @@ -159,3 +159,38 @@ def run(cache): assert module.route_counts.get(f"chunked_q{window}", 0) + module.route_counts.get(f"fused_q{window}", 0) == 2 for a, b in zip(stock, seg): assert _top_ulps(a, b) <= 4 + + +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +@pytest.mark.parametrize("cold_rows", [3000, 9000]) +def test_no_gather_route_fires_in_a_cold_turn_or_in_follow_up_turns(lane, cold_rows) -> None: + """One path with the switch on: a cold prompt (below and above the packed threshold) and the + follow-up turns after it never gather the history. A single-segment call runs the stock + route over a zero-copy view (``single_q*``); several segments take the kernels.""" + import mtplx.segmented_kv as module + + module.route_counts.clear() + cache = SegmentedKVCache() + for seed, n in ((1, cold_rows // 2), (2, cold_rows - cold_rows // 2)): # cold prefill, two chunks + lane(_x(seed, n), mask="causal", cache=cache) + for seed in (3, 4): + lane(_x(seed, 4), mask="causal", cache=cache) # cold verify windows + cache.trim(2) + lane(_x(5, 1), mask=None, cache=cache) # cold decode + mx.eval(cache.keys) + assert not any(k.startswith("gather") for k in module.route_counts), module.route_counts + assert any(k.startswith(("single_q", "fused_q")) for k in module.route_counts) + for turn in range(3): # follow-up turns on restored caches + state = cache.state + cache = SegmentedKVCache() + cache.state = state + lane(_x(10 + turn, 300), mask="causal", cache=cache) # prefill over history + for seed in (20 + turn, 30 + turn): + lane(_x(seed, 4), mask="causal", cache=cache) + cache.trim(2) + lane(_x(40 + turn, 9), mask="causal", cache=cache) # context-copy window + cache.trim(5) + lane(_x(50 + turn, 1), mask=None, cache=cache) + mx.eval(cache.keys) + assert not any(k.startswith("gather") for k in module.route_counts), module.route_counts + assert module.route_counts.get("prefill_lse", 0) >= 1 diff --git a/tests/test_segmented_kv_bank.py b/tests/test_segmented_kv_bank.py index a845cf7e9..37c2779f8 100644 --- a/tests/test_segmented_kv_bank.py +++ b/tests/test_segmented_kv_bank.py @@ -278,7 +278,7 @@ def test_health_block_is_absent_with_the_switch_off_and_reports_segments_with_it assert block["enabled"] is True and block["entries"] assert block["max_segments"] >= 2 assert block["sealed_bytes_unique"] > 0 - assert {"route_counts", "ssd_counts", "merges", "seals", "min_tokens"} <= set(block) + assert {"route_counts", "ssd_counts", "merges", "seals"} <= set(block) assert all(e["segments"] >= 1 and e["sealed_bytes"] > 0 for e in block["entries"]) import json diff --git a/tests/test_segmented_kv_boundary.py b/tests/test_segmented_kv_boundary.py deleted file mode 100644 index c47758165..000000000 --- a/tests/test_segmented_kv_boundary.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Segmented KV only above the compiled-verify boundary (MTPLX_SEGMENTED_KV_MIN_TOKENS).""" - -from __future__ import annotations - -import mlx.core as mx -import pytest -from mlx_lm.models.cache import KVCache - -from mtplx.cache_state import ( - configure_tail_owned_attention_kv_cache, - restore_cache, - snapshot_cache, -) -from mtplx.segmented_kv import ( - SegmentedKVCache, - SegmentedKVState, - segmented_kv_min_tokens, -) - -HEADS = 2 -DIM = 8 - - -@pytest.fixture(autouse=True) -def _env(monkeypatch): - monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") - monkeypatch.delenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", raising=False) - monkeypatch.delenv("MTPLX_COMPILED_VERIFY", raising=False) - monkeypatch.delenv("MTPLX_COMPILED_VERIFY_MAX_CONTEXT", raising=False) - - -def _rows(seed: int, n: int): - k = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed)).astype(mx.bfloat16) - v = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed + 50_000)).astype(mx.bfloat16) - return k, v - - -def _same(a, b) -> bool: - ka, va = a - kb, vb = b - return bool(mx.array_equal(ka, kb).item() and mx.array_equal(va, vb).item()) - - -def _target_cache(n_layers: int = 2) -> list: - cache = [KVCache() for _ in range(n_layers)] - configure_tail_owned_attention_kv_cache(cache) - return cache - - -def test_boundary_follows_the_compiled_verify_settings(monkeypatch) -> None: - assert segmented_kv_min_tokens() == 0 # compiled verify off: segments from the first row - monkeypatch.setenv("MTPLX_COMPILED_VERIFY", "1") - assert segmented_kv_min_tokens() == 6144 # compiled verify's own default ceiling - monkeypatch.setenv("MTPLX_COMPILED_VERIFY_MAX_CONTEXT", "32768") - assert segmented_kv_min_tokens() == 32768 - monkeypatch.setenv("MTPLX_COMPILED_VERIFY_MAX_CONTEXT", "0") - assert segmented_kv_min_tokens() > 10**12 # compiled at every length: never segmented - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "123") - assert segmented_kv_min_tokens() == 123 - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "0") - assert segmented_kv_min_tokens() == 0 - - -def test_boundary_zero_installs_segments_at_creation() -> None: - cache = _target_cache() - assert all(isinstance(c, SegmentedKVCache) for c in cache) - - -def test_below_the_boundary_the_cache_stays_stock(monkeypatch) -> None: - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") - cache = _target_cache() - assert all(type(c) is KVCache for c in cache) - cache[0].update_and_fetch(*_rows(1, 60)) - configure_tail_owned_attention_kv_cache(cache) # what repage runs after a prefill - assert type(cache[0]) is KVCache - - -def test_a_prefill_above_the_boundary_becomes_one_sealed_segment(monkeypatch) -> None: - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") - cache = _target_cache() - control = KVCache() - for seed, n in ((1, 70), (2, 70)): - cache[0].update_and_fetch(*_rows(seed, n)) - control.update_and_fetch(*_rows(seed, n)) - stats_before = type(cache[0]) - configure_tail_owned_attention_kv_cache(cache) - assert stats_before is KVCache and isinstance(cache[0], SegmentedKVCache) - assert type(cache[1]) is KVCache # empty layer: nothing to convert - assert cache[0].offset == 140 and cache[0].segment_count == 1 - assert _same(cache[0].gather(), control.state) - # growth after the conversion goes to a new tail; contents equal the stock cache - cache[0].update_and_fetch(*_rows(3, 33)) - control.update_and_fetch(*_rows(3, 33)) - assert cache[0].segment_count == 2 and _same(cache[0].gather(), control.state) - - -def test_restore_picks_the_layout_by_rows(monkeypatch) -> None: - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") - long_src = _target_cache(1) - long_src[0].update_and_fetch(*_rows(1, 150)) - configure_tail_owned_attention_kv_cache(long_src) - assert isinstance(long_src[0], SegmentedKVCache) - long_snap = snapshot_cache(long_src) - assert isinstance(long_snap.states[0], SegmentedKVState) - - short_src = _target_cache(1) - short_src[0].update_and_fetch(*_rows(2, 80)) - short_snap = snapshot_cache(short_src) - assert type(short_src[0]) is KVCache - - # segmented snapshot, long -> segments (a reference, no copy of the rows) - dest = _target_cache(1) - restore_cache(dest, long_snap, clone_states=False) - assert isinstance(dest[0], SegmentedKVCache) - assert dest[0].attention_segments()[0][0] is long_src[0].attention_segments()[0][0] - # stock snapshot, short -> stays stock - dest = _target_cache(1) - restore_cache(dest, short_snap) - assert type(dest[0]) is KVCache and dest[0].offset == 80 - # stock snapshot above the boundary -> one sealed segment - stock_long = _target_cache(1) - stock_long[0].update_and_fetch(*_rows(5, 130)) - dest = _target_cache(1) - restore_cache(dest, snapshot_cache(stock_long)) - assert isinstance(dest[0], SegmentedKVCache) and dest[0].offset == 130 - assert _same(dest[0].gather(), stock_long[0].state) - # segmented snapshot of a short conversation (boundary raised) -> stock rows - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "500") - dest = _target_cache(1) - restore_cache(dest, long_snap) - assert type(dest[0]) is KVCache and dest[0].offset == 150 - assert _same(dest[0].state, long_src[0].gather()) - - -def test_restored_segments_extend_like_a_stock_cache(monkeypatch) -> None: - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") - control = KVCache() - src = _target_cache(1) - for seed, n in ((1, 90), (2, 90)): - src[0].update_and_fetch(*_rows(seed, n)) - control.update_and_fetch(*_rows(seed, n)) - snap = snapshot_cache(src) # stock, 180 rows - dest = _target_cache(1) - restore_cache(dest, snap) - assert isinstance(dest[0], SegmentedKVCache) - dest[0].update_and_fetch(*_rows(9, 17)) - control.update_and_fetch(*_rows(9, 17)) - assert _same(dest[0].gather(), control.state) - dest[0].trim(5) - control.trim(5) - assert _same(dest[0].gather(), control.state) - - -def test_untagged_caches_are_never_converted(monkeypatch) -> None: - # An MTP history cache is a list of stock KVCache that never went through the install. - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") - mtp = [KVCache()] - mtp[0].update_and_fetch(*_rows(1, 150)) - snap = snapshot_cache(mtp) - dest = [KVCache()] - restore_cache(dest, snap) - assert type(dest[0]) is KVCache - configure_tail_owned_attention_kv_cache(mtp) - assert type(mtp[0]) is KVCache - - -def test_conversion_and_snapshot_do_not_copy_the_history(monkeypatch) -> None: - rows = 4096 - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") - cache = _target_cache(1) - cache[0].update_and_fetch(*_rows(1, rows)) - mx.eval(cache[0].keys, cache[0].values) - history = 2 * HEADS * rows * DIM * 2 - mx.synchronize() - before = mx.get_active_memory() - configure_tail_owned_attention_kv_cache(cache) - snap = snapshot_cache(cache) - cache[0].append_rows(*_rows(2, 64)) - mx.eval(cache[0]._tail.keys, cache[0]._tail.values) - grew = mx.get_active_memory() - before - assert isinstance(snap.states[0], SegmentedKVState) - assert grew < history // 4 - - -def test_a_converted_stock_buffer_is_stored_as_exact_rows(monkeypatch) -> None: - """A stock buffer has capacity beyond its rows; a segment of its rows must not keep a strided view - (MLX copies a non-contiguous kernel input on every launch).""" - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") - cache = _target_cache(1) - cache[0].update_and_fetch(*_rows(1, 300)) # stock capacity: 512 rows - assert cache[0].keys.shape[2] > 300 - configure_tail_owned_attention_kv_cache(cache) - segment = cache[0].attention_segments()[0][0] - assert segment.shape[2] == 300 - flat = mx.contiguous(segment) - mx.eval(flat) - assert mx.array_equal(flat, segment).item() diff --git a/tests/test_segmented_kv_one_path.py b/tests/test_segmented_kv_one_path.py new file mode 100644 index 000000000..88a27e4b4 --- /dev/null +++ b/tests/test_segmented_kv_one_path.py @@ -0,0 +1,153 @@ +"""With MTPLX_SEGMENTED_KV on, the full-attention layers are segment caches from the first token: one path, no context fence.""" + +from __future__ import annotations + +import mlx.core as mx +import pytest +from mlx_lm.models.cache import KVCache + +from mtplx.cache_state import ( + configure_tail_owned_attention_kv_cache, + restore_cache, + snapshot_cache, +) +from mtplx.segmented_kv import ( + SegmentedKVCache, + SegmentedKVState, + route_counts, +) + +HEADS = 2 +DIM = 8 + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") + + +def _rows(seed: int, n: int): + k = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed)).astype(mx.bfloat16) + v = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed + 50_000)).astype(mx.bfloat16) + return k, v + + +def _same(a, b) -> bool: + ka, va = a + kb, vb = b + return bool(mx.array_equal(ka, kb).item() and mx.array_equal(va, vb).item()) + + +def _target_cache(n_layers: int = 2) -> list: + cache = [KVCache() for _ in range(n_layers)] + configure_tail_owned_attention_kv_cache(cache) + return cache + + +def test_segments_from_the_first_token_whatever_the_compiled_verify_settings(monkeypatch) -> None: + for compiled in (None, "1"): + if compiled: + monkeypatch.setenv("MTPLX_COMPILED_VERIFY", compiled) + monkeypatch.setenv("MTPLX_COMPILED_VERIFY_MAX_CONTEXT", "32768") + cache = _target_cache() + assert all(isinstance(c, SegmentedKVCache) for c in cache) + assert all(c.offset == 0 for c in cache) + + +def test_there_is_no_boundary_setting() -> None: + import mtplx.segmented_kv as module + + assert not hasattr(module, "segmented_kv_min_tokens") + assert "min_tokens" not in module.segmented_kv_health([]) + + +def test_a_layer_that_already_holds_rows_stays_stock() -> None: + cache = [KVCache()] + cache[0].update_and_fetch(*_rows(1, 60)) + configure_tail_owned_attention_kv_cache(cache) + assert type(cache[0]) is KVCache + + +def test_segmented_snapshot_restores_as_a_reference() -> None: + src = _target_cache(1) + src[0].update_and_fetch(*_rows(1, 150)) + snap = snapshot_cache(src) + assert isinstance(snap.states[0], SegmentedKVState) + dest = _target_cache(1) + restore_cache(dest, snap, clone_states=False) + assert isinstance(dest[0], SegmentedKVCache) + assert dest[0].attention_segments()[0][0] is src[0].attention_segments()[0][0] + + +def test_stock_snapshot_restores_into_a_segment_cache_and_extends_like_stock() -> None: + control = KVCache() + stock = [KVCache()] + for seed, n in ((1, 90), (2, 90)): + stock[0].update_and_fetch(*_rows(seed, n)) + control.update_and_fetch(*_rows(seed, n)) + dest = _target_cache(1) + before = route_counts.get("restored_stock_as_segment", 0) + restore_cache(dest, snapshot_cache(stock)) + assert route_counts.get("restored_stock_as_segment", 0) == before + 1 + assert isinstance(dest[0], SegmentedKVCache) and dest[0].offset == 180 + assert _same(dest[0].gather(), control.state) + dest[0].update_and_fetch(*_rows(9, 17)) + control.update_and_fetch(*_rows(9, 17)) + assert _same(dest[0].gather(), control.state) + dest[0].trim(5) + control.trim(5) + assert _same(dest[0].gather(), control.state) + + +def test_a_stock_buffer_restored_as_a_segment_is_stored_as_exact_rows() -> None: + """A stock buffer has capacity beyond its rows; the segment must not keep a strided view + (MLX copies a non-contiguous kernel input on every launch).""" + stock = [KVCache()] + stock[0].update_and_fetch(*_rows(1, 300)) # stock capacity: 512 rows + assert stock[0].keys.shape[2] > 300 + dest = _target_cache(1) + restore_cache(dest, snapshot_cache(stock)) + segment = dest[0].attention_segments()[0][0] + assert segment.shape[2] == 300 + flat = mx.contiguous(segment) + mx.eval(flat) + assert mx.array_equal(flat, segment).item() + + +def test_a_segmented_snapshot_restores_as_rows_into_an_untagged_stock_cache() -> None: + src = _target_cache(1) + src[0].update_and_fetch(*_rows(1, 150)) + snap = snapshot_cache(src) + dest = [KVCache()] # an MTP history cache never went through the install + restore_cache(dest, snap) + assert type(dest[0]) is KVCache and dest[0].offset == 150 + assert _same(dest[0].state, src[0].gather()) + + +def test_untagged_caches_are_never_converted() -> None: + mtp = [KVCache()] + mtp[0].update_and_fetch(*_rows(1, 150)) + snap = snapshot_cache(mtp) + dest = [KVCache()] + restore_cache(dest, snap) + assert type(dest[0]) is KVCache + configure_tail_owned_attention_kv_cache(mtp) + assert type(mtp[0]) is KVCache + + +def test_snapshot_does_not_copy_the_history() -> None: + rows = 4096 + cache = _target_cache(1) + cache[0].update_and_fetch(*_rows(1, rows)) + mx.eval(cache[0].keys, cache[0].values) + history = 2 * HEADS * rows * DIM * 2 + mx.synchronize() + before = mx.get_active_memory() + snap = snapshot_cache(cache) + cache[0].append_rows(*_rows(2, 64)) + mx.eval(cache[0]._tail.keys, cache[0]._tail.values) + grew = mx.get_active_memory() - before + assert isinstance(snap.states[0], SegmentedKVState) + assert grew < history // 4 diff --git a/tests/test_segmented_kv_ssd.py b/tests/test_segmented_kv_ssd.py index 0ab6e9be1..61308ca4c 100644 --- a/tests/test_segmented_kv_ssd.py +++ b/tests/test_segmented_kv_ssd.py @@ -215,7 +215,6 @@ def test_a_stock_entry_on_disk_restores_with_the_switch_on_as_segments(monkeypat if (tier.base_dir / "blobs" / next(iter(names.values()))["sha256"][:2]).exists() else None decoded = decode_payload(payload["payload_spec"], (lambda n: blobs[n]) if blobs else None) monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100") from mtplx.cache_state import configure_tail_owned_attention_kv_cache cache = [KVCache(), KVCache(), KVCache()] cache = cache[:2] @@ -318,7 +317,6 @@ def test_ssd_restore_builds_the_requests_cache_layout_with_and_without_the_switc its own tests are in tests/test_ssd_restore_cache_layout.py).""" from mtplx.session_bank import SessionBank - monkeypatch.setenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", "100000") # stock layers stay stock if flag: monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") else: diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index bf3246028..48248b945 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -27,7 +27,6 @@ def _cpu_and_clean(monkeypatch): reset_model_support() monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") - monkeypatch.delenv("MTPLX_SEGMENTED_KV_MIN_TOKENS", raising=False) monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: True) yield reset_model_support() From 01a7db584e46248b294fc8e5d0b90d87bd4b44bf Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Thu, 8 Oct 2026 23:40:21 +0200 Subject: [PATCH 20/37] Segmented KV test: force the LSE prefill route so the no-gather test also holds on stock MLX --- tests/test_segmented_kv_attention.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_segmented_kv_attention.py b/tests/test_segmented_kv_attention.py index f962c2307..8c3dc2c8d 100644 --- a/tests/test_segmented_kv_attention.py +++ b/tests/test_segmented_kv_attention.py @@ -163,12 +163,13 @@ def run(cache): @pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") @pytest.mark.parametrize("cold_rows", [3000, 9000]) -def test_no_gather_route_fires_in_a_cold_turn_or_in_follow_up_turns(lane, cold_rows) -> None: +def test_no_gather_route_fires_in_a_cold_turn_or_in_follow_up_turns(lane, cold_rows, monkeypatch) -> None: """One path with the switch on: a cold prompt (below and above the packed threshold) and the follow-up turns after it never gather the history. A single-segment call runs the stock route over a zero-copy view (``single_q*``); several segments take the kernels.""" import mtplx.segmented_kv as module + monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "lse") # the LSE route also runs (reference) without the MLX kernel module.route_counts.clear() cache = SegmentedKVCache() for seed, n in ((1, cold_rows // 2), (2, cold_rows - cold_rows // 2)): # cold prefill, two chunks From eb59b70235c0658f278664be452f932161b061b2 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 02:40:17 +0200 Subject: [PATCH 21/37] Segmented KV: one log line when the MLX build has no return_lse (prefill over segments then gathers the history) --- mtplx/segmented_kv.py | 14 ++++++++++++-- tests/test_segmented_kv_support.py | 17 +++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 65062fcbf..cb58ebbd7 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -61,6 +61,7 @@ def segmented_kv_requested() -> bool: #: Verdict on the loaded model (``evaluate_model_support``); None until a model was checked. _MODEL_SUPPORT: dict[str, Any] | None = None _UNSUPPORTED_LOGGED = False +_LSE_NOTE_LOGGED = False def segmented_kv_enabled() -> bool: @@ -76,9 +77,10 @@ def model_support() -> dict[str, Any] | None: def reset_model_support() -> None: - global _MODEL_SUPPORT, _UNSUPPORTED_LOGGED + global _MODEL_SUPPORT, _UNSUPPORTED_LOGGED, _LSE_NOTE_LOGGED _MODEL_SUPPORT = None _UNSUPPORTED_LOGGED = False + _LSE_NOTE_LOGGED = False def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | None = None) -> dict[str, Any]: @@ -94,7 +96,7 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No from .kernels.sdpa_segmented import SUPPORTED_HEAD_DIMS, segments_supported from .nax_verify import nax_available - global _MODEL_SUPPORT, _UNSUPPORTED_LOGGED + global _MODEL_SUPPORT, _UNSUPPORTED_LOGGED, _LSE_NOTE_LOGGED layers = list(layers) flags = list(hooked) if hooked is not None else [True] * len(layers) reasons: list[str] = [] @@ -142,6 +144,14 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No + "): keeping the stock KV cache", flush=True, ) + if not reasons and segmented_kv_requested() and prefill_route() == "gather" and not _LSE_NOTE_LOGGED: + _LSE_NOTE_LOGGED = True + print( + "[mtplx] MTPLX_SEGMENTED_KV is on, but this MLX build has no logsumexp output on the fused " + "SDPA (return_lse): a prefill chunk over existing segments gathers the history first " + "(measured about 5 % slower prefill on follow-up turns; decode and verify are unaffected)", + flush=True, + ) return verdict diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index 48248b945..f0f29ed34 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -208,3 +208,20 @@ def test_the_hook_covers_plain_attention_in_the_support_check() -> None: os.environ.pop("MTPLX_GQA_PACKED_SDPA") assert stats["segmented_kv_supported"] is True assert model_support()["head_dims"] == [128] and model_support()["gqa"] == [4] + + +def test_without_the_lse_kernel_one_line_says_prefill_gathers_and_the_feature_stays_on(monkeypatch, capsys) -> None: + monkeypatch.setattr(module, "_SDPA_LSE", False) + monkeypatch.delenv("MTPLX_SEGMENTED_KV_PREFILL", raising=False) + verdict = evaluate_model_support([_attn()]) + assert verdict["supported"] and segmented_kv_enabled() + assert module.prefill_route() == "gather" + assert capsys.readouterr().out.count("no logsumexp output") == 1 + evaluate_model_support([_attn()]) + assert "no logsumexp output" not in capsys.readouterr().out # once per process + + +def test_with_the_lse_kernel_there_is_no_note(monkeypatch, capsys) -> None: + monkeypatch.setattr(module, "_SDPA_LSE", True) + evaluate_model_support([_attn()]) + assert "no logsumexp output" not in capsys.readouterr().out From 3c36a06b2f1435a30ca35158bfeb774f4bc819d1 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 16:21:11 +0200 Subject: [PATCH 22/37] 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 The NOMASK template parameter changed the stock kernel's source and instantiation name without a function: _fused_source replaced the expression anyway. It now replaces the stock `gp <= row_limit` directly by `(nomask_seg || gp <= row_limit)`, so the file the switch-off path runs is byte-identical to upstream. The kernel limits are named constants (MAX_M_ROWS = 32, MAX_Q_LEN = 10) shared with the route gate, and the docstrings no longer point at a gather_attention that does not exist. --- mtplx/kernels/sdpa_nax_flash_dsplit.py | 10 ++------- mtplx/kernels/sdpa_segmented.py | 29 ++++++++++++++++---------- 2 files changed, 20 insertions(+), 19 deletions(-) diff --git a/mtplx/kernels/sdpa_nax_flash_dsplit.py b/mtplx/kernels/sdpa_nax_flash_dsplit.py index 3baa7acfb..627ab45ef 100644 --- a/mtplx/kernels/sdpa_nax_flash_dsplit.py +++ b/mtplx/kernels/sdpa_nax_flash_dsplit.py @@ -34,12 +34,7 @@ # its bails) — the one-line receipt the #459 reports needed. nax_flash_dsplit_dispatch_counts: dict[str, int] = {} -# Template params: InT, PartT, D, QL, GQA_F, NOMASK. -# NOMASK=1 makes every key below ``offset`` visible to every query row (no tail-causal -# mask): a sealed segment of a segmented KV cache sits entirely before the rows being -# verified (mtplx/kernels/sdpa_segmented.py). The kernel derives ``tail_lo`` from the same -# ``offset`` that bounds the keys, so the mask cannot be switched off through the inputs. -# NOMASK=0 (every other caller) compiles to the previous kernel. +# Template params: InT, PartT, D, QL, GQA_F. _SOURCE = r""" constexpr int TK = 32; constexpr int MROWS = 16; @@ -177,7 +172,7 @@ for (short j = 0; j < kElemCols; j++) { const int gp = t0 + hh * 16 + sc.x + j; const float raw = s_raw[hh][i * kElemCols + j]; - const bool vis = live && gp < kv_end && (NOMASK != 0 || gp <= row_limit); + const bool vis = live && gp < kv_end && gp <= row_limit; const float s = vis ? raw * scale : -1e38f; s_p[hh][i * kElemCols + j] = s; tmax[i] = metal::max(tmax[i], s); @@ -371,7 +366,6 @@ def sdpa_nax_flash_dsplit( ("D", d), ("QL", q_len), ("GQA_F", gqa_factor), - ("NOMASK", 0), ], grid=(hk * nthreads, 1, blocks), threadgroup=(nthreads, 1, 1), diff --git a/mtplx/kernels/sdpa_segmented.py b/mtplx/kernels/sdpa_segmented.py index 7273a6fc1..ceefc5a3a 100644 --- a/mtplx/kernels/sdpa_segmented.py +++ b/mtplx/kernels/sdpa_segmented.py @@ -25,10 +25,10 @@ (0.012 to 0.016 with one launch per segment); a verify forward over 6 or 20 segments at 50K and 80K costs 0 to 1.1 % more than the contiguous one. -Routes that cannot take segments yet fall back to ``segmented_kv.gather_attention`` (one -contiguous copy of the history per call, then the stock kernel): the wide (q_len 6 to 16) -and dsplit4 variants are not in this tree, and the head-dim-split kernel serves q_len 1 to 5 -at GQA 6 (``gqa_factor * q_len <= 32``). +The kernel takes a window of ``q_len <= MAX_Q_LEN`` rows with ``gqa_factor * q_len <= +MAX_M_ROWS``; ``segmented_kv.decode_segments_attention`` runs a longer verify window (up to 32 +rows) as sub-windows within both limits. A call no segment route takes (a mask, a kernel +bail) runs the stock route on the gathered rows (``segmented_kv.gathered_view``). """ from __future__ import annotations @@ -52,6 +52,10 @@ #: Metal allows 31 buffer bindings: queries + 2 per segment + meta + scale + 3 outputs. MAX_FUSED_SEGMENTS = 12 +#: Query rows per launch: the head-dim-split kernel's tile takes at most 32 rows of +#: ``gqa_factor * q_len`` and at most 10 verify rows. +MAX_M_ROWS = 32 +MAX_Q_LEN = 10 #: The reduce kernel binds one partial buffer per launch group plus sums, maxs, starts, out. MAX_GROUPS = 20 @@ -75,7 +79,11 @@ def _fused_source(nseg: int) -> str: """The partial kernel over ``nseg`` segment buffers in one launch. Same body as ``sdpa_nax_flash_dsplit._SOURCE``; only the prologue (segment lookup, - buffer selection, per-segment mask flag) and the partial/stat indexing change. + buffer selection, per-segment mask flag) and the partial/stat indexing change, and the + tail-causal mask is skipped for every segment before the last one (``nomask_seg``: a + sealed segment lies entirely before the rows being verified; the stock kernel derives the + mask from the same ``offset`` that bounds the keys, so it cannot be switched off through + the inputs). ``meta`` = bstart[nseg+1] | n[nseg] | cap[nseg] (int32). """ src = _dsplit._SOURCE @@ -96,7 +104,7 @@ def _fused_source(nseg: int) -> str: {sel}""" src = src.replace(_OLD_HEAD, new_head) for old, new in ( - ("(NOMASK != 0 || gp <= row_limit)", "(nomask_seg || gp <= row_limit)"), + ("gp < kv_end && gp <= row_limit", "gp < kv_end && (nomask_seg || gp <= row_limit)"), ("((size_t)hq_row * n_blocks + block_idx) * D + db", "((size_t)hq_row * B_total + gblock) * D + db"), ("sums[hq_row * n_blocks + block_idx]", "sums[hq_row * B_total + gblock]"), ("maxs[hq_row * n_blocks + block_idx]", "maxs[hq_row * B_total + gblock]"), @@ -216,7 +224,7 @@ def sdpa_nax_flash_dsplit_segments( ``keys``/``values`` are (1, Hk, capacity, D) buffers, D in SUPPORTED_HEAD_DIMS, of which the first ``n`` rows are live. Every segment but the last is fully visible; the last holds the rows being verified as its final ``q_len`` rows and keeps the tail-causal mask. None when the contract is not - met (the caller falls back to ``segmented_kv.gather_attention``). + met (the caller then runs the stock route, see ``segmented_kv.decode_segments_attention``). """ if os.environ.get("MTPLX_NAX_FLASH_DSPLIT", "1") == "0": return _bail("env_disabled") @@ -235,9 +243,9 @@ def sdpa_nax_flash_dsplit_segments( if hk <= 0 or hq % hk: return _bail("gqa_heads") gqa_factor = hq // hk - if gqa_factor * q_len > 32: + if gqa_factor * q_len > MAX_M_ROWS: return _bail("m_rows_gt_32") - if q_len < 1 or q_len > 10: + if q_len < 1 or q_len > MAX_Q_LEN: return _bail("q_len") if queries.dtype not in (mx.bfloat16, mx.float16): return _bail("query_dtype") @@ -303,7 +311,6 @@ def sdpa_nax_flash_dsplit_segments( ("D", d), ("QL", q_len), ("GQA_F", gqa_factor), - ("NOMASK", 0), ("NLAST", nseg - 1 if last_group else nseg), ], grid=(hk * nthreads, 1, total_blocks), @@ -341,5 +348,5 @@ def sdpa_nax_flash_dsplit_segments( def segments_supported(q_len: int, gqa_factor: int, head_dim: int) -> bool: """Static part of the contract (no device work), for the route gate in attention_split.""" - return head_dim in SUPPORTED_HEAD_DIMS and 1 <= q_len <= 10 and gqa_factor * q_len <= 32 + return head_dim in SUPPORTED_HEAD_DIMS and 1 <= q_len <= MAX_Q_LEN and gqa_factor * q_len <= MAX_M_ROWS From 7095bd9baf3ef5b501fb15219e9aba944796055b Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 16:21:44 +0200 Subject: [PATCH 23/37] Segmented KV: sub-windows respect the kernel's 10-row limit too (GQA 1 and 2 gathered windows of 11 to 32 rows) The sub-window size was 32 // gqa: 32 rows at GQA 1 and 16 at GQA 2, both above the kernel's q_len <= 10, so segments_supported refused and the window gathered the history. It is now min(MAX_Q_LEN, MAX_M_ROWS // gqa). The measured models (GQA 4, 6, 8) were not affected. Tests: the sub-window sizes at GQA 1 to 8 (CPU) and GQA 1 and 2 windows of 12 and 32 rows against an fp32 reference on the kernel (GPU). Docstring of decode_segments_attention describes the current route (1 to 32 rows, sub-windows, when the stock ladder runs). --- mtplx/segmented_kv.py | 25 +++++++---- tests/test_sdpa_segmented_hd128.py | 70 ++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 9 deletions(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index cb58ebbd7..c38f5dd8a 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -772,17 +772,23 @@ def decode_segments_attention( mask: Any, packed_threshold: int, ) -> mx.array | None: - """Verify/decode attention straight over the segments, or None for the gather fallback. + """Verify/decode attention straight over the segments, or None for the stock ladder. Engages under the same conditions as the packed route over a dense cache, with the total length where that route reads the buffer capacity: MTPLX_NAX_FLASH_ROUTE on, a verify - window of 1 to 8 rows, a causal or no mask, and total length at least ``packed_threshold`` when the cache is one segment (several segments take the kernel at any length). - The head-dim-split kernel serves windows of up to 5 rows at GQA 6 (wide and dsplit4 are not - in this tree); a longer window, an unsupported shape or a kernel bail returns None and the - call runs on the gathered rows through the stock ladder. + window of 1 to ``VERIFY_WINDOW_MAX`` (32) rows, a causal or no mask, and a total length of + at least ``packed_threshold`` when the cache is one segment (several segments take the + kernel at any length). The head-dim-split segment kernel takes a window whole when + ``gqa * q_len <= 32`` and ``q_len <= MAX_Q_LEN`` (10); a longer window runs as sub-windows + of the rows the kernel takes, each over the segments, so the history is never gathered. + None (the call then runs the stock ladder: on a zero-copy view when the cache is one + segment, on gathered rows otherwise) comes from a mask, an unsupported shape or a kernel + bail. """ from .kernel_selfcheck import lane_disabled from .kernels.sdpa_segmented import ( + MAX_M_ROWS, + MAX_Q_LEN, sdpa_nax_flash_dsplit_segments, segments_supported, ) @@ -814,10 +820,11 @@ def decode_segments_attention( if out is not None: _count(f"fused_q{q_len}") return out - # A window the kernel does not take whole (6 to 32 rows without a wide route): sub-windows of - # the rows it does take, each over the keys up to its own last row, so the history is read - # once per sub-window and never gathered. Row i of the window sees keys up to n - Q + i. - sub = max(1, 32 // gqa) + # A window the kernel does not take whole: sub-windows of the rows it does take (both + # limits: gqa * rows <= 32 and rows <= 10, so GQA 1 and 2 split at 10 rows, GQA 4 at 8, + # GQA 6 at 5), each over the keys up to its own last row, so the history is read once per + # sub-window and never gathered. Row i of the window sees keys up to n - Q + i. + sub = max(1, min(MAX_Q_LEN, MAX_M_ROWS // gqa)) if not segments_supported(sub, gqa, d): return None tail_keys, tail_values, tail_n = segments[-1] diff --git a/tests/test_sdpa_segmented_hd128.py b/tests/test_sdpa_segmented_hd128.py index e7ec614bc..ce9e32b41 100644 --- a/tests/test_sdpa_segmented_hd128.py +++ b/tests/test_sdpa_segmented_hd128.py @@ -178,3 +178,73 @@ def run(cache): top = float(mx.maximum(mx.abs(a).max(), mx.abs(b).max()).item()) ulp = 2.0 ** (math.floor(math.log2(top)) - 7) assert float(mx.abs(a.astype(mx.float32) - b.astype(mx.float32)).max().item()) / ulp <= 4 + + +def _low_gqa_cache(hk: int, d: int, window: int, seed: int = 11): + """Two sealed segments plus a tail whose last ``window`` rows are the verify window.""" + from mtplx.segmented_kv import SegmentedKVCache + + cache = SegmentedKVCache() + for i, n in enumerate((5000, 700, 40 + window)): + k = (mx.random.normal((1, hk, n, d), key=mx.random.key(seed + i)) * 0.5).astype(mx.bfloat16) + v = (mx.random.normal((1, hk, n, d), key=mx.random.key(seed + 50 + i)) * 0.5).astype(mx.bfloat16) + cache.append_rows(k, v) + if i < 2: + cache.seal() + return cache + + +@pytest.mark.parametrize("gqa, expected_sub", [(1, 10), (2, 10), (4, 8), (6, 5), (8, 4)]) +@pytest.mark.parametrize("window", [11, 12, 17, 32]) +def test_sub_windows_respect_both_kernel_limits_at_every_gqa(monkeypatch, gqa, expected_sub, window) -> None: + """GQA 1 and 2 used sub-windows of 32 and 16 rows, which the kernel refuses (q_len <= 10): the + window then fell back to a gather. Every sub-window must be one the kernel takes.""" + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") + hk, d = 2, 128 + cache = _low_gqa_cache(hk, d, window) + rows: list[int] = [] + + def fake_kernel(*, queries, segments, scale): + q_len = int(queries.shape[2]) + assert S.segments_supported(q_len, gqa, d), (q_len, gqa) + rows.append(q_len) + return mx.zeros(queries.shape, queries.dtype) + + monkeypatch.setattr(S, "sdpa_nax_flash_dsplit_segments", fake_kernel) + monkeypatch.setattr("mtplx.kernel_selfcheck.lane_disabled", lambda name: False) + q = mx.zeros((1, hk * gqa, window, d), mx.bfloat16) + out = module.decode_segments_attention(q, cache, scale=0.1, mask="causal", packed_threshold=0) + assert out is not None and tuple(out.shape) == tuple(q.shape) + assert sum(rows) == window and max(rows) <= expected_sub + if window > expected_sub: + assert max(rows) == expected_sub + finally: + mx.set_default_device(previous) + + +@needs_nax +@pytest.mark.parametrize("gqa", [1, 2]) +@pytest.mark.parametrize("window", [12, 32]) +def test_low_gqa_wide_windows_run_on_the_segments_and_match_the_reference(monkeypatch, gqa, window) -> None: + """GQA 1 and 2 (4 and 8 query heads over 4 KV heads, head_dim 128) with a window above the + kernel's 10 rows: sub-windows over the segments, no gather, exact against fp32.""" + monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") + hk, d = 4, 128 + cache = _low_gqa_cache(hk, d, window) + hq = hk * gqa + q = (mx.random.normal((1, hq, window, d), key=mx.random.key(99)) * 0.5).astype(mx.bfloat16) + module.route_counts.clear() + out = module.decode_segments_attention(q, cache, scale=1.0 / math.sqrt(d), mask="causal", packed_threshold=0) + assert out is not None + assert module.route_counts.get(f"chunked_q{window}") == 1, module.route_counts + keys, values = cache.gather() + total = int(keys.shape[2]) + kf = mx.repeat(keys.astype(mx.float32), gqa, axis=1) + vf = mx.repeat(values.astype(mx.float32), gqa, axis=1) + s = (q.astype(mx.float32) @ kf.transpose(0, 1, 3, 2)) / math.sqrt(d) + vis = mx.arange(total)[None, :] <= (total - window + mx.arange(window))[:, None] + ref = mx.softmax(mx.where(vis[None, None], s, -mx.inf), axis=-1) @ vf + assert float(mx.abs(out.astype(mx.float32) - ref).max()) < 5e-3 From 4fc08cadcf8e97beb13b4d8ad00dea0c6d421bbf Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 16:22:03 +0200 Subject: [PATCH 24/37] Segmented KV: no verdict on the model means off (fail closed); every runtime load starts without one segmented_kv_enabled() was True while no model had been checked. The verdict is made in configure_split_full_attention, which the Laguna load route skips: such a model got segment caches without the split-attention hook, and its attention gathered every segment on every call. Without a verdict the switch now behaves as off (stock caches, bank and SSD path), /health reports model_not_checked, and runtime.load resets the verdict so a model never inherits the previous load's. The test that encoded the old behaviour now asserts the new one; the other segmented test modules give their caches a supported verdict, as a real load does. The SSD-restore test checks that a request's factory builds segment caches with the switch on. --- mtplx/runtime.py | 6 +++++ mtplx/segmented_kv.py | 19 ++++++++++---- tests/test_segmented_kv.py | 7 +++++ tests/test_segmented_kv_bank.py | 8 ++++++ tests/test_segmented_kv_one_path.py | 3 +++ tests/test_segmented_kv_ssd.py | 17 ++++++++++-- tests/test_segmented_kv_support.py | 40 ++++++++++++++++++++++++++--- 7 files changed, 90 insertions(+), 10 deletions(-) diff --git a/mtplx/runtime.py b/mtplx/runtime.py index 58476ec6d..033c66130 100644 --- a/mtplx/runtime.py +++ b/mtplx/runtime.py @@ -830,6 +830,12 @@ def load( whole_moe_plan = None selfcheck_report = None router_report: dict[str, Any] = {} + # MTPLX_SEGMENTED_KV: the verdict on the model belongs to this load. A route that + # skips configure_split_full_attention (Laguna below) leaves none, and without a + # verdict the segmented cache stays off (mtplx/segmented_kv.py). + from .segmented_kv import reset_model_support + + reset_model_support() # Laguna skips the qwen3-next kernel stack entirely; its own env-gated # fused lanes install right before runtime construction below. if not _is_laguna_s_2_1_mlx_4bit_config(config): diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index c38f5dd8a..2fcc6ca6a 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -65,11 +65,16 @@ def segmented_kv_requested() -> bool: def segmented_kv_enabled() -> bool: - """MTPLX_SEGMENTED_KV, unless the loaded model cannot use segments (then everything behaves - as with the switch off: stock caches, stock bank, stock SSD path).""" + """MTPLX_SEGMENTED_KV, once the loaded model was checked and can use segments. + + Without a verdict everything behaves as with the switch off (stock caches, stock bank, stock + SSD path): the verdict is made in ``configure_split_full_attention``, and a model that loads + without it (Laguna, a script that builds its own model) has no split-attention hook, so its + attention would gather every segment on every call. + """ if not segmented_kv_requested(): return False - return _MODEL_SUPPORT is None or bool(_MODEL_SUPPORT.get("supported", True)) + return _MODEL_SUPPORT is not None and bool(_MODEL_SUPPORT.get("supported", False)) def model_support() -> dict[str, Any] | None: @@ -1013,8 +1018,12 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: if not segmented_kv_requested(): return None if not segmented_kv_enabled(): - # Asked for, but the loaded model cannot use it: say so (and why) instead of staying silent. - return {"enabled": False, "requested": True, "model_support": model_support()} + # Asked for, but the loaded model cannot use it or was never checked: say so (and why) + # instead of staying silent. + support = model_support() + if support is None: + support = {"supported": False, "reasons": ["model_not_checked"]} + return {"enabled": False, "requested": True, "model_support": support} per_entry: list[dict[str, Any]] = [] unique: dict[int, int] = {} for entry in entries: diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py index 584a1c046..1387eae8d 100644 --- a/tests/test_segmented_kv.py +++ b/tests/test_segmented_kv.py @@ -13,6 +13,7 @@ import pytest from mlx_lm.models.cache import KVCache +import mtplx.segmented_kv as segmented_kv_module from mtplx.cache_state import ( configure_tail_owned_attention_kv_cache, restore_cache, @@ -43,6 +44,12 @@ def _no_merge_by_default(monkeypatch): monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") +@pytest.fixture(autouse=True) +def _model_checked(monkeypatch): + # segmented_kv_enabled() needs a verdict on the loaded model (without one it stays off). + monkeypatch.setattr(segmented_kv_module, "_MODEL_SUPPORT", {"supported": True, "reasons": []}) + + def _rows(seed: int, n: int) -> tuple[mx.array, mx.array]: k = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed)).astype(mx.bfloat16) v = mx.random.normal((1, HEADS, n, DIM), key=mx.random.key(seed + 50_000)).astype(mx.bfloat16) diff --git a/tests/test_segmented_kv_bank.py b/tests/test_segmented_kv_bank.py index 37c2779f8..b8f15a1ad 100644 --- a/tests/test_segmented_kv_bank.py +++ b/tests/test_segmented_kv_bank.py @@ -11,8 +11,10 @@ from pathlib import Path import mlx.core as mx +import pytest from mlx_lm.models.cache import KVCache +import mtplx.segmented_kv as segmented_kv_module from mtplx.segmented_kv import SegmentedKVCache, SegmentedKVState from mtplx.session_bank import SessionBank, entry_has_segmented_kv @@ -21,6 +23,12 @@ ROW_BYTES = LAYERS * HEADS * DIM * 2 * 2 # K and V, bf16 +@pytest.fixture(autouse=True) +def _model_checked(monkeypatch): + # segmented_kv_enabled() needs a verdict on the loaded model (without one it stays off). + monkeypatch.setattr(segmented_kv_module, "_MODEL_SUPPORT", {"supported": True, "reasons": []}) + + class Runtime: model_path = Path("models/example") mtp_enabled = False diff --git a/tests/test_segmented_kv_one_path.py b/tests/test_segmented_kv_one_path.py index 88a27e4b4..ff4706d55 100644 --- a/tests/test_segmented_kv_one_path.py +++ b/tests/test_segmented_kv_one_path.py @@ -6,6 +6,7 @@ import pytest from mlx_lm.models.cache import KVCache +import mtplx.segmented_kv as segmented_kv_module from mtplx.cache_state import ( configure_tail_owned_attention_kv_cache, restore_cache, @@ -26,6 +27,8 @@ def _env(monkeypatch): monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") + # segmented_kv_enabled() needs a verdict on the loaded model (without one it stays off). + monkeypatch.setattr(segmented_kv_module, "_MODEL_SUPPORT", {"supported": True, "reasons": []}) def _rows(seed: int, n: int): diff --git a/tests/test_segmented_kv_ssd.py b/tests/test_segmented_kv_ssd.py index 61308ca4c..f9123101b 100644 --- a/tests/test_segmented_kv_ssd.py +++ b/tests/test_segmented_kv_ssd.py @@ -10,6 +10,7 @@ import pytest from mlx_lm.models.cache import KVCache +import mtplx.segmented_kv as segmented_kv_module from mtplx.cache_bank.codec import decode_payload, encode_payload from mtplx.cache_bank.cold_tier import SessionBankColdTier from mtplx.cache_state import CacheSnapshot, restore_cache @@ -23,6 +24,12 @@ HEADS, DIM = 2, 16 +@pytest.fixture(autouse=True) +def _model_checked(monkeypatch): + # segmented_kv_enabled() needs a verdict on the loaded model (without one it stays off). + monkeypatch.setattr(segmented_kv_module, "_MODEL_SUPPORT", {"supported": True, "reasons": []}) + + def _rand(seed, n): rng = np.random.default_rng(seed) return mx.array(rng.standard_normal((1, HEADS, n, DIM)).astype(np.float32)).astype(mx.bfloat16) @@ -314,7 +321,9 @@ def make_cache(self): @pytest.mark.parametrize("flag", ["1", ""]) def test_ssd_restore_builds_the_requests_cache_layout_with_and_without_the_switch(monkeypatch, tmp_path, flag) -> None: """The SSD restore builds the request's cache layout, with or without segments (the fix itself and - its own tests are in tests/test_ssd_restore_cache_layout.py).""" + its own tests are in tests/test_ssd_restore_cache_layout.py): with the switch on that layout is the + segmented one, and the stock entry on disk becomes its rows.""" + from mtplx.cache_state import configure_tail_owned_attention_kv_cache from mtplx.session_bank import SessionBank if flag: @@ -335,9 +344,13 @@ def test_ssd_restore_builds_the_requests_cache_layout_with_and_without_the_switc def factory(): made_by_factory.append(1) - return [KVCache() for _ in range(2)] + layers = [KVCache() for _ in range(2)] + configure_tail_owned_attention_kv_cache(layers) # as rt.make_cache() does for the dense layout + return layers restored = fresh.restore(runtime, tokens + [701, 702], session_id="s", cache_factory=factory) assert restored is not None and restored.restore_mode == "ssd_clone" assert made_by_factory + assert all(isinstance(c, SegmentedKVCache) for c in restored.cache) == bool(flag) assert restored.cache[0].offset == 700 + assert mx.array_equal(restored.cache[0].keys[..., :700, :], cache[0].state[0]).item() diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index f0f29ed34..804d47d6f 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -107,9 +107,43 @@ def test_health_of_a_supported_model_carries_the_verdict() -> None: assert block["enabled"] is True and block["model_support"]["supported"] is True -def test_nothing_is_decided_before_a_model_is_checked() -> None: - assert model_support() is None and segmented_kv_enabled() is True - assert module._MODEL_SUPPORT is None +def test_without_a_verdict_on_the_model_the_feature_stays_off() -> None: + """A model that loads without configure_split_full_attention (Laguna, a script) has no hook: with + segment caches its attention would gather every segment on every call, so no verdict means off.""" + assert model_support() is None and module._MODEL_SUPPORT is None + assert segmented_kv_enabled() is False and segmented_ssd_enabled() is False + + +def test_a_load_route_without_a_verdict_builds_stock_caches_and_health_says_why() -> None: + cache = [KVCache(), KVCache()] + stats = configure_tail_owned_attention_kv_cache(cache) + assert all(type(c) is KVCache for c in cache) + assert stats.get("mode") != "segmented" + block = segmented_kv_health([]) + assert block["enabled"] is False and block["requested"] is True + assert block["model_support"]["reasons"] == ["model_not_checked"] + evaluate_model_support([_attn()]) # the verdict a supported model gets at load + cache = [KVCache()] + configure_tail_owned_attention_kv_cache(cache) + from mtplx.segmented_kv import SegmentedKVCache + + assert isinstance(cache[0], SegmentedKVCache) + + +def test_every_runtime_load_starts_without_a_verdict() -> None: + """runtime.load resets the verdict before the model's own check, so a Laguna load after a + supported model does not inherit that model's verdict.""" + import inspect + + from mtplx import runtime + + evaluate_model_support([_attn()]) + assert segmented_kv_enabled() is True + source = inspect.getsource(runtime) + reset_at = source.index("reset_model_support()") + assert "if not _is_laguna_s_2_1_mlx_4bit_config(config):" in source[reset_at : reset_at + 400] + reset_model_support() + assert segmented_kv_enabled() is False def test_configure_split_full_attention_decides_from_the_hooked_layers(monkeypatch) -> None: From 059b095a993dfb3b47ec728b0fe0b3155cf8733b Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 16:22:18 +0200 Subject: [PATCH 25/37] 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 - Reference counting (KVSegment holders, hold/release, refcount) steered nothing since the merge also merges held segments; removed with its test. Segments are shared by plain references; the module docstring says so. - Unused in production and removed: SegmentedKVState.prefix/segments/__len__, _prefix_refs, SegmentedKVCache.tail_rows/stats() and its per-cache counters (route_counts has them), compact(max_segments=), is_segmented, the _mtplx_segmentable tag. - "Unique bytes over segment id" lives in one place (unique_state_bytes over _unique_segment_bytes), used by the bank's overcount and /health; SegmentedKVCache.gather goes through _gather_refs. - adapt_layer_for_restore(entry, state) -> state; restore_cache reads the switch once per call instead of importing and reading it per layer. - One _env_on() for the module's switches; attention_split asks segmented_kv_requested(); the route counter is public (count_route) and the codec uses its own _count_ssd everywhere; a known segment block and a reuse hit share one block entry in the codec. - Constants instead of literals (VERIFY_WINDOW_MAX in the counter label, MAX_FUSED_SEGMENTS as the default segment limit), and one sentence on why small segmented tensors are concatenated in the codec. - With the switch off, SessionBank.total_nbytes and the admission no longer walk the snapshot trees for segments. - The state getter documents its side effect (it seals the tail). - Tests: the test of the never-shipped MIN_TOKENS setting is gone, the merge test no longer says "unshared", counters are read from route_counts; ruff I001 (attention_split imports, a test) and RUF059 fixed. --- mtplx/attention_split.py | 27 ++-- mtplx/cache_bank/codec.py | 63 ++++---- mtplx/cache_state.py | 16 +- mtplx/segmented_kv.py | 226 ++++++++++------------------ mtplx/server/openai.py | 2 +- mtplx/session_bank.py | 20 +-- tests/test_segmented_kv.py | 40 ++--- tests/test_segmented_kv_bank.py | 7 +- tests/test_segmented_kv_one_path.py | 11 +- tests/test_segmented_kv_support.py | 9 +- 10 files changed, 159 insertions(+), 262 deletions(-) diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index 3126e9b1c..941106452 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -14,20 +14,21 @@ ) from .attention_math import attention_gate from .compile_state import in_compiled_step_body, is_compile_trace_error +from .rope_origin import ( + cache_owns_rotary_origin, + note_unowned_rotary_origin, + rope_offset_of, +) from .segmented_kv import ( - SegmentedKVCache, VERIFY_WINDOW_MAX, + SegmentedKVCache, attend_segments_lse, decode_segments_attention, gathered_view, - _count as segmented_count, - prefill_route as segmented_prefill_route, -) -from .rope_origin import ( - cache_owns_rotary_origin, - note_unowned_rotary_origin, - rope_offset_of, + 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: @@ -465,8 +466,10 @@ def split_call( keys = self.rope(keys, offset=rope_offset) if isinstance(cache, SegmentedKVCache): # Sealed segments plus a tail (MTPLX_SEGMENTED_KV): the new rows go to - # the tail. A verify window attends straight over the segments; every - # other call runs the ladder below on one gathered view of the rows. + # 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 @@ -498,7 +501,7 @@ def split_call( q_rows = int(queries.shape[2]) segmented_count( f"{'single' if cache.segment_count == 1 else 'gather'}_q" - f"{q_rows if q_rows <= 32 else 'prefill'}" + f"{q_rows if q_rows <= VERIFY_WINDOW_MAX else 'prefill'}" ) view = gathered_view(cache) keys, values = view.keys, view.values @@ -1002,7 +1005,7 @@ def configure_split_full_attention( attn._mtplx_split_full_attention_calls = 0 stats["layers"] += 1 stats["exact_gather_layers"] += int(exact_gather_layer) - if _env_enabled("MTPLX_SEGMENTED_KV"): + 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 diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index 5f7fe639b..06337dd6d 100644 --- a/mtplx/cache_bank/codec.py +++ b/mtplx/cache_bank/codec.py @@ -156,9 +156,9 @@ def content_fingerprints(value: Any, *, rows: int | None = None) -> list[int] | def _count_ssd(route: str) -> None: - from mtplx.segmented_kv import _count + from mtplx.segmented_kv import count_route - _count(route) + count_route(route) def _is_segmented_state(value: Any) -> bool: @@ -274,11 +274,12 @@ def encode(self, value: Any) -> Any: return self._encode_tensor(value) if _is_segmented_state(value): # MTPLX_SEGMENTED_KV: the same spec and blobs as the stock (keys, values) tuple of - # the same rows; the history is sliced block by block, never gathered. - from mtplx.segmented_kv import _count, segmented_state_tensors - - _count("ssd_layer_states_encoded") + # the same rows; the history is sliced block by block, never gathered. Below the + # block threshold of _encode_tensor (2 blocks of rows) a tensor is stored as one + # blob, which the codec reads as one mx.array: those few rows are concatenated. + from mtplx.segmented_kv import segmented_state_tensors + _count_ssd("ssd_layer_states_encoded") return { "kind": "tuple", "items": [ @@ -412,38 +413,30 @@ def _encode_tensor_blocks( key = (dtype, tuple(chunk_shape), int(start), block_fps[block_index]) self.fingerprints[name] = key hit = self.reuse.get(key) if self.reuse is not None else None - if hit is not None: - nbytes = int(hit["nbytes"]) - self.reused[name] = dict(hit) - total += nbytes - blocks.append( - { - "name": name, - "start": int(start), - "end": int(end), - "shape": chunk_shape, - "nbytes": nbytes, - } - ) - continue + else: + hit = None seg_rows = value if hasattr(value, "known_digest") and self.known_blob is not None else None - if seg_rows is not None: + if hit is None and seg_rows is not None: + # MTPLX_SEGMENTED_KV spill: a block of a sealed segment that an earlier spill + # already wrote is referenced like a reuse hit (known_blob confirms the blob). known = seg_rows.known_digest(start, end) - known = self.known_blob(known) if known is not None else None - if known is not None: + hit = self.known_blob(known) if known is not None else None + if hit is not None: _count_ssd("ssd_blocks_reused") - self.reused[name] = dict(known) - total += int(known["nbytes"]) - blocks.append( - { - "name": name, - "start": int(start), - "end": int(end), - "shape": [int(dim) for dim in chunk.shape], - "nbytes": int(known["nbytes"]), - } - ) - continue + if hit is not None: + nbytes = int(hit["nbytes"]) + self.reused[name] = dict(hit) + total += nbytes + blocks.append( + { + "name": name, + "start": int(start), + "end": int(end), + "shape": [int(dim) for dim in chunk.shape], + "nbytes": nbytes, + } + ) + continue raw = self._materialize("block", chunk) self.tensors[name] = raw if seg_rows is not None and self.on_block is not None: diff --git a/mtplx/cache_state.py b/mtplx/cache_state.py index 551b73ff4..b6ca84b1c 100644 --- a/mtplx/cache_state.py +++ b/mtplx/cache_state.py @@ -5021,14 +5021,14 @@ def restore_cache( restore_meta_state: bool = True, clone_states: bool = True, ) -> None: - for idx, (entry, state, meta_state) in enumerate(zip(list(cache), snapshot.states, snapshot.meta_states)): - if state is not None: - from .segmented_kv import adapt_layer_for_restore, segmented_kv_enabled - - if segmented_kv_enabled(): - # MTPLX_SEGMENTED_KV: a segment cache takes either snapshot kind; an - # untagged stock layer gets a segmented snapshot as rows. - entry, state = adapt_layer_for_restore(cache, idx, state) + from .segmented_kv import adapt_layer_for_restore, segmented_kv_enabled + + segmented = segmented_kv_enabled() + for entry, state, meta_state in zip(cache, snapshot.states, snapshot.meta_states): + if segmented and state is not None: + # MTPLX_SEGMENTED_KV: a segment cache takes either snapshot kind; a stock + # layer (MTP, recurrent) gets a segmented snapshot as contiguous rows. + state = adapt_layer_for_restore(entry, state) if state is not None: install_view = not clone_states and _is_trimmable(entry) _restore_state_preserving_container(entry, state, clone=not install_view) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 2fcc6ca6a..1e17af0f1 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -33,17 +33,17 @@ GatedDeltaNet (recurrent) layers keep their state objects as they are. -Reference counts. A segment counts its holders (live caches and ``SegmentedKVState`` -snapshots) in a weak set: ``refcount`` is the number of holders still alive, an evicted or -dropped bank entry releases its segments by going away, and the merge policy only merges -segments that no snapshot holds. +Sharing. Segments are shared by plain Python references: a segment lives as long as a cache +or a ``SegmentedKVState`` snapshot references it, so an evicted or dropped bank entry releases +its segments by going away. The merge policy (``SegmentedKVCache.compact``) also merges small +segments that snapshots still hold: the snapshots keep their own pieces and the merged copy is +the only new memory. """ from __future__ import annotations import itertools import os -import weakref from collections.abc import Iterable, Sequence from dataclasses import dataclass from typing import Any @@ -53,9 +53,19 @@ _SEGMENT_IDS = itertools.count(1) +def _env_on(name: str, *, default: bool = False) -> bool: + """A boolean switch: 1/true/yes/on and 0/false/no/off, anything else (or unset) is ``default``.""" + raw = os.environ.get(name, "").strip().lower() + if raw in {"1", "true", "yes", "on"}: + return True + if raw in {"0", "false", "no", "off"}: + return False + return default + + def segmented_kv_requested() -> bool: """MTPLX_SEGMENTED_KV (default off), read on every call: what the operator asked for.""" - return os.environ.get("MTPLX_SEGMENTED_KV", "").strip().lower() in {"1", "true", "yes", "on"} + return _env_on("MTPLX_SEGMENTED_KV") #: Verdict on the loaded model (``evaluate_model_support``); None until a model was checked. @@ -127,7 +137,7 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No reasons.append(f"gqa_{q_heads}_{kv_heads}_unsupported") if not all(getattr(a, "_mtplx_gqa_packed_sdpa_enabled", False) for a in layers): reasons.append("gqa_packed_sdpa_off") - if os.environ.get("MTPLX_NAX_FLASH_ROUTE", "").strip().lower() not in {"1", "true", "yes", "on"}: + if not _env_on("MTPLX_NAX_FLASH_ROUTE"): reasons.append("nax_flash_route_off") if not nax_available(): reasons.append("nax_unavailable") @@ -172,7 +182,8 @@ def _int_env(name: str, default: int) -> int: route_counts: dict[str, int] = {} -def _count(route: str) -> None: +def count_route(route: str) -> None: + """One more call served by ``route`` (``route_counts``, reported on /health).""" route_counts[route] = route_counts.get(route, 0) + 1 @@ -224,14 +235,13 @@ class KVSegment: The tail of a cache is an unsealed segment: only its owner writes it. """ - __slots__ = ("__weakref__", "_holders", "block_digests", "id", "keys", "sealed", "values") + __slots__ = ("block_digests", "id", "keys", "sealed", "values") def __init__(self, keys: mx.array, values: mx.array, *, sealed: bool = False) -> None: self.id = next(_SEGMENT_IDS) self.keys = keys self.values = values self.sealed = bool(sealed) - self._holders: weakref.WeakSet[Any] = weakref.WeakSet() # SSD tier: sha256 and size of the 256-row blocks of this (immutable) segment that a # spill already hashed, keyed (is_values, first_row, last_row) in segment rows. self.block_digests: dict[tuple[bool, int, int], dict[str, Any]] = {} @@ -244,17 +254,6 @@ def capacity(self) -> int: def nbytes(self) -> int: return int(self.keys.nbytes) + int(self.values.nbytes) - @property - def refcount(self) -> int: - """Holders alive right now: live caches and snapshot states.""" - return len(self._holders) - - def hold(self, holder: Any) -> None: - self._holders.add(holder) - - def release(self, holder: Any) -> None: - self._holders.discard(holder) - @dataclass(frozen=True) class SegRef: @@ -267,14 +266,12 @@ class SegRef: class SegmentedKVState: """An immutable snapshot of a segmented cache: references, no buffers of its own. - This is what the session bank stores as the cache snapshot of a segmented layer. It holds - its segments (reference counts) for as long as it lives. + This is what the session bank stores as the cache snapshot of a segmented layer. Its + references keep the segments alive for as long as it lives. """ def __init__(self, refs: Iterable[SegRef]) -> None: self.refs: tuple[SegRef, ...] = tuple(refs) - for ref in self.refs: - ref.segment.hold(self) @property def rows(self) -> int: @@ -283,44 +280,27 @@ def rows(self) -> int: @property def nbytes(self) -> int: """Bytes of the segments this state references, counted as if it were alone.""" - return sum(ref.segment.nbytes for ref in _unique_refs(self.refs)) - - def segments(self) -> list[KVSegment]: - return [ref.segment for ref in self.refs] - - def prefix(self, rows: int) -> SegmentedKVState: - """The first ``rows`` rows as a new state (shortens the last reference, no copy).""" - return SegmentedKVState(_prefix_refs(self.refs, rows)) + return _unique_segment_bytes(self.refs) def to_arrays(self) -> tuple[mx.array, mx.array] | None: """One contiguous (keys, values) pair, the state a stock KV cache would give.""" return _gather_refs(self.refs) - def __len__(self) -> int: - return len(self.refs) - -def _unique_refs(refs: Iterable[SegRef]) -> list[SegRef]: - seen: set[int] = set() - out: list[SegRef] = [] +def _unique_segment_bytes(refs: Iterable[SegRef]) -> int: + """Bytes of the segments behind ``refs``, each segment counted once.""" + seen: dict[int, int] = {} for ref in refs: - if ref.segment.id not in seen: - seen.add(ref.segment.id) - out.append(ref) - return out + seen.setdefault(ref.segment.id, ref.segment.nbytes) + return sum(seen.values()) -def _prefix_refs(refs: Sequence[SegRef], rows: int) -> list[SegRef]: - rows = max(0, int(rows)) - out: list[SegRef] = [] - left = rows - for ref in refs: - if left <= 0: - break - take = min(int(ref.n), left) - out.append(SegRef(ref.segment, take) if take != ref.n else ref) - left -= take - return out +def unique_state_bytes(states: Iterable[Any]) -> int: + """Bytes of the segments referenced by ``states`` (SegmentedKVState objects; anything else is + skipped), each segment counted once however many states share it.""" + return _unique_segment_bytes( + ref for state in states if isinstance(state, SegmentedKVState) for ref in state.refs + ) def _rows_view(array: mx.array, n: int) -> mx.array: @@ -359,10 +339,6 @@ def __init__(self, *, step: int = 256) -> None: self._sealed: list[SegRef] = [] self._tail: KVSegment | None = None self._tail_n = 0 - self.seal_copies = 0 - self.seal_copy_rows = 0 - self.merges = 0 - self.merge_rows = 0 # ---- shape ---------------------------------------------------------------------------- @@ -379,10 +355,6 @@ def offset(self, value: int) -> None: raise ValueError("a segmented KV cache cannot grow its offset without rows") self.trim(current - value) - @property - def tail_rows(self) -> int: - return self._tail_n - @property def segment_count(self) -> int: return len(self._sealed) + (1 if self._tail_n else 0) @@ -399,7 +371,7 @@ def is_trimmable(self) -> bool: @property def nbytes(self) -> int: """Bytes this cache keeps allocated: its unique segments plus the tail buffer.""" - total = sum(ref.segment.nbytes for ref in _unique_refs(self._sealed)) + total = _unique_segment_bytes(self._sealed) if self._tail is not None: total += self._tail.nbytes return total @@ -461,14 +433,10 @@ def attention_segments(self) -> list[tuple[mx.array, mx.array, int]]: def gather(self) -> tuple[mx.array, mx.array] | None: """One contiguous (keys, values) over all rows (lazy; a view when there is one segment).""" - parts = self.attention_segments() - if not parts: - return None - keys = [_rows_view(k, n) for k, _, n in parts] - values = [_rows_view(v, n) for _, v, n in parts] - if len(parts) == 1: - return keys[0], values[0] - return mx.concatenate(keys, axis=2), mx.concatenate(values, axis=2) + refs = list(self._sealed) + if self._tail is not None and self._tail_n: + refs.append(SegRef(self._tail, self._tail_n)) + return _gather_refs(refs) @property def keys(self) -> mx.array | None: @@ -495,8 +463,6 @@ def trim(self, n: int) -> int: ref = self._sealed[-1] if ref.n <= left: self._sealed.pop() - if not any(r.segment is ref.segment for r in self._sealed): - ref.segment.release(self) left -= int(ref.n) else: self._sealed[-1] = SegRef(ref.segment, int(ref.n) - left) @@ -517,33 +483,32 @@ def seal(self) -> bool: if rows == 0: self._tail = None return False - _count("seal") + count_route("seal") if rows == tail.capacity: keys, values = tail.keys, tail.values mx.eval(keys, values) else: - _count("seal_copy") + count_route("seal_copy") keys = _own_copy(tail.keys[..., :rows, :]) values = _own_copy(tail.values[..., :rows, :]) - self.seal_copies += 1 - self.seal_copy_rows += rows self._tail = None self._tail_n = 0 - segment = KVSegment(keys, values, sealed=True) - segment.hold(self) - self._sealed.append(SegRef(segment, rows)) + self._sealed.append(SegRef(KVSegment(keys, values, sealed=True), rows)) return True @property def state(self) -> SegmentedKVState: - """The snapshot of this cache: seals the tail, then the list of references.""" + """The snapshot of this cache: seals the tail, then the list of references. + + Reading it has a side effect: the tail is sealed (a copy of the tail's rows). A caller + about to replace the state does not need to read it first; restores go to fresh caches, + where there is no tail to seal. + """ self.seal() return SegmentedKVState(self._sealed) @state.setter def state(self, value: Any) -> None: - for ref in self._sealed: - ref.segment.release(self) self._sealed = [] self._tail = None self._tail_n = 0 @@ -562,21 +527,19 @@ def state(self, value: Any) -> None: refs = [SegRef(segment, int(keys.shape[2]))] else: raise ValueError(f"unsupported segmented KV state: {type(value).__name__}") - for ref in refs: - ref.segment.hold(self) self._sealed = refs # ---- merge policy --------------------------------------------------------------------- - def compact(self, *, max_segments: int | None = None) -> int: + def compact(self) -> int: """Merge small recent sealed segments among themselves (geometric tiers). Two adjacent sealed segments merge when the older is at most ``MTPLX_SEGMENTED_KV_TIER`` times (default 2) the newer one and the result stays within ``MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS`` (default 16384) rows. The big history segment is never merged here: a full copy would bring the duplicate back while a snapshot, a - subagent or a lease references it. A segment that bank snapshots also hold (reference - count above 1) is merged too: the merged copy is new memory (at most the merged rows, + subagent or a lease references it. A segment that bank snapshots also hold is merged + too: the merged copy is new memory (at most the merged rows, released when the snapshots holding the originals go), and without it nothing would ever merge, because every turn's snapshot holds that turn's segments. The segment count is thereby bounded by the tier structure (logarithmic in the rows added) instead of @@ -584,9 +547,11 @@ def compact(self, *, max_segments: int | None = None) -> int: 4.1 ms (4096 + 4096 rows) for 16 layers (phase 1 measurement). Returns the number of merges. """ + from .kernels.sdpa_segmented import MAX_FUSED_SEGMENTS + tier = max(1, _int_env("MTPLX_SEGMENTED_KV_TIER", 2)) cap_rows = _int_env("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", 16384) - limit = max_segments if max_segments is not None else _int_env("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", 12) + limit = _int_env("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", MAX_FUSED_SEGMENTS) merged = 0 while len(self._sealed) >= 2: older, newer = self._sealed[-2], self._sealed[-1] @@ -608,27 +573,8 @@ def _merge_last_two(self) -> None: [_rows_view(older.segment.values, older.n), _rows_view(newer.segment.values, newer.n)], axis=2 ) mx.eval(keys, values) - segment = KVSegment(keys, values, sealed=True) - segment.hold(self) - for ref in (older, newer): - ref.segment.release(self) - self._sealed[-2:] = [SegRef(segment, int(older.n) + int(newer.n))] - self.merges += 1 - _count("merge") - self.merge_rows += int(older.n) + int(newer.n) - - def stats(self) -> dict[str, int]: - return { - "segments": self.segment_count, - "sealed": len(self._sealed), - "tail_rows": self._tail_n, - "rows": self.offset, - "bytes": self.nbytes, - "seal_copies": self.seal_copies, - "seal_copy_rows": self.seal_copy_rows, - "merges": self.merges, - "merge_rows": self.merge_rows, - } + self._sealed[-2:] = [SegRef(KVSegment(keys, values, sealed=True), int(older.n) + int(newer.n))] + count_route("merge") class GatheredKVView: @@ -727,11 +673,11 @@ def segment_sdpa_lse(queries, keys, values, n, *, scale, causal): except (ValueError, RuntimeError): # The fused kernel does not take this shape (the detection probe uses a tiny head # dim): the exact reference route serves the call, once counted. - _count("sdpa_lse_kernel_unsupported") + count_route("sdpa_lse_kernel_unsupported") else: - _count("sdpa_lse_kernel") + count_route("sdpa_lse_kernel") return out - _count("sdpa_lse_reference") + count_route("sdpa_lse_reference") return _segment_sdpa_lse_reference(queries, keys, values, n, scale=scale, causal=causal) @@ -798,7 +744,7 @@ def decode_segments_attention( segments_supported, ) - if os.environ.get("MTPLX_NAX_FLASH_ROUTE", "").strip().lower() not in {"1", "true", "yes", "on"}: + if not _env_on("MTPLX_NAX_FLASH_ROUTE"): return None if lane_disabled("nax_flash_dsplit_sdpa"): return None @@ -823,7 +769,7 @@ def decode_segments_attention( if segments_supported(q_len, gqa, d): out = sdpa_nax_flash_dsplit_segments(queries=queries, segments=segments, scale=scale) if out is not None: - _count(f"fused_q{q_len}") + count_route(f"fused_q{q_len}") return out # A window the kernel does not take whole: sub-windows of the rows it does take (both # limits: gqa * rows <= 32 and rows <= 10, so GQA 1 and 2 split at 10 rows, GQA 4 at 8, @@ -845,38 +791,33 @@ def decode_segments_attention( if part is None: return None outs.append(part) - _count(f"chunked_q{q_len}") + count_route(f"chunked_q{q_len}") return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=2) # ---- install -------------------------------------------------------------------------------- -def is_segmented(entry: Any) -> bool: - return isinstance(entry, SegmentedKVCache) - - def _is_stock_plain_kv(entry: Any) -> bool: return type(entry).__name__ == "KVCache" and getattr(entry, "_idx", None) is None -def adapt_layer_for_restore(cache: list[Any], idx: int, state: Any) -> tuple[Any, Any]: - """Prepare layer ``idx`` for a restore of ``state``; returns (entry, state). +def adapt_layer_for_restore(entry: Any, state: Any) -> Any: + """The state to restore into ``entry``. A segment cache takes either snapshot kind (a stock snapshot, from the SSD tier or a stock - bank entry, becomes its rows). A stock layer that cannot hold segments (a layer that was - never tagged: MTP caches, recurrent layers, or one that already had rows at install) gets a - segmented snapshot as contiguous rows. The entry is never replaced: with the switch on, the - layout is decided once, when the cache is built (``install_segmented_attention_kv_cache``). + bank entry, becomes its rows). A stock layer (an MTP cache, a recurrent layer, or a layer + that already had rows when the cache was built) gets a segmented snapshot as contiguous + rows. The entry itself is never replaced: with the switch on, the layout is decided once, + when the cache is built (``install_segmented_attention_kv_cache``). """ - entry = cache[idx] if isinstance(entry, SegmentedKVCache): if not isinstance(state, SegmentedKVState) and isinstance(state, (tuple, list)) and len(state) == 2 and state[0] is not None: - _count("restored_stock_as_segment") - return entry, state + count_route("restored_stock_as_segment") + return state if isinstance(state, SegmentedKVState): - state = state.to_arrays() - return entry, state + return state.to_arrays() + return state def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = None) -> dict[str, int | str]: @@ -904,9 +845,7 @@ def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = and rows == 0 and getattr(entry, "keys", None) is None ): - new = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) - new._mtplx_segmentable = True - cache[idx] = new + cache[idx] = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) stats["entries"] = int(stats["entries"]) + 1 else: stats["skipped"] = int(stats["skipped"]) + 1 @@ -920,7 +859,7 @@ def segmented_ssd_enabled() -> bool: """Segmented snapshots may go to the SSD tier (MTPLX_SEGMENTED_KV_SSD, default on with the switch).""" if not segmented_kv_enabled(): return False - return os.environ.get("MTPLX_SEGMENTED_KV_SSD", "").strip().lower() not in {"0", "false", "no", "off"} + return _env_on("MTPLX_SEGMENTED_KV_SSD", default=True) class SegmentedRows: @@ -1025,7 +964,7 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: support = {"supported": False, "reasons": ["model_not_checked"]} return {"enabled": False, "requested": True, "model_support": support} per_entry: list[dict[str, Any]] = [] - unique: dict[int, int] = {} + all_states: list[SegmentedKVState] = [] for entry in entries: counts: list[int] = [] sealed_bytes = 0 @@ -1033,8 +972,7 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: for state in walk_segmented_states(snapshot.states) if snapshot is not None else (): counts.append(len(state.refs)) sealed_bytes += int(state.nbytes) - for ref in state.refs: - unique.setdefault(ref.segment.id, int(ref.segment.nbytes)) + all_states.append(state) live = getattr(entry, "cache_ref", None) for layer in live or (): if isinstance(layer, SegmentedKVCache): @@ -1058,7 +996,7 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: "prefill_route": prefill_route(), "entries": per_entry, "max_segments": max((e["segments"] for e in per_entry), default=0), - "sealed_bytes_unique": sum(unique.values()), + "sealed_bytes_unique": unique_state_bytes(all_states), "merges": counts_now.get("merge", 0), "seals": counts_now.get("seal", 0), "seal_copies": counts_now.get("seal_copy", 0), @@ -1070,16 +1008,6 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: # ---- accounting ----------------------------------------------------------------------------- -def unique_state_bytes(states: Iterable[Any]) -> int: - """Bytes of the segments referenced by ``states`` (SegmentedKVState objects), each counted once.""" - seen: dict[int, int] = {} - for state in states: - if isinstance(state, SegmentedKVState): - for ref in state.refs: - seen.setdefault(ref.segment.id, ref.segment.nbytes) - return sum(seen.values()) - - def walk_segmented_states(value: Any) -> Iterable[SegmentedKVState]: """Every SegmentedKVState inside a snapshot tree (tuples, lists, dicts).""" if isinstance(value, SegmentedKVState): diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 3af799f70..df0693797 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -21780,7 +21780,7 @@ def _admission_restore_copies_prefix( return True from mtplx.session_bank import entry_has_segmented_kv - if entry_has_segmented_kv(entry): + if segmented_kv_enabled() and entry_has_segmented_kv(entry): # A segmented snapshot is a list of references to sealed segments: the restored # cache writes its turn into a new tail, so the history is never copied # (mtplx/segmented_kv.py; +4.95 GiB at 80K with a stock snapshot, +0.125 GiB with diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py index 7d99e5c36..0f6811a53 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -38,8 +38,10 @@ from .segmented_kv import ( SegmentedKVCache, SegmentedKVState, + segmented_kv_enabled, segmented_kv_health, segmented_ssd_enabled, + unique_state_bytes, walk_segmented_states, ) @@ -426,14 +428,9 @@ def _cache_has_segmented_kv(cache: list[Any] | None) -> bool: def _shared_segment_overcount(entries: list[Any]) -> int: """Bytes counted more than once because several snapshots reference the same segment.""" - naive = 0 - unique: dict[int, int] = {} - for entry in entries: - for state in _entry_segmented_states(entry): - naive += int(state.nbytes) - for ref in state.refs: - unique.setdefault(ref.segment.id, int(ref.segment.nbytes)) - return max(0, naive - sum(unique.values())) + states = [state for entry in entries for state in _entry_segmented_states(entry)] + naive = sum(int(state.nbytes) for state in states) + return max(0, naive - unique_state_bytes(states)) def _snapshot_nbytes(snapshot: CacheSnapshot) -> int: @@ -1306,8 +1303,11 @@ def total_nbytes(self) -> int: # dynamic ceiling (working set = active - weights - this total) took # the lease for working set and evicted the durable snapshots instead. entries = list(self._entries.values()) - # A segment that several entries reference is held once (mtplx/segmented_kv.py). - return sum(entry.held_nbytes for entry in entries) - _shared_segment_overcount(entries) + total = sum(entry.held_nbytes for entry in entries) + if segmented_kv_enabled(): + # A segment that several entries reference is held once (mtplx/segmented_kv.py). + total -= _shared_segment_overcount(entries) + return total @property def lease_entries(self) -> int: diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py index 1387eae8d..35c9025f2 100644 --- a/tests/test_segmented_kv.py +++ b/tests/test_segmented_kv.py @@ -7,8 +7,6 @@ from __future__ import annotations -import gc - import mlx.core as mx import pytest from mlx_lm.models.cache import KVCache @@ -28,6 +26,7 @@ attend_segments_lse, install_segmented_attention_kv_cache, merge_segment_outputs, + route_counts, segment_sdpa_lse, segmented_kv_enabled, unique_state_bytes, @@ -144,7 +143,7 @@ def test_verify_rollback_stays_in_the_tail_and_never_touches_sealed_segments() - cache.update_and_fetch(*_rows(3, 4)) # a verify window assert cache.offset == 414 assert cache.trim(3) == 3 # accept one token, roll back three - assert cache.offset == 411 and cache.tail_rows == 11 + assert cache.offset == 411 and cache._tail_n == 11 cache.update_and_fetch(*_rows(4, 4)) assert cache.offset == 415 assert cache._sealed[0].segment is sealed @@ -177,17 +176,18 @@ def test_seal_copies_only_the_tail_and_gives_it_its_exact_capacity() -> None: cache = SegmentedKVCache() _fill(cache, [(1, 300)]) assert cache._tail.capacity > 300 # growth slack + copies = route_counts.get("seal_copy", 0) assert cache.seal() is True assert cache.seal() is False # nothing left to seal ref = cache._sealed[0] assert ref.segment.sealed and ref.segment.capacity == ref.n == 300 - assert cache.seal_copies == 1 and cache.seal_copy_rows == 300 + assert route_counts.get("seal_copy", 0) == copies + 1 # the history is not copied by a later seal _fill(cache, [(2, 40)]) history = cache._sealed[0].segment.keys cache.seal() assert cache._sealed[0].segment.keys is history - assert cache.seal_copy_rows == 340 + assert cache._sealed[1].segment.capacity == 40 def test_state_is_a_list_of_references_and_a_snapshot_never_aliases_the_tail() -> None: @@ -246,7 +246,6 @@ def test_fork_inside_a_segment_is_a_reference_without_a_copy() -> None: ref = KVCache() ref.update_and_fetch(*_rows(1, 400)) assert _same(owner, ref) - assert segment.refcount >= 3 # owner, state, fork def test_state_setter_accepts_a_contiguous_pair() -> None: @@ -260,28 +259,7 @@ def test_state_setter_accepts_a_contiguous_pair() -> None: assert cache.offset == 0 -def test_refcount_follows_the_holders_and_eviction_frees_the_segment() -> None: - cache = SegmentedKVCache() - _fill(cache, [(1, 300)]) - state = cache.state - segment = state.refs[0].segment - assert segment.refcount == 2 # the cache and the snapshot - other = SegmentedKVState(state.refs) - assert segment.refcount == 3 - del other - gc.collect() - assert segment.refcount == 2 - cache.state = None - assert segment.refcount == 1 - import weakref - - probe = weakref.ref(segment) - del segment, state - gc.collect() - assert probe() is None # nobody holds it: the buffers are free - - -def test_tiered_merge_only_merges_small_recent_unshared_segments(monkeypatch) -> None: +def test_tiered_merge_merges_small_recent_segments_and_leaves_the_history_alone(monkeypatch) -> None: monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "16384") cache = SegmentedKVCache() _fill(cache, [(1, 4096)]) @@ -312,9 +290,10 @@ def test_merge_also_merges_segments_a_snapshot_holds_and_leaves_the_snapshot_int held_segments = [ref.segment for ref in held.refs] before = [ref.segment for ref in cache._sealed] assert before == held_segments + merges = route_counts.get("merge", 0) _fill(cache, [(3, 300)]) # new tail: compact() merges the two equal-sized held segments assert len(cache._sealed) == 1 and cache._sealed[0].n == 600 - assert cache.merges == 1 and cache.merge_rows == 600 + assert route_counts.get("merge", 0) == merges + 1 assert held.rows == 600 and [ref.segment for ref in held.refs] == held_segments # the snapshot is untouched control = KVCache() for seed in (1, 2): @@ -332,6 +311,7 @@ def test_segment_count_stays_bounded_over_many_turns_with_snapshots(monkeypatch) cache = SegmentedKVCache() _fill(cache, [(1, 20000)]) cache.seal() + merges = route_counts.get("merge", 0) snapshots = [cache.state] peak = 0 for turn in range(2, 26): @@ -342,7 +322,7 @@ def test_segment_count_stays_bounded_over_many_turns_with_snapshots(monkeypatch) peak = max(peak, cache.segment_count) assert peak <= 12, peak assert cache.segment_count <= 12 - assert cache.merges > 0 + assert route_counts.get("merge", 0) > merges def test_nbytes_counts_each_shared_segment_once() -> None: diff --git a/tests/test_segmented_kv_bank.py b/tests/test_segmented_kv_bank.py index b8f15a1ad..ab3788976 100644 --- a/tests/test_segmented_kv_bank.py +++ b/tests/test_segmented_kv_bank.py @@ -226,10 +226,8 @@ def test_eviction_releases_the_segments(monkeypatch) -> None: gc.collect() mx.clear_cache() held = mx.get_active_memory() - segments = [state.refs[0].segment for state in entry.cache_snapshot.states] - assert all(segment.refcount >= 1 for segment in segments) assert bank.clear(session_id="s") == 1 - del entry, segments + del entry gc.collect() mx.clear_cache() assert held - mx.get_active_memory() >= BASE * ROW_BYTES * 0.9 @@ -263,6 +261,7 @@ def put_entry(self, *a, **k): def test_admission_prices_no_history_copy_for_a_segmented_entry(monkeypatch) -> None: from mtplx.server.openai import _admission_restore_copies_prefix + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") # segmented entries only exist with the switch on bank = _bank(monkeypatch) _, seg_entry = _turn(bank, Runtime("segmented"), list(range(BASE)), BASE, 1, "x") stock_bank = _bank(monkeypatch) @@ -277,7 +276,7 @@ def test_health_block_is_absent_with_the_switch_off_and_reports_segments_with_it bank = _bank(monkeypatch) runtime = Runtime("segmented") tokens = list(range(BASE)) - _, entry = _turn(bank, runtime, tokens, BASE, 1, "s") + _turn(bank, runtime, tokens, BASE, 1, "s") tokens2 = tokens + list(range(BASE * 10, BASE * 10 + TURN)) _turn(bank, runtime, tokens2, TURN, 2, "s", prev=tokens) assert "segmented_kv" not in bank.to_dict() diff --git a/tests/test_segmented_kv_one_path.py b/tests/test_segmented_kv_one_path.py index ff4706d55..4c3587e09 100644 --- a/tests/test_segmented_kv_one_path.py +++ b/tests/test_segmented_kv_one_path.py @@ -59,13 +59,6 @@ def test_segments_from_the_first_token_whatever_the_compiled_verify_settings(mon assert all(c.offset == 0 for c in cache) -def test_there_is_no_boundary_setting() -> None: - import mtplx.segmented_kv as module - - assert not hasattr(module, "segmented_kv_min_tokens") - assert "min_tokens" not in module.segmented_kv_health([]) - - def test_a_layer_that_already_holds_rows_stays_stock() -> None: cache = [KVCache()] cache[0].update_and_fetch(*_rows(1, 60)) @@ -119,7 +112,7 @@ def test_a_stock_buffer_restored_as_a_segment_is_stored_as_exact_rows() -> None: assert mx.array_equal(flat, segment).item() -def test_a_segmented_snapshot_restores_as_rows_into_an_untagged_stock_cache() -> None: +def test_a_segmented_snapshot_restores_as_rows_into_a_stock_cache() -> None: src = _target_cache(1) src[0].update_and_fetch(*_rows(1, 150)) snap = snapshot_cache(src) @@ -129,7 +122,7 @@ def test_a_segmented_snapshot_restores_as_rows_into_an_untagged_stock_cache() -> assert _same(dest[0].state, src[0].gather()) -def test_untagged_caches_are_never_converted() -> None: +def test_stock_caches_with_rows_are_never_converted() -> None: mtp = [KVCache()] mtp[0].update_and_fetch(*_rows(1, 150)) snap = snapshot_cache(mtp) diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index 804d47d6f..7ce3624e7 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -85,7 +85,6 @@ def test_an_unsupported_model_keeps_stock_caches_and_never_builds_a_segmented_on cache = [KVCache(), KVCache()] configure_tail_owned_attention_kv_cache(cache) assert all(type(c) is KVCache for c in cache) - assert not any(getattr(c, "_mtplx_segmentable", False) for c in cache) for c in cache: c.update_and_fetch(mx.zeros((1, 2, 9000, 8)), mx.zeros((1, 2, 9000, 8))) configure_tail_owned_attention_kv_cache(cache) # what the repage after a prefill runs @@ -196,9 +195,11 @@ def _qwen3_attention(head_dim=32, q=4, kv=2, hidden=64): def test_plain_qwen3_attention_is_hooked_for_segmented_caches_only_and_matches_the_stock_forward() -> None: """A plain q/k-norm attention (no gate) goes through the hook body only on a segmented cache; with a stock cache it is the unchanged mlx-lm forward, so models without the switch keep their code.""" - from mlx_lm.models.cache import KVCache - - from mtplx.attention_split import _attention_has_gated_q_proj, _attention_has_plain_q_proj, _install_split_attention_hook + from mtplx.attention_split import ( + _attention_has_gated_q_proj, + _attention_has_plain_q_proj, + _install_split_attention_hook, + ) from mtplx.segmented_kv import SegmentedKVCache attn = _qwen3_attention() From 36f79fa3e83326732586e8b274f8576c8cfabd05 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 16:22:18 +0200 Subject: [PATCH 26/37] Changelog: segmented KV entry as a heading with sub-bullets, current behaviour only (no 32K fence history, windows 1 to 32 rows) --- CHANGELOG.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c19cb0475..1f30f8411 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,12 @@ All notable user-facing changes to MTPLX. The format is based on ### 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 (a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K on Qwen3.8 27B at the first write of the next turn; phase 1 measured +0.125 GiB with segments). The session bank stores the list of segments as the snapshot, shares segments between entries by reference count, and the admission gate prices no history copy for them. Verify windows of 2 to 5 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`; 0 to 1.1 % on a verify forward at 6 or 20 segments, phase 1 measurement); a call over several segments that no kernel takes gathers the rows and runs the stock route (counter `gather_q*`; none fires in a normal turn). Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers otherwise. With the switch on the full-attention layers are segment caches from the first token: one path, no context fence (an earlier version kept the stock cache below `MTPLX_COMPILED_VERIFY_MAX_CONTEXT` so the compiled verify could run; measured on Qwen3.8 27B at 8K to 32K, segments with the eager verify are 3 to 12 % faster per verify and lower the peak memory by 1 to 4.5 GiB at every context, so the fence and its `MTPLX_SEGMENTED_KV_MIN_TOKENS` setting are gone). A cache with one segment (a cold prompt) runs the stock route on that segment's rows, a zero-copy view; calls over several segments take the kernels at any length. Segmented entries persist to the SSD tier in the existing format (`MTPLX_SEGMENTED_KV_SSD=0` skips them as before): the codec encodes a segmented snapshot as the 256-row block blobs of the stock tensors with the same rows, slicing across segments without gathering the history, so shared prefixes dedupe through the content-addressed store, entries restore with the switch off and stock entries restore with it on. The streaming spill hashes each sealed segment's blocks once (the digest is kept on the segment, and the blob is confirmed on disk before it is referenced), so a follow-up turn no longer rereads and rehashes the whole history. The tiered merge of small recent segments also merges segments that bank snapshots hold (before, nothing ever merged because every turn's snapshot holds its segments), which keeps a conversation at a handful of segments (3 to 5 over 12 follow-up turns at 50K) instead of two more per turn, below the 12-segment fused launch. An SSD restore builds the request's dense cache layout instead of the paged default (with the switch on). `/health` gets a `segmented_kv` block (segments and sealed bytes per entry, merges, route and SSD counters), absent when the switch is off. Not covered: the compiled verify (segmented caches verify eager; keep `MTPLX_VERIFY_ASYNC_CHUNK_LAYERS=8`, it still helps by 4 to 19 % with segments), paged and quantized KV. +- **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, the NAX flash route on). 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. + - 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. ### Fixed From 6fdd948451dd402796bb28adf00e3c7e71e690a1 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 18:39:48 +0200 Subject: [PATCH 27/37] Segmented KV: Llama-style attention (Llama, Mistral) through the split-attention hook mlx-lm's llama Attention has no q/k norm, so the hook did not take it and the support check reported attention_not_hooked (Mistral 7B v0.3 kept the stock cache). The hook body now serves an ungated attention with or without q/k norm, on a segmented cache only; with a stock cache the forward is mlx-lm's own, switch on or off. Tests: detection of the three forms, flag-off and stock-cache identity (with and without bias), support verdict, and a GPU test at Mistral shapes (32/8 heads, head_dim 128) through turns and rollback with no gather. --- mtplx/attention_split.py | 40 +++++++++++++---- tests/test_sdpa_segmented_hd128.py | 54 +++++++++++++++++++++++ tests/test_segmented_kv_support.py | 69 ++++++++++++++++++++++++++++++ 3 files changed, 155 insertions(+), 8 deletions(-) diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index 941106452..720bdcbc6 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -358,6 +358,27 @@ def _attention_has_plain_q_proj(attn: Any) -> bool: 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): @@ -395,11 +416,12 @@ def split_call( # and float32 values (mtplx/attention_math.py). gated_q_proj = _attention_has_gated_q_proj(self) if not gated_q_proj and not ( - isinstance(cache, SegmentedKVCache) and _attention_has_plain_q_proj(self) + 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, a plain q/k-norm attention (Qwen3, head_dim 128) on a - # segmented cache: the stock forward would gather the history on every call. + # 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 @@ -420,10 +442,12 @@ def split_call( 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, n_kv_heads, -1) - ).transpose(0, 2, 1, 3) + 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, @@ -981,7 +1005,7 @@ def configure_split_full_attention( # 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_plain_q_proj(attn)) + 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( diff --git a/tests/test_sdpa_segmented_hd128.py b/tests/test_sdpa_segmented_hd128.py index ce9e32b41..2367397e6 100644 --- a/tests/test_sdpa_segmented_hd128.py +++ b/tests/test_sdpa_segmented_hd128.py @@ -180,6 +180,60 @@ def run(cache): assert float(mx.abs(a.astype(mx.float32) - b.astype(mx.float32)).max().item()) / ulp <= 4 +@needs_nax +def test_llama_style_attention_at_head_dim_128_matches_the_stock_cache_through_turns_and_rollback(monkeypatch) -> None: + """Mistral-7B shapes (32 query / 8 KV heads, head_dim 128, no q/k norm) through the attention_split ladder.""" + from types import SimpleNamespace + + from mlx.utils import tree_map + from mlx_lm.models.cache import KVCache + from mlx_lm.models.llama import Attention, ModelArgs + + from mtplx.attention_split import configure_split_full_attention + from mtplx.segmented_kv import SegmentedKVCache + + for name in ("MTPLX_GQA_PACKED_SDPA", "MTPLX_NAX_FLASH_ROUTE"): + monkeypatch.setenv(name, "1") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + hidden = 256 + args = ModelArgs( + model_type="mistral", hidden_size=hidden, num_hidden_layers=1, intermediate_size=64, num_attention_heads=HQ, + rms_norm_eps=1e-5, vocab_size=64, head_dim=D, num_key_value_heads=HK, rope_theta=1e6, tie_word_embeddings=False, + ) + mx.random.seed(0) + attn = Attention(args) + attn.update(tree_map(lambda p: p.astype(mx.bfloat16), attn.parameters())) + mx.eval(attn.parameters()) + configure_split_full_attention(SimpleNamespace(model=SimpleNamespace(layers=[SimpleNamespace(self_attn=attn)]))) + + def x(seed, n): + return (mx.random.normal((1, n, hidden), key=mx.random.key(seed)) * 0.5).astype(mx.bfloat16) + + def run(cache): + outs = [attn(x(1, 9000), mask="causal", cache=cache)] + if isinstance(cache, SegmentedKVCache): + state = cache.state + cache = SegmentedKVCache() + cache.state = state + outs.append(attn(x(2, 300), mask="causal", cache=cache)) + for seed, window in ((4, 4), (5, 4), (6, 9), (7, 1)): + outs.append(attn(x(seed, window), mask="causal" if window > 1 else None, cache=cache)) + cache.trim(max(1, window - 2)) + mx.eval(outs) + return outs, cache + + stock, _ = run(KVCache()) + module.route_counts.clear() + seg, cache = run(SegmentedKVCache()) + assert cache.segment_count >= 2 + assert not any(k.startswith("gather_q") and k[8:].isdigit() for k in module.route_counts), module.route_counts + assert any(k.startswith("fused_q") for k in module.route_counts), module.route_counts + for a, b in zip(stock, seg): + top = float(mx.maximum(mx.abs(a).max(), mx.abs(b).max()).item()) + ulp = 2.0 ** (math.floor(math.log2(top)) - 7) + assert float(mx.abs(a.astype(mx.float32) - b.astype(mx.float32)).max().item()) / ulp <= 4 + + def _low_gqa_cache(hk: int, d: int, window: int, seed: int = 11): """Two sealed segments plus a tail whose last ``window`` rows are the verify window.""" from mtplx.segmented_kv import SegmentedKVCache diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index 7ce3624e7..fccafff1d 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -260,3 +260,72 @@ def test_with_the_lse_kernel_there_is_no_note(monkeypatch, capsys) -> None: monkeypatch.setattr(module, "_SDPA_LSE", True) evaluate_model_support([_attn()]) assert "no logsumexp output" not in capsys.readouterr().out + + +def _llama_attention(head_dim=32, q=4, kv=2, hidden=64, bias=False): + from mlx_lm.models.llama import Attention, ModelArgs + + args = ModelArgs( + model_type="mistral", hidden_size=hidden, num_hidden_layers=1, intermediate_size=64, num_attention_heads=q, + rms_norm_eps=1e-5, vocab_size=64, head_dim=head_dim, num_key_value_heads=kv, rope_theta=1e6, + attention_bias=bias, tie_word_embeddings=False, + ) + mx.random.seed(0) + return Attention(args) + + +def test_llama_style_attention_is_recognised_and_other_forms_are_not() -> None: + from mtplx.attention_split import ( + _attention_has_gated_q_proj, + _attention_has_plain_q_proj, + _attention_has_unnormed_q_proj, + ) + + llama = _llama_attention() + assert _attention_has_unnormed_q_proj(llama) + assert not _attention_has_plain_q_proj(llama) and not _attention_has_gated_q_proj(llama) + assert not _attention_has_unnormed_q_proj(_qwen3_attention()) # has q/k norm: the plain form + assert not _attention_has_unnormed_q_proj(SimpleNamespace(q_proj=None, n_heads=4, head_dim=32)) + + +@pytest.mark.parametrize("bias", [False, True]) +def test_llama_style_attention_is_hooked_for_segmented_caches_only_and_matches_the_stock_forward(bias) -> None: + """Llama/Mistral attention (no q/k norm) goes through the hook body only on a segmented cache; with + a stock cache, switch on or off, it is the unchanged mlx-lm forward bit for bit.""" + from mtplx.attention_split import _install_split_attention_hook + from mtplx.segmented_kv import SegmentedKVCache + + attn = _llama_attention(bias=bias) + + def run(cache): + outs = [] + for seed, n in ((1, 40), (2, 9), (3, 1)): + x = mx.random.normal((1, n, 64), key=mx.random.key(seed)) + outs.append(attn(x, mask="causal" if n > 1 else None, cache=cache)) + mx.eval(outs) + return outs + + before = run(KVCache()) + _install_split_attention_hook(attn) + for enabled in (False, True): + attn._mtplx_split_full_attention_enabled = enabled + for a, b in zip(before, run(KVCache())): + assert mx.array_equal(a, b).item() + attn._mtplx_split_full_attention_enabled = True + cache = SegmentedKVCache() + seg = run(cache) + cache.seal() + for a, b in zip(before, seg): + assert float(mx.abs(a - b).max().item()) < 1e-4 # same math on the gathered rows (CPU: no kernel) + assert cache.offset == 50 + + +def test_the_hook_covers_llama_style_attention_in_the_support_check(monkeypatch) -> None: + from mtplx.attention_split import configure_split_full_attention + + attn = _llama_attention(head_dim=128, q=8, kv=2, hidden=128) + model = SimpleNamespace(model=SimpleNamespace(layers=[SimpleNamespace(self_attn=attn)])) + monkeypatch.setenv("MTPLX_GQA_PACKED_SDPA", "1") + stats = configure_split_full_attention(model) + assert stats["segmented_kv_supported"] is True + assert model_support()["head_dims"] == [128] and model_support()["gqa"] == [4] From 8e15baadee90e225abfde573ca73d6b0728ded45 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 20:28:08 +0200 Subject: [PATCH 28/37] Changelog: segmented KV serves Llama-style attention (Mistral 7B v0.3 numbers) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f30f8411..eaf2482f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ All notable user-facing changes to MTPLX. The format is based on - **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, the NAX flash route on). Any other model, and a model loaded without that check, keeps the stock cache; `/health` says why. + - 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, the NAX flash route on). 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. - 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. From 715517731cc0214645daefd0a57dbbbdb6ce7d8e Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Fri, 9 Oct 2026 23:26:48 +0200 Subject: [PATCH 29/37] Segmented KV under any profile: the attention switches are kill switches only, the self-check probes the segment kernel Segments only engaged with MTPLX_GQA_PACKED_SDPA and MTPLX_NAX_FLASH_ROUTE on, which only the turbo profile sets; under sustained (the recommended profile of Qwen3.6-35B-A3B and Apodex-1.1-mini) the support check switched them off. Two holes besides: the verdict was made before the kernel self-check, so a self-check that disabled the kernel left segmented caches that gathered the history on every multi-segment call; and under a profile without an attention route the hook sent a segmented cache to the stock forward (a gather per call). - Verdict: neither switch is required; either one set to 0 still switches segments off (gqa_packed_sdpa_off, nax_flash_route_off). nax_unavailable stays (M1 to M4: off). - Hook: a segmented cache always takes the segment route, whatever the profile. - Self-check: runs whenever MTPLX_SEGMENTED_KV is on and probes the segment kernel (lane segmented_kv_sdpa, two segments against stock SDPA) at the model's head_dim and GQA factor. A failure, or a self-check that cannot run, turns the verdict off before the first cache is built: stock caches, reason segment_kernel_selfcheck_failed on /health, one log line. With the switch off nothing changes: every path sits behind segmented_kv_requested() or a segmented cache. Tests: verdict with the switches unset, on and 0; self-check probe shapes, failure (mismatch, raise, bail, no run) to stock caches; hook routing without the switches for gated, q/k-norm and Llama-style attention; flag off leaves the self-check and verdict alone; GPU: the segment kernel serves without the switches, no gather, and the probe passes on the real kernel at head_dim 128 and 256. --- mtplx/attention_split.py | 15 ++- mtplx/kernel_selfcheck.py | 59 +++++++++ mtplx/segmented_kv.py | 81 +++++++++--- tests/test_segmented_kv_attention.py | 44 ++++++- tests/test_segmented_kv_support.py | 182 ++++++++++++++++++++++++--- 5 files changed, 336 insertions(+), 45 deletions(-) diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index 720bdcbc6..33b9a10da 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -393,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 @@ -414,6 +418,9 @@ 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). + # 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) @@ -506,11 +513,7 @@ def split_call( ) segmented_route = "segments_lse" segmented_count("prefill_lse") - elif ( - getattr(self, "_mtplx_gqa_packed_sdpa_enabled", False) - and not blockwise_enabled - and not vllm_metal_paged_enabled - ): + else: segmented_output = decode_segments_attention( queries, cache, diff --git a/mtplx/kernel_selfcheck.py b/mtplx/kernel_selfcheck.py index c42aa50e1..9e7da4000 100644 --- a/mtplx/kernel_selfcheck.py +++ b/mtplx/kernel_selfcheck.py @@ -95,6 +95,8 @@ def selfcheck_enabled(*, prism_ternary: bool = False, flash_next: bool = False) ``flash_next``: the model is a Flash-Next (qwen4_exp) load, likewise for its four prefill kernels. """ + from .segmented_kv import segmented_kv_requested + raw = str(os.environ.get("MTPLX_KERNEL_SELFCHECK", "")).strip().lower() if raw in {"0", "false", "off", "no"}: return False @@ -102,6 +104,7 @@ def selfcheck_enabled(*, prism_ternary: bool = False, flash_next: bool = False) return True return ( prism_ternary + or segmented_kv_requested() or flash_next or _env_on("MTPLX_NAX_VERIFY") or _env_on("MTPLX_GQA_PACKED_SDPA") @@ -277,6 +280,38 @@ def _check_gqa_packed(mx, dtype, kernel=None, *, d=128, q_len=4) -> float: return _max_abs_diff(mx, out, ref) +def _check_segments(mx, dtype, *, d: int, gqa: int) -> float: + """The segment kernel over two segments against stock SDPA on the concatenated rows.""" + from .kernels.sdpa_segmented import ( + MAX_M_ROWS, + MAX_Q_LEN, + sdpa_nax_flash_dsplit_segments, + ) + + hk = 2 + hq = hk * gqa + q_len = max(1, min(4, MAX_Q_LEN, MAX_M_ROWS // gqa)) + rows, capacity = (192, 136), 256 + scale = d**-0.5 + mx.random.seed(13) + + def buffer(): + return (mx.random.normal((1, hk, capacity, d), dtype=mx.float32) * 0.5).astype(dtype) + + queries = (mx.random.normal((1, hq, q_len, d), dtype=mx.float32) * 0.5).astype(dtype) + segments = [(buffer(), buffer(), n) for n in rows] + out = sdpa_nax_flash_dsplit_segments(queries=queries, segments=segments, scale=scale) + if out is None: + return float("inf") + keys = mx.concatenate([k[:, :, :n, :] for k, _, n in segments], axis=2) + values = mx.concatenate([v[:, :, :n, :] for _, v, n in segments], axis=2) + total = sum(rows) + # Tail-causal: query row j attends to rows n <= total - q_len + j. + mask = mx.arange(total).reshape(1, total) <= (total - q_len + mx.arange(q_len).reshape(q_len, 1)) + ref = mx.fast.scaled_dot_product_attention(queries, keys, values, scale=scale, mask=mask) + return _max_abs_diff(mx, out, ref) + + def _check_fused_add_rmsnorm(mx, dtype) -> float: from .gdn_capture import _fused_post_norm_tg_override from .kernels.fused_norm import fused_add_rmsnorm @@ -934,6 +969,21 @@ def _record(lane: str, tolerance: float, probe) -> None: else: lanes[lane] = _STATUS_SKIPPED + # MTPLX_SEGMENTED_KV: the segment kernel at the loaded model's head_dim and GQA factor + # (any profile). A failure turns segments off for this load (stock caches), since a + # segmented cache without its kernel would gather the history on every call. + from .segmented_kv import kernel_probe_shapes + + segment_shapes = kernel_probe_shapes() + if segment_shapes: + _record( + "segmented_kv_sdpa", + _SDPA_TOLERANCE, + lambda: max(_check_segments(mx, dtype, d=d, gqa=g) for d, g in segment_shapes), + ) + else: + lanes["segmented_kv_sdpa"] = _STATUS_SKIPPED + if _env_on("MTPLX_FUSE_POST_NORM_RESIDUAL"): # Bitwise gate: this lane's contract is exact identity with the # unfused reference (#319). _NORM_TOLERANCE stays loose only for @@ -1049,6 +1099,10 @@ def _record(lane: str, tolerance: float, probe) -> None: _DISABLED_LANES.update( lane for lane, status in lanes.items() if status == _STATUS_FALLBACK ) + if lanes.get("segmented_kv_sdpa") == _STATUS_FALLBACK: + from .segmented_kv import note_kernel_selfcheck_failed + + note_kernel_selfcheck_failed() dtype_tag = {mx.bfloat16: "bfloat16", mx.float16: "float16"}.get(dtype, str(dtype)) report = { "lanes": dict(lanes), @@ -1161,4 +1215,9 @@ def maybe_run_model_selfcheck(model: Any) -> dict[str, Any] | None: return report except Exception as exc: # noqa: BLE001 - probe must never block serving logger.warning("[mtplx] kernel selfcheck failed to run: %s", exc) + from .segmented_kv import kernel_probe_shapes, note_kernel_selfcheck_failed + + if kernel_probe_shapes(): + # The segment kernel went unchecked: stock caches, never an unprobed kernel. + note_kernel_selfcheck_failed() return None diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 1e17af0f1..6d0649099 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -68,6 +68,20 @@ def segmented_kv_requested() -> bool: return _env_on("MTPLX_SEGMENTED_KV") +def segment_route_switched_off() -> list[str]: + """The attention switches an operator set to 0 explicitly (kill switches for the segment route). + + Segments do not need MTPLX_GQA_PACKED_SDPA or MTPLX_NAX_FLASH_ROUTE on (they run their own + kernel, under any profile); only an explicit 0 on either turns them off. + """ + reasons = [] + if not _env_on("MTPLX_GQA_PACKED_SDPA", default=True): + reasons.append("gqa_packed_sdpa_off") + if not _env_on("MTPLX_NAX_FLASH_ROUTE", default=True): + reasons.append("nax_flash_route_off") + return reasons + + #: Verdict on the loaded model (``evaluate_model_support``); None until a model was checked. _MODEL_SUPPORT: dict[str, Any] | None = None _UNSUPPORTED_LOGGED = False @@ -111,7 +125,7 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No from .kernels.sdpa_segmented import SUPPORTED_HEAD_DIMS, segments_supported from .nax_verify import nax_available - global _MODEL_SUPPORT, _UNSUPPORTED_LOGGED, _LSE_NOTE_LOGGED + global _MODEL_SUPPORT, _LSE_NOTE_LOGGED layers = list(layers) flags = list(hooked) if hooked is not None else [True] * len(layers) reasons: list[str] = [] @@ -135,10 +149,7 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No reasons.append(f"head_dim_{head_dim}_unsupported") elif q_heads % kv_heads or not segments_supported(1, q_heads // kv_heads, head_dim): reasons.append(f"gqa_{q_heads}_{kv_heads}_unsupported") - if not all(getattr(a, "_mtplx_gqa_packed_sdpa_enabled", False) for a in layers): - reasons.append("gqa_packed_sdpa_off") - if not _env_on("MTPLX_NAX_FLASH_ROUTE"): - reasons.append("nax_flash_route_off") + reasons.extend(segment_route_switched_off()) if not nax_available(): reasons.append("nax_unavailable") reasons = list(dict.fromkeys(reasons)) @@ -151,14 +162,7 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No "full_attention_layers": len(layers), } _MODEL_SUPPORT = verdict - if reasons and segmented_kv_requested() and not _UNSUPPORTED_LOGGED: - _UNSUPPORTED_LOGGED = True - print( - "[mtplx] MTPLX_SEGMENTED_KV is on but this model cannot use it (" - + ", ".join(reasons) - + "): keeping the stock KV cache", - flush=True, - ) + _log_unsupported(reasons) if not reasons and segmented_kv_requested() and prefill_route() == "gather" and not _LSE_NOTE_LOGGED: _LSE_NOTE_LOGGED = True print( @@ -170,6 +174,47 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No return verdict +def _log_unsupported(reasons: Sequence[str]) -> None: + global _UNSUPPORTED_LOGGED + if reasons and segmented_kv_requested() and not _UNSUPPORTED_LOGGED: + _UNSUPPORTED_LOGGED = True + print( + "[mtplx] MTPLX_SEGMENTED_KV is on but this model cannot use it (" + + ", ".join(reasons) + + "): keeping the stock KV cache", + flush=True, + ) + + +def kernel_probe_shapes() -> list[tuple[int, int]]: + """(head_dim, gqa) pairs the kernel self-check probes the segment kernel at. + + Empty unless segments are requested and the loaded model passed the support check: the + self-check then tests the kernel at the model's own shapes before any cache is built. + """ + from .kernels.sdpa_segmented import segments_supported + + if not segmented_kv_enabled(): + return [] + support = _MODEL_SUPPORT or {} + pairs = [(int(d), int(g)) for d in support.get("head_dims", []) for g in support.get("gqa", [])] + return [(d, g) for d, g in pairs if segments_supported(1, g, d)] + + +def note_kernel_selfcheck_failed() -> None: + """The self-check found the segment kernel wrong or failing on this GPU: stock caches from now on. + + Without the kernel every call over several segments would gather the history, so the verdict + turns off (reason on /health, one log line) before the first cache is built. + """ + global _MODEL_SUPPORT + if _MODEL_SUPPORT is None: + return + reasons = list(dict.fromkeys([*_MODEL_SUPPORT.get("reasons", []), "segment_kernel_selfcheck_failed"])) + _MODEL_SUPPORT = {**_MODEL_SUPPORT, "supported": False, "reasons": reasons} + _log_unsupported(reasons) + + def _int_env(name: str, default: int) -> int: raw = os.environ.get(name, "").strip() try: @@ -725,9 +770,9 @@ def decode_segments_attention( ) -> mx.array | None: """Verify/decode attention straight over the segments, or None for the stock ladder. - Engages under the same conditions as the packed route over a dense cache, with the total - length where that route reads the buffer capacity: MTPLX_NAX_FLASH_ROUTE on, a verify - window of 1 to ``VERIFY_WINDOW_MAX`` (32) rows, a causal or no mask, and a total length of + Engages under any profile unless MTPLX_GQA_PACKED_SDPA or MTPLX_NAX_FLASH_ROUTE is set to 0 + or the kernel self-check failed the segment kernel: a verify window of 1 to + ``VERIFY_WINDOW_MAX`` (32) rows, a causal or no mask, and a total length of at least ``packed_threshold`` when the cache is one segment (several segments take the kernel at any length). The head-dim-split segment kernel takes a window whole when ``gqa * q_len <= 32`` and ``q_len <= MAX_Q_LEN`` (10); a longer window runs as sub-windows @@ -744,9 +789,7 @@ def decode_segments_attention( segments_supported, ) - if not _env_on("MTPLX_NAX_FLASH_ROUTE"): - return None - if lane_disabled("nax_flash_dsplit_sdpa"): + if segment_route_switched_off() or lane_disabled("segmented_kv_sdpa"): return None q_len = int(queries.shape[2]) if not 1 <= q_len <= VERIFY_WINDOW_MAX: diff --git a/tests/test_segmented_kv_attention.py b/tests/test_segmented_kv_attention.py index 8c3dc2c8d..7925ae3f4 100644 --- a/tests/test_segmented_kv_attention.py +++ b/tests/test_segmented_kv_attention.py @@ -110,11 +110,14 @@ def run(cache): @pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") -def test_without_the_route_the_gather_fallback_is_the_stock_result(monkeypatch) -> None: - """No packed lane configured: segments gather into one array and the stock SDPA runs on it.""" +def test_with_the_route_switched_off_the_gather_fallback_is_the_stock_result(monkeypatch) -> None: + """Both attention switches set to 0 (the kill switch): segments gather into one array and the + stock SDPA runs on it. The support check keeps such a model on stock caches; this pins the + fallback's correctness.""" for name in ("MTPLX_GQA_PACKED_SDPA", "MTPLX_NAX_FLASH_ROUTE"): - monkeypatch.delenv(name, raising=False) + monkeypatch.setenv(name, "0") monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "gather") # the LSE prefill route is not bitwise stock model, attn = _model() configure_split_full_attention(model) stock, _ = _run_turns(attn, KVCache(), False) @@ -123,6 +126,41 @@ def test_without_the_route_the_gather_fallback_is_the_stock_result(monkeypatch) assert mx.array_equal(a, b).item() +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +@pytest.mark.parametrize("profile_env", [{}, {"MTPLX_VLLM_METAL_PAGED_ATTN": "1"}]) +def test_without_the_attention_switches_the_segment_kernel_serves(monkeypatch, profile_env) -> None: + """The sustained profile sets neither MTPLX_GQA_PACKED_SDPA nor MTPLX_NAX_FLASH_ROUTE: a segmented + cache still takes the segment kernel and never gathers, and agrees with the stock cache.""" + import mtplx.segmented_kv as module + + for name in ("MTPLX_GQA_PACKED_SDPA", "MTPLX_NAX_FLASH_ROUTE", "MTPLX_VLLM_METAL_PAGED_ATTN"): + monkeypatch.delenv(name, raising=False) + for name, value in profile_env.items(): + monkeypatch.setenv(name, value) + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "lse") + model, attn = _model() + configure_split_full_attention(model) + stock, _ = _run_turns(attn, KVCache(), False) + S.segmented_dispatch_counts.clear() + module.route_counts.clear() + seg, _ = _run_turns(attn, SegmentedKVCache(), False) + assert not any(k.startswith("gather") for k in module.route_counts), module.route_counts + assert S.segmented_dispatch_counts.get("dispatched", 0) >= 3 + for a, b in zip(stock, seg): + assert bool(mx.all(mx.isfinite(b)).item()) + assert _top_ulps(a, b) <= 8 + + +@pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") +@pytest.mark.parametrize("head_dim, gqa", [(256, 6), (128, 4), (128, 1)]) +def test_the_selfcheck_probe_passes_on_the_real_segment_kernel(head_dim, gqa) -> None: + from mtplx.kernel_selfcheck import _SDPA_TOLERANCE, _check_segments + + for dtype in (mx.bfloat16, mx.float16): + assert _check_segments(mx, dtype, d=head_dim, gqa=gqa) <= _SDPA_TOLERANCE + + @pytest.mark.skipif(not mx.metal.is_available() or not nax_available(), reason="TensorOps unavailable") def test_lse_prefill_route_matches_the_gather_route(lane, monkeypatch) -> None: monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "gather") diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index fccafff1d..a6e7e404f 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -25,7 +25,9 @@ def _cpu_and_clean(monkeypatch): previous = mx.default_device() mx.set_default_device(mx.cpu) reset_model_support() - monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") + # The sustained profile's attention switches: neither is set. + monkeypatch.delenv("MTPLX_GQA_PACKED_SDPA", raising=False) + monkeypatch.delenv("MTPLX_NAX_FLASH_ROUTE", raising=False) monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: True) yield @@ -52,7 +54,6 @@ def test_a_supported_model_stays_enabled() -> None: ([_attn(head_dim=64, q=16, kv=4)], None, "head_dim_64_unsupported"), ([_attn(head_dim=256, q=40, kv=1)], None, "gqa_40_1_unsupported"), ([_attn()], [False], "attention_not_hooked"), - ([_attn(packed=False)], None, "gqa_packed_sdpa_off"), ([], None, "no_full_attention_layers"), ([SimpleNamespace(head_dim=None)], None, "attention_shape_unknown"), ], @@ -72,10 +73,28 @@ def test_the_log_line_appears_once(capsys) -> None: assert capsys.readouterr().out.count("keeping the stock KV cache") == 1 -def test_missing_flash_route_or_nax_is_a_reason(monkeypatch) -> None: - monkeypatch.delenv("MTPLX_NAX_FLASH_ROUTE") - assert "nax_flash_route_off" in evaluate_model_support([_attn()])["reasons"] - monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") +@pytest.mark.parametrize("value", [None, "1"]) +def test_the_attention_switches_unset_or_on_leave_the_verdict_supported(monkeypatch, value) -> None: + """Segments run their own kernel under any profile: sustained sets neither switch.""" + for name in ("MTPLX_GQA_PACKED_SDPA", "MTPLX_NAX_FLASH_ROUTE"): + if value is not None: + monkeypatch.setenv(name, value) + assert evaluate_model_support([_attn(packed=False)])["supported"] is True + assert segmented_kv_enabled() is True + + +@pytest.mark.parametrize( + "name, reason", [("MTPLX_GQA_PACKED_SDPA", "gqa_packed_sdpa_off"), ("MTPLX_NAX_FLASH_ROUTE", "nax_flash_route_off")] +) +@pytest.mark.parametrize("value", ["0", "off"]) +def test_an_attention_switch_set_to_0_is_a_kill_switch(monkeypatch, name, reason, value) -> None: + monkeypatch.setenv(name, value) + verdict = evaluate_model_support([_attn()]) + assert not verdict["supported"] and verdict["reasons"] == [reason] + assert segmented_kv_enabled() is False + + +def test_missing_nax_is_a_reason(monkeypatch) -> None: monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: False) assert "nax_unavailable" in evaluate_model_support([_attn()])["reasons"] @@ -150,8 +169,6 @@ def test_configure_split_full_attention_decides_from_the_hooked_layers(monkeypat from mtplx.attention_split import configure_split_full_attention - monkeypatch.setenv("MTPLX_GQA_PACKED_SDPA", "1") - def model(head_dim): args = SimpleNamespace( hidden_size=64, num_attention_heads=4, num_key_value_heads=2, head_dim=head_dim, attention_bias=False, @@ -174,7 +191,7 @@ def model(head_dim): def test_head_dim_is_read_from_the_scale_when_the_module_has_no_attribute() -> None: """mlx-lm's Qwen3 attention (head_dim 128) keeps n_heads, n_kv_heads and scale only.""" - attn = SimpleNamespace(n_heads=32, n_kv_heads=8, scale=128**-0.5, _mtplx_gqa_packed_sdpa_enabled=True) + attn = SimpleNamespace(n_heads=32, n_kv_heads=8, scale=128**-0.5) verdict = evaluate_model_support([attn]) assert verdict["supported"], verdict assert verdict["head_dims"] == [128] and verdict["gqa"] == [4] @@ -234,13 +251,7 @@ def test_the_hook_covers_plain_attention_in_the_support_check() -> None: attn = _qwen3_attention(head_dim=128, q=8, kv=2, hidden=128) model = SimpleNamespace(model=SimpleNamespace(layers=[SimpleNamespace(is_linear=False, self_attn=attn)])) - import os - - os.environ["MTPLX_GQA_PACKED_SDPA"] = "1" - try: - stats = configure_split_full_attention(model) - finally: - os.environ.pop("MTPLX_GQA_PACKED_SDPA") + stats = configure_split_full_attention(model) assert stats["segmented_kv_supported"] is True assert model_support()["head_dims"] == [128] and model_support()["gqa"] == [4] @@ -325,7 +336,144 @@ def test_the_hook_covers_llama_style_attention_in_the_support_check(monkeypatch) attn = _llama_attention(head_dim=128, q=8, kv=2, hidden=128) model = SimpleNamespace(model=SimpleNamespace(layers=[SimpleNamespace(self_attn=attn)])) - monkeypatch.setenv("MTPLX_GQA_PACKED_SDPA", "1") stats = configure_split_full_attention(model) assert stats["segmented_kv_supported"] is True assert model_support()["head_dims"] == [128] and model_support()["gqa"] == [4] + + +# ---- the kernel self-check decides before the first cache (any profile) -------------------- + + +@pytest.fixture +def selfcheck(): + from mtplx import kernel_selfcheck + + kernel_selfcheck._reset_for_tests() + yield kernel_selfcheck + kernel_selfcheck._reset_for_tests() + + +def test_requesting_segments_runs_the_selfcheck_under_a_profile_without_turbo_lanes(selfcheck, monkeypatch) -> None: + monkeypatch.delenv("MTPLX_KERNEL_SELFCHECK", raising=False) + for name in ("MTPLX_NAX_VERIFY", "MTPLX_QWEN_ROW_OWNED_ROUTER", "MTPLX_QWEN_COMBINE_TAIL", + "MTPLX_FUSE_GDN_POST_CONV", "MTPLX_A3B_WHOLE_MOE_FUSION"): + monkeypatch.delenv(name, raising=False) + assert selfcheck.selfcheck_enabled() is True + monkeypatch.delenv("MTPLX_SEGMENTED_KV") + assert selfcheck.selfcheck_enabled() is False # flag off: unchanged + + +def test_the_selfcheck_probes_the_segment_kernel_at_the_models_shapes(selfcheck, monkeypatch) -> None: + seen = [] + monkeypatch.setattr(selfcheck, "_check_segments", lambda mx, dtype, *, d, gqa: seen.append((d, gqa)) or 0.0) + evaluate_model_support([_attn(head_dim=128, q=32, kv=8)]) + report = selfcheck.run_kernel_selfcheck(mx.bfloat16, 4, 64) + assert seen == [(128, 4)] + assert report["lanes"]["segmented_kv_sdpa"] == "ok" + assert segmented_kv_enabled() is True + + +@pytest.mark.parametrize("failure", ["mismatch", "raises", "bails"]) +def test_a_failed_segment_kernel_means_stock_caches_with_a_reason(selfcheck, monkeypatch, capsys, failure) -> None: + from mtplx.segmented_kv import SegmentedKVCache + + def probe(mx, dtype, *, d, gqa): + if failure == "raises": + raise RuntimeError("no such kernel") + return 1.0 if failure == "mismatch" else float("inf") + + monkeypatch.setattr(selfcheck, "_check_segments", probe) + evaluate_model_support([_attn()]) + assert segmented_kv_enabled() is True + report = selfcheck.run_kernel_selfcheck(mx.bfloat16, 4, 64) + assert report["lanes"]["segmented_kv_sdpa"] == "fallback" + assert segmented_kv_enabled() is False and segmented_ssd_enabled() is False + assert "segment_kernel_selfcheck_failed" in model_support()["reasons"] + cache = [KVCache(), KVCache()] + configure_tail_owned_attention_kv_cache(cache) + assert all(type(c) is KVCache for c in cache) and not any(isinstance(c, SegmentedKVCache) for c in cache) + block = segmented_kv_health([]) + assert block["enabled"] is False and "segment_kernel_selfcheck_failed" in block["model_support"]["reasons"] + selfcheck.run_kernel_selfcheck(mx.bfloat16, 4, 64) # a second run (a later selfcheck) keeps it off + assert segmented_kv_enabled() is False + assert capsys.readouterr().out.count("MTPLX_SEGMENTED_KV is on but this model cannot use it") == 1 + + +def test_a_selfcheck_that_cannot_run_leaves_segments_off(selfcheck, monkeypatch) -> None: + monkeypatch.setattr(selfcheck, "_model_quant_signature", lambda model: (_ for _ in ()).throw(RuntimeError("x"))) + evaluate_model_support([_attn()]) + assert selfcheck.maybe_run_model_selfcheck(object()) is None + assert segmented_kv_enabled() is False + assert "segment_kernel_selfcheck_failed" in model_support()["reasons"] + + +def test_with_the_flag_off_the_selfcheck_skips_the_segment_lane_and_leaves_no_verdict(selfcheck, monkeypatch) -> None: + monkeypatch.delenv("MTPLX_SEGMENTED_KV") + monkeypatch.setattr(selfcheck, "_check_segments", lambda *a, **k: pytest.fail("probed with the flag off")) + evaluate_model_support([_attn()]) + report = selfcheck.run_kernel_selfcheck(mx.bfloat16, 4, 64) + assert report["lanes"]["segmented_kv_sdpa"] == "skipped" + assert segmented_kv_enabled() is False + + +# ---- the hook routes a segmented cache to the segment route under any profile ---------------- + + +def _qwen3_next_attention(): + from mlx_lm.models.qwen3_next import Qwen3NextAttention + + args = SimpleNamespace( + hidden_size=64, num_attention_heads=4, num_key_value_heads=2, head_dim=32, attention_bias=False, + rms_norm_eps=1e-6, partial_rotary_factor=0.25, rope_theta=10000.0, rope_scaling=None, + max_position_embeddings=1024, + ) + mx.random.seed(0) + return Qwen3NextAttention(args) + + +@pytest.mark.parametrize("make", [_qwen3_next_attention, _qwen3_attention, _llama_attention]) +@pytest.mark.parametrize("profile_env", [{}, {"MTPLX_VLLM_METAL_PAGED_ATTN": "1"}]) +def test_a_segmented_cache_takes_the_segment_route_without_the_attention_switches(monkeypatch, make, profile_env) -> None: + """Sustained-like: no packed SDPA, no flash route, and either no attention route at all or only + the paged one. A verify window over several segments goes to decode_segments_attention (never + the stock forward or a gather); its output matches the stock cache's.""" + import mtplx.attention_split as split + from mtplx.attention_split import configure_split_full_attention + from mtplx.segmented_kv import SegmentedKVCache, gathered_view + + for name, value in profile_env.items(): + monkeypatch.setenv(name, value) + attn = make() + configure_split_full_attention(SimpleNamespace(model=SimpleNamespace(layers=[SimpleNamespace(is_linear=False, self_attn=attn)]))) + assert attn._mtplx_gqa_packed_sdpa_enabled is False + assert attn._mtplx_split_full_attention_enabled is bool(profile_env) # route off, or the paged route only + calls = [] + + def segment_route(queries, cache, *, scale, mask, packed_threshold): + # CPU stand-in for the kernel: exact attention over the rows, read without a gather count. + calls.append((int(queries.shape[2]), cache.segment_count)) + view = gathered_view(cache) + q = int(queries.shape[2]) + causal = mx.arange(view.offset)[None, :] <= (view.offset - q + mx.arange(q)[:, None]) + return mx.fast.scaled_dot_product_attention(queries, view.keys, view.values, scale=scale, mask=causal) + + monkeypatch.setattr(split, "decode_segments_attention", segment_route) + + def run(cache): + outs = [attn(mx.random.normal((1, 40, 64), key=mx.random.key(1)), mask="causal", cache=cache)] + if isinstance(cache, SegmentedKVCache): + state = cache.state # turn end: the history becomes a sealed segment + cache = SegmentedKVCache() + cache.state = state + outs.append(attn(mx.random.normal((1, 9, 64), key=mx.random.key(2)), mask="causal", cache=cache)) + outs.append(attn(mx.random.normal((1, 4, 64), key=mx.random.key(3)), mask="causal", cache=cache)) + mx.eval(outs) + return outs + + stock = run(KVCache()) + module.route_counts.clear() + seg = run(SegmentedKVCache()) + assert (4, 2) in calls and (9, 2) in calls + assert not any(k.startswith("gather_q") and k[8:].isdigit() for k in module.route_counts), module.route_counts + for a, b in zip(stock, seg): + assert float(mx.abs(a - b).max().item()) < 1e-4 From 0ae05ae4a78f7f9169a814d6e44dcf25d2ccdb8d Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Sat, 10 Oct 2026 00:12:03 +0200 Subject: [PATCH 30/37] Changelog: segmented KV under any profile (attention switches only as kill switches, self-check on the segment kernel, Qwen3.6-35B-A3B sustained numbers) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eaf2482f8..ff297b749 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ All notable user-facing changes to MTPLX. The format is based on - **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, the NAX flash route on). 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. + - 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. - 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. From 7be91ff2e61230843d16ee2122316acb0c6211e2 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Sat, 10 Oct 2026 09:39:42 +0200 Subject: [PATCH 31/37] Segmented KV: no dense-decode ceiling, a long conversation stays on segments instead of repaging Past the dense-decode ceiling (auto: 15 % of RAM over 64 KiB a token, 157,286 tokens on a 64 GB Mac) the auto layout repaged the prefill cache into VllmMetalPagedKVCache, which is built from each layer's ``keys``: on a segment cache that gathers every segment into one paged copy of the history. The admission priced that copy correctly and refused the turn. 2026-10-10, Qwen3.8-27B, 64 GB: a follow-up turn of 161,835 tokens (161,751 restored, 84 new) was priced 11.6 GiB of growth for the repage (repage_copy_bytes = (prompt + 16,387 reserved rows) x 68 KiB) and refused with 10.1 GiB free, eleven retries in a row; the turn before it had decoded from 157,089 to 161,751 tokens on segments at 19 tok/s. With segments on (and KV unquantized) the auto layout is contiguous_dense_decode at any length, so the restore keeps its segments, nothing repages and the admission prices the new rows only (about 1.1 GiB for that turn). The stock cache, quantized KV and an explicit MTPLX_SUSTAINED_PREFILL_LAYOUT keep their repage. --- mtplx/generation.py | 12 +++ tests/test_segmented_kv_layout.py | 125 ++++++++++++++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 tests/test_segmented_kv_layout.py diff --git a/mtplx/generation.py b/mtplx/generation.py index ebbd8eca7..e1d6e0863 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -2447,6 +2447,18 @@ def _sustained_prefill_layout(context_tokens: int | None = None) -> str: if paged_kv_quant_mode_from_env() != "off": return "contiguous_then_repage" + from .segmented_kv import segmented_kv_enabled + + if segmented_kv_enabled(): + # Segments are the long-context layout, so there is no dense-decode + # ceiling: past it the repage would gather every segment into a paged + # copy of the history (VllmMetalPagedKVCache.from_cache reads + # ``keys``), the copy segments exist to avoid. 2026-10-10, + # Qwen3.8-27B on a 64 GB Mac (ceiling 157,286 tokens): a follow-up + # turn of 161,835 tokens, 84 of them new, was priced 11.6 GiB for + # that repage and refused with 10.1 GiB free, after the turn before + # it had decoded to 161,751 tokens on segments. + return "contiguous_dense_decode" if context_tokens is None: context_tokens = _env_int("MTPLX_CURRENT_PREFILL_CONTEXT_TOKENS", 0) dense_max = _dense_decode_max_context() diff --git a/tests/test_segmented_kv_layout.py b/tests/test_segmented_kv_layout.py new file mode 100644 index 000000000..72d84fb33 --- /dev/null +++ b/tests/test_segmented_kv_layout.py @@ -0,0 +1,125 @@ +"""A segmented conversation past the dense-decode ceiling stays on segments (MTPLX_SEGMENTED_KV). + +2026-10-10, Qwen3.8-27B on a 64 GB Mac: the auto ceiling is 157,286 tokens. A follow-up turn of +161,835 tokens (161,751 restored, 84 new) was laid out ``contiguous_then_repage``; the repage +builds paged caches from each layer's ``keys``, which a segment cache answers by gathering every +segment, and the admission priced that copy at 11.6 GiB and refused the turn with 10.1 GiB free. +""" + +from __future__ import annotations + +import pytest + +import mtplx.segmented_kv as segmented_kv_module +from mtplx.generation import ( + _maybe_repage_target_prefill_cache, + _sustained_prefill_layout, + prefill_cache_layout, +) +from mtplx.segmented_kv import SegmentedKVCache + +GIB = 1024**3 +CEILING = 157_286 +PROMPT, RESTORED = 161_835, 161_751 +# Qwen3.8-27B: 16 full-attention layers x K+V x 4 KV heads x 256 x bf16, plus 4 KiB of MTP history. +ROW = 65_536 + 4_096 +OUTPUT_ROWS = 16_384 + 3 + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setenv("MTPLX_SUSTAINED_PREFILL_LAYOUT", "auto") + monkeypatch.setenv("MTPLX_SUSTAINED_DENSE_DECODE_MAX_CONTEXT", str(CEILING)) + monkeypatch.delenv("MTPLX_VLLM_METAL_PAGED_KV_QUANT", raising=False) + monkeypatch.delenv("MTPLX_PAGED_KV_QUANT", raising=False) + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + # segmented_kv_enabled() needs a verdict on the loaded model (without one it stays off). + monkeypatch.setattr(segmented_kv_module, "_MODEL_SUPPORT", {"supported": True, "reasons": []}) + + +def _switch_off(monkeypatch) -> None: + monkeypatch.delenv("MTPLX_SEGMENTED_KV", raising=False) + + +def test_segments_keep_the_dense_layout_past_the_ceiling(monkeypatch) -> None: + assert _sustained_prefill_layout(CEILING) == "contiguous_dense_decode" + assert _sustained_prefill_layout(PROMPT) == "contiguous_dense_decode" + assert prefill_cache_layout(None, PROMPT) == "contiguous_dense_decode" + monkeypatch.setenv("MTPLX_CURRENT_PREFILL_CONTEXT_TOKENS", str(PROMPT)) + assert _sustained_prefill_layout() == "contiguous_dense_decode" + + +def test_the_stock_cache_still_repages_past_the_ceiling(monkeypatch) -> None: + _switch_off(monkeypatch) + assert _sustained_prefill_layout(CEILING) == "contiguous_dense_decode" + assert _sustained_prefill_layout(PROMPT) == "contiguous_then_repage" + monkeypatch.setenv("MTPLX_SEGMENTED_KV", "1") + monkeypatch.setattr(segmented_kv_module, "_MODEL_SUPPORT", None) # no verdict: stock + assert _sustained_prefill_layout(PROMPT) == "contiguous_then_repage" + + +def test_quantized_kv_keeps_its_repage_with_segments(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_PAGED_KV_QUANT", "q8") + assert _sustained_prefill_layout(PROMPT) == "contiguous_then_repage" + + +def test_an_explicit_layout_is_left_alone(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_SUSTAINED_PREFILL_LAYOUT", "contiguous_then_repage") + assert _sustained_prefill_layout(PROMPT) == "contiguous_then_repage" + + +class _Runtime: + def __init__(self) -> None: + self.repages = 0 + + def repage_target_prefill_cache(self, cache) -> bool: + self.repages += 1 + return True + + +def test_a_segmented_prefill_past_the_ceiling_is_not_repaged(monkeypatch) -> None: + monkeypatch.setenv("MTPLX_CURRENT_PREFILL_CONTEXT_TOKENS", str(PROMPT)) + runtime = _Runtime() + cache = [SegmentedKVCache()] + assert _maybe_repage_target_prefill_cache(runtime, cache) == 0.0 + assert runtime.repages == 0 + assert isinstance(cache[0], SegmentedKVCache) + + +def _incident_growth(layout: str) -> dict: + from mtplx.server.openai import _admission_growth, _AdmissionGeometry + + geometry = _AdmissionGeometry( + live_bytes_per_token=ROW, + paged_bytes_per_token=ROW, + context_transient_bytes_per_token=0, + flat_transient_bytes=0, + weights_bytes=0, + aux_bytes_per_token=4_096, + ) + return _admission_growth( + geometry, + prompt_tokens=PROMPT, + reused_tokens=RESTORED, + restore_copies_prefix=False, + layout=layout, + source_layout=None, + output_tokens=OUTPUT_ROWS, + publish=True, + scratch_bytes=GIB // 2, + lease={"paged": False}, # the live segmented cache the lease extends + segmented_kv=True, + ) + + +def test_the_incident_turn_is_priced_without_a_history_copy() -> None: + model = _incident_growth(prefill_cache_layout(None, PROMPT)) + assert model["repage_copy_bytes"] == 0 + assert model["restore_copy_bytes"] == 0 + # The new rows, the answer's rows the snapshot seals, and the prefill scratch: about 1 GiB. + assert model["publish_copy_bytes"] == (PROMPT - RESTORED + OUTPUT_ROWS) * ROW + assert model["growth_bytes"] < 2 * GIB + # The repage layout the turn got: a paged copy of the whole history. + repaged = _incident_growth("contiguous_then_repage") + assert repaged["repage_copy_bytes"] == (PROMPT + OUTPUT_ROWS) * ROW + assert repaged["growth_bytes"] > 11 * GIB From 31464dbbd605b9aa4272e26289b32087bf47ff54 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Sat, 10 Oct 2026 09:39:55 +0200 Subject: [PATCH 32/37] Changelog: segmented KV stays on segments past the dense-decode ceiling --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ff297b749..f4a251dea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ All notable user-facing changes to MTPLX. The format is based on - 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. From 434800528eaec4195fc5e7280826dbd865520eea Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Tue, 13 Oct 2026 14:16:14 +0200 Subject: [PATCH 33/37] Segmented KV: use mx.fast.scaled_dot_product_attention_with_lse when present, keep return_lse as fallback All fused logsumexp calls go through sdpa_with_lse; the probe records which API answered and /health reports it as lse_api. The per-segment merge and the reference route no longer produce NaN for rows without a visible key (lse -inf): such rows get weight 0 and merge to zeros. --- mtplx/segmented_kv.py | 77 +++++++++++++++++++++++++++++--------- tests/test_segmented_kv.py | 76 +++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 17 deletions(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index 6d0649099..a26d1aad8 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -167,7 +167,7 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No _LSE_NOTE_LOGGED = True print( "[mtplx] MTPLX_SEGMENTED_KV is on, but this MLX build has no logsumexp output on the fused " - "SDPA (return_lse): a prefill chunk over existing segments gathers the history first " + "SDPA (neither mx.fast.scaled_dot_product_attention_with_lse nor return_lse): a prefill chunk over existing segments gathers the history first " "(measured about 5 % slower prefill on follow-up turns; decode and verify are unaffected)", flush=True, ) @@ -233,22 +233,58 @@ def count_route(route: str) -> None: _SDPA_LSE: bool | None = None +_SDPA_LSE_API: str | None = None + +#: The two MLX spellings of "fused SDPA that also returns the logsumexp", in order of preference. +LSE_API_FUNCTION = "scaled_dot_product_attention_with_lse" +LSE_API_KEYWORD = "scaled_dot_product_attention(return_lse=True)" + + +def _lse_api_name() -> str | None: + """Which logsumexp API this MLX build offers: the separate function (proposed upstream + public API) when present, else the ``return_lse=True`` keyword of the patched build, else None. + The keyword form cannot be told apart by attribute: ``_probe_sdpa_lse`` runs it.""" + if hasattr(mx.fast, LSE_API_FUNCTION): + return LSE_API_FUNCTION + return LSE_API_KEYWORD + + +def sdpa_with_lse(queries, keys, values, *, scale, mask=None): + """Fused SDPA returning ``(out, lse)``; lse is natural log, float32, (B, Hq, Q, 1). + + The one place that knows the API. Prefers ``mx.fast.scaled_dot_product_attention_with_lse`` + and falls back to ``scaled_dot_product_attention(..., return_lse=True)``. Stock released + MLX raises TypeError on the keyword, which callers treat as "no fused lse". + """ + kwargs = {} if mask is None else {"mask": mask} + function = getattr(mx.fast, LSE_API_FUNCTION, None) + if function is not None: + return function(queries, keys, values, scale=scale, **kwargs) + return mx.fast.scaled_dot_product_attention( + queries, keys, values, scale=scale, return_lse=True, **kwargs + ) def _probe_sdpa_lse() -> bool: - """Whether this MLX build's fused SDPA can return the logsumexp (fork branch sdpa-lse). + """Whether this MLX build's fused SDPA can return the logsumexp. - The one place that knows the API: ``mx.fast.scaled_dot_product_attention(..., return_lse=True)`` - returning ``(out, lse)``. Stock released MLX raises TypeError on the keyword, which is the - "absent" answer; any other failure also counts as absent. + Records which API answered in ``_SDPA_LSE_API`` (``lse_api_found``). Any failure counts as absent. """ + global _SDPA_LSE_API try: q = mx.zeros((1, 1, 1, 64), dtype=mx.bfloat16) - result = mx.fast.scaled_dot_product_attention(q, q, q, scale=1.0, return_lse=True) + result = sdpa_with_lse(q, q, q, scale=1.0) mx.eval(result) - return isinstance(result, (tuple, list)) and len(result) == 2 + ok = isinstance(result, (tuple, list)) and len(result) == 2 except Exception: # noqa: BLE001 - return False + ok = False + _SDPA_LSE_API = _lse_api_name() if ok else None + return ok + + +def lse_api_found() -> str | None: + """Name of the logsumexp API the probe found (None when absent or not probed yet).""" + return _SDPA_LSE_API if sdpa_lse_available() else None def sdpa_lse_available() -> bool: @@ -690,7 +726,8 @@ def _segment_sdpa_lse_reference( visible = mx.arange(int(n))[None, :] <= limit[:, None] scores = mx.where(visible[None, None], scores, -mx.inf) lse = mx.logsumexp(scores, axis=-1) - p = mx.exp(scores - lse[..., None]) + # A row without any visible key has lse -inf: weights 0 and output 0 (not -inf - -inf = NaN). + p = mx.exp(scores - _finite_or_zero(lse)[..., None]) out = (p @ v).reshape(b, hq, rows, d) outs.append(out.astype(queries.dtype)) lses.append(lse.reshape(b, hq, rows)) @@ -700,12 +737,9 @@ def _segment_sdpa_lse_reference( def _segment_sdpa_lse_fast(queries, keys, values, n, *, scale, causal): - """Per-segment attention on the MLX kernel that returns the logsumexp (fork only).""" + """Per-segment attention on the MLX kernel that returns the logsumexp (``sdpa_with_lse``).""" k, v = _rows_view(keys, n), _rows_view(values, n) - kwargs = {"mask": "causal"} if causal else {} - out, lse = mx.fast.scaled_dot_product_attention( - queries, k, v, scale=scale, return_lse=True, **kwargs - ) + out, lse = sdpa_with_lse(queries, k, v, scale=scale, mask="causal" if causal else None) # The kernel's logsumexp is (B, Hq, Q, 1); the merge works on (B, Hq, Q). return out, lse.astype(mx.float32).reshape(out.shape[:-1]) @@ -726,16 +760,24 @@ def segment_sdpa_lse(queries, keys, values, n, *, scale, causal): return _segment_sdpa_lse_reference(queries, keys, values, n, scale=scale, causal=causal) +def _finite_or_zero(x: mx.array) -> mx.array: + """x with +-inf replaced by 0, so ``exp(lse - x)`` is 0 instead of NaN for lse = -inf.""" + return mx.where(mx.isinf(x), 0.0, x) + + def merge_segment_outputs(outs: Sequence[mx.array], lses: Sequence[mx.array]) -> mx.array: """Exact merge of per-segment (out, lse): sum_s exp(lse_s - L) * out_s with L = logsumexp(lse).""" if len(outs) == 1: return outs[0] stacked = mx.stack([l.astype(mx.float32) for l in lses], axis=0) - total = mx.logsumexp(stacked, axis=0) + total = _finite_or_zero(mx.logsumexp(stacked, axis=0)) merged = None for out, lse in zip(outs, lses): - weight = mx.exp(lse.astype(mx.float32) - total)[..., None] - term = out.astype(mx.float32) * weight + # lse -inf (no visible key in this segment): weight 0 and the segment's out is ignored, + # even if a kernel left NaN there. A row that is -inf in every segment merges to zeros. + lse = lse.astype(mx.float32) + weight = mx.exp(lse - total)[..., None] + term = mx.where(mx.isinf(lse)[..., None], 0.0, out.astype(mx.float32) * weight) merged = term if merged is None else merged + term return merged.astype(outs[0].dtype) @@ -1037,6 +1079,7 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: "model_support": model_support(), "ssd": segmented_ssd_enabled(), "prefill_route": prefill_route(), + "lse_api": lse_api_found(), "entries": per_entry, "max_segments": max((e["segments"] for e in per_entry), default=0), "sealed_bytes_unique": unique_state_bytes(all_states), diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py index 35c9025f2..8fde138fd 100644 --- a/tests/test_segmented_kv.py +++ b/tests/test_segmented_kv.py @@ -484,3 +484,79 @@ def test_real_lse_kernel_matches_the_reference_route() -> None: keys, values = cache.gather() want = mx.fast.scaled_dot_product_attention(q, keys, values, scale=0.25, mask="causal") assert float(mx.abs(got.astype(mx.float32) - want.astype(mx.float32)).max()) < 2e-2 + + +def _fake_lse_api(calls, tag): + def fake(q, k, v, *, scale, mask=None, **extra): + calls.append((tag, mask, extra)) + out = mx.zeros(q.shape, dtype=q.dtype) + return out, mx.zeros((*q.shape[:-1], 1), dtype=mx.float32) + + return fake + + +def test_sdpa_with_lse_prefers_the_separate_function(monkeypatch) -> None: + import mtplx.segmented_kv as module + + calls: list = [] + monkeypatch.setattr(mx.fast, "scaled_dot_product_attention_with_lse", _fake_lse_api(calls, "func"), raising=False) + monkeypatch.setattr(mx.fast, "scaled_dot_product_attention", _fake_lse_api(calls, "kw")) + q = mx.zeros((1, 1, 1, 64), dtype=mx.bfloat16) + module.sdpa_with_lse(q, q, q, scale=1.0, mask="causal") + module.sdpa_with_lse(q, q, q, scale=1.0) + assert calls == [("func", "causal", {}), ("func", None, {})] + monkeypatch.setattr(module, "_SDPA_LSE", None) + assert module.sdpa_lse_available() is True + assert module.lse_api_found() == module.LSE_API_FUNCTION + + +def test_sdpa_with_lse_falls_back_to_the_return_lse_keyword(monkeypatch) -> None: + import mtplx.segmented_kv as module + + calls: list = [] + monkeypatch.delattr(mx.fast, "scaled_dot_product_attention_with_lse", raising=False) + monkeypatch.setattr(mx.fast, "scaled_dot_product_attention", _fake_lse_api(calls, "kw")) + q = mx.zeros((1, 1, 1, 64), dtype=mx.bfloat16) + module.sdpa_with_lse(q, q, q, scale=1.0, mask="causal") + assert calls == [("kw", "causal", {"return_lse": True})] + monkeypatch.setattr(module, "_SDPA_LSE", None) + assert module.sdpa_lse_available() is True + assert module.lse_api_found() == module.LSE_API_KEYWORD + + +def test_no_lse_api_reports_unavailable(monkeypatch) -> None: + import mtplx.segmented_kv as module + + monkeypatch.delattr(mx.fast, "scaled_dot_product_attention_with_lse", raising=False) + + def stock(q, k, v, *, scale, mask=None): + raise TypeError("unexpected keyword return_lse") + + monkeypatch.setattr(mx.fast, "scaled_dot_product_attention", stock) + monkeypatch.setattr(module, "_SDPA_LSE", None) + assert module.sdpa_lse_available() is False + assert module.lse_api_found() is None + + +def test_merge_of_fully_masked_rows_is_zero_not_nan() -> None: + from mtplx.segmented_kv import merge_segment_outputs + + ninf = float("-inf") + out_a = mx.array([[[[1.0, 2.0], [float("nan"), float("nan")], [3.0, 4.0]]]]) + out_b = mx.array([[[[5.0, 6.0], [float("nan"), float("nan")], [0.0, 0.0]]]]) + lse_a = mx.array([[[0.0, ninf, 0.0]]]) + lse_b = mx.array([[[0.0, ninf, ninf]]]) + merged = merge_segment_outputs([out_a, out_b], [lse_a, lse_b]).astype(mx.float32) + assert merged.tolist() == [[[[3.0, 4.0], [0.0, 0.0], [3.0, 4.0]]]] + + +def test_reference_route_fully_masked_row_is_zero_not_nan() -> None: + import mtplx.segmented_kv as module + + q = mx.random.normal((1, 2, 4, 64), key=mx.random.key(1)).astype(mx.bfloat16) + k = mx.random.normal((1, 2, 2, 64), key=mx.random.key(2)).astype(mx.bfloat16) + # 4 queries over a 2-row causal segment: the first two rows see nothing. + out, lse = module._segment_sdpa_lse_reference(q, k, k, 2, scale=0.125, causal=True) + assert lse[..., :2].tolist() == [[[float("-inf")] * 2] * 2] + assert not mx.isnan(out.astype(mx.float32)).any().item() + assert float(mx.abs(out[:, :, :2].astype(mx.float32)).max()) == 0.0 From 257af12a1ca4352725e7ff9edd0b85a18c4c8fb2 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Wed, 14 Oct 2026 06:47:47 +0200 Subject: [PATCH 34/37] Segmented KV lse route: keep per-segment outputs in float32, round once after the merge The reference route rounded every segment's output to the query dtype before the logsumexp merge, and the merge rounded again. On Qwen3.8-27B tensors at 20K (one verify window, 16 layers) that put 26 to 29% of the attention elements one ulp away from the exact result rounded once, against 15 to 22% for the stock kernel and 3.5 to 7% with a single rounding. The per-segment outputs now stay float32 and merge_segment_outputs rounds once to the query dtype (MTPLX_SEGMENTED_KV_PREFILL=lse, reference route; the fused sdpa_with_lse kernel still returns the query dtype per segment). The existing exactness test compared bfloat16 output with a loose tolerance and could not see the second rounding; the new tests require the route to equal the float32 attention rounded once and to be no further from stock SDPA than that exact result is. --- mtplx/segmented_kv.py | 23 +++++++++---- tests/test_segmented_kv.py | 70 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 7 deletions(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index a26d1aad8..b32b1fc2c 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -699,7 +699,10 @@ def _segment_sdpa_lse_reference( ) -> tuple[mx.array, mx.array]: """Exact attention of ``queries`` over the first ``n`` rows of one segment, with its logsumexp. - ``queries`` (1, Hq, Q, D); returns (out in the query dtype, lse float32 (1, Hq, Q)). + ``queries`` (1, Hq, Q, D); returns (out float32, lse float32 (1, Hq, Q)). The output stays + float32 so that the merge rounds once, like the stock attention, instead of rounding every + segment to the query dtype first (that second rounding put about one element in three one + bf16 ulp away from the stock output at 20K, against about one in six for the exact result). ``causal=True`` marks the segment holding the Q new rows as its last Q rows (query row i sees keys up to ``n - Q + i``); otherwise every key is visible. Scores run in float32 and are bounded by chunking the query rows: this is the slow, exact reference route (QK^T plus @@ -729,7 +732,7 @@ def _segment_sdpa_lse_reference( # A row without any visible key has lse -inf: weights 0 and output 0 (not -inf - -inf = NaN). p = mx.exp(scores - _finite_or_zero(lse)[..., None]) out = (p @ v).reshape(b, hq, rows, d) - outs.append(out.astype(queries.dtype)) + outs.append(out) lses.append(lse.reshape(b, hq, rows)) if len(outs) == 1: return outs[0], lses[0] @@ -765,10 +768,16 @@ def _finite_or_zero(x: mx.array) -> mx.array: return mx.where(mx.isinf(x), 0.0, x) -def merge_segment_outputs(outs: Sequence[mx.array], lses: Sequence[mx.array]) -> mx.array: - """Exact merge of per-segment (out, lse): sum_s exp(lse_s - L) * out_s with L = logsumexp(lse).""" +def merge_segment_outputs( + outs: Sequence[mx.array], lses: Sequence[mx.array], *, dtype: mx.Dtype | None = None +) -> mx.array: + """Exact merge of per-segment (out, lse): sum_s exp(lse_s - L) * out_s with L = logsumexp(lse). + + The sum runs in float32 and is rounded once, to ``dtype`` (default: the first output's dtype). + """ + dtype = outs[0].dtype if dtype is None else dtype if len(outs) == 1: - return outs[0] + return outs[0] if outs[0].dtype == dtype else outs[0].astype(dtype) stacked = mx.stack([l.astype(mx.float32) for l in lses], axis=0) total = _finite_or_zero(mx.logsumexp(stacked, axis=0)) merged = None @@ -779,7 +788,7 @@ def merge_segment_outputs(outs: Sequence[mx.array], lses: Sequence[mx.array]) -> weight = mx.exp(lse - total)[..., None] term = mx.where(mx.isinf(lse)[..., None], 0.0, out.astype(mx.float32) * weight) merged = term if merged is None else merged + term - return merged.astype(outs[0].dtype) + return merged.astype(dtype) def attend_segments_lse(queries: mx.array, cache: SegmentedKVCache, *, scale: float) -> mx.array: @@ -792,7 +801,7 @@ def attend_segments_lse(queries: mx.array, cache: SegmentedKVCache, *, scale: fl ) outs.append(out) lses.append(lse) - return merge_segment_outputs(outs, lses) + return merge_segment_outputs(outs, lses, dtype=queries.dtype) # ---- the decode/verify route used by attention_split ----------------------------------- diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py index 8fde138fd..b5c961630 100644 --- a/tests/test_segmented_kv.py +++ b/tests/test_segmented_kv.py @@ -386,6 +386,76 @@ def test_lse_merge_equals_attention_over_the_concatenated_rows() -> None: assert float(mx.abs(got.astype(mx.float32) - ref).max()) < 0.02 +@pytest.fixture +def cpu(): + """The exactness tests run on the CPU: MLX float32 GEMM on a tensor-unit GPU rounds like TF32.""" + previous = mx.default_device() + mx.set_default_device(mx.cpu) + yield + mx.set_default_device(previous) + + +def _rounding_case(sizes, q_len, dim=64, heads_q=8): + cache = SegmentedKVCache() + for i, n in enumerate(sizes[:-1]): + k = mx.random.normal((1, HEADS, n, dim), key=mx.random.key(10 + i)).astype(mx.bfloat16) + v = mx.random.normal((1, HEADS, n, dim), key=mx.random.key(50 + i)).astype(mx.bfloat16) + cache.update_and_fetch(k, v) + cache.seal() + k = mx.random.normal((1, HEADS, sizes[-1], dim), key=mx.random.key(7)).astype(mx.bfloat16) + v = mx.random.normal((1, HEADS, sizes[-1], dim), key=mx.random.key(8)).astype(mx.bfloat16) + cache.update_and_fetch(k, v) + q = mx.random.normal((1, heads_q, q_len, dim), key=mx.random.key(9)).astype(mx.bfloat16) + keys, values = cache.gather() + scale = dim ** -0.5 + kf = mx.repeat(keys.astype(mx.float32), heads_q // HEADS, axis=1) + vf = mx.repeat(values.astype(mx.float32), heads_q // HEADS, axis=1) + s = (q.astype(mx.float32) @ kf.transpose(0, 1, 3, 2)) * scale + total = keys.shape[2] + vis = mx.arange(total)[None, :] <= (total - q_len + mx.arange(q_len))[:, None] + s = mx.where(vis[None, None], s, -mx.inf) + exact = (mx.softmax(s, axis=-1) @ vf).astype(mx.bfloat16) + return cache, q, keys, values, scale, exact + + +def _differing(a, b) -> float: + return float(mx.mean(a.astype(mx.float32) != b.astype(mx.float32)).item()) + + +@pytest.mark.parametrize("q_len", [1, 6]) +@pytest.mark.parametrize("sizes", [[3000, 200], [2000, 1024, 600, 40]]) +def test_reference_lse_route_rounds_once(cpu, monkeypatch, q_len, sizes) -> None: + """The reference route must be the float32 attention rounded once to the query dtype. + + A tolerance of one bf16 ulp does not see a second rounding (every segment rounded to bf16 + before the merge), which puts a quarter or more of the elements one ulp off. + """ + monkeypatch.setattr(segmented_kv_module, "sdpa_lse_available", lambda: False) + cache, q, _, _, scale, exact = _rounding_case(sizes, q_len) + got = attend_segments_lse(q, cache, scale=scale) + assert got.dtype == mx.bfloat16 + assert _differing(got, exact) <= 0.005 + + +@pytest.mark.parametrize("q_len", [1, 6]) +def test_reference_lse_route_is_as_close_to_stock_sdpa_as_the_exact_result(cpu, monkeypatch, q_len) -> None: + monkeypatch.setattr(segmented_kv_module, "sdpa_lse_available", lambda: False) + cache, q, keys, values, scale, exact = _rounding_case([3000, 1000, 200], q_len) + got = attend_segments_lse(q, cache, scale=scale) + stock = mx.fast.scaled_dot_product_attention(q, keys, values, scale=scale, mask="causal" if q_len > 1 else None) + assert _differing(got, stock) <= _differing(exact, stock) + 0.005 + + +def test_merge_rounds_float32_segment_outputs_once_to_the_requested_dtype(cpu) -> None: + a = mx.random.normal((1, 2, 3, 4), key=mx.random.key(1)) + b = mx.random.normal((1, 2, 3, 4), key=mx.random.key(2)) + la, lb = mx.zeros((1, 2, 3)), mx.zeros((1, 2, 3)) + merged = merge_segment_outputs([a, b], [la, lb], dtype=mx.bfloat16) + assert merged.dtype == mx.bfloat16 + assert mx.array_equal(merged, ((a + b) / 2).astype(mx.bfloat16)).item() + assert merge_segment_outputs([a], [la], dtype=mx.bfloat16).dtype == mx.bfloat16 + + def test_merge_of_one_segment_is_the_segment_output() -> None: out = mx.ones((1, 2, 3, 4), dtype=mx.bfloat16) lse = mx.zeros((1, 2, 3)) From c695fe3f5ea454d2ae9cecf3b06555bb71180344 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Tue, 13 Oct 2026 14:17:02 +0200 Subject: [PATCH 35/37] Segmented KV on Macs without tensor units: per-segment attention with a logsumexp merge The segment kernel needs the Metal 4 tensor units of an M5-class GPU. On M1 to M4 the support check reported nax_unavailable and kept the stock cache. Add a decode route that attends per segment (exact float32 attention with its logsumexp, then one exact merge), selected automatically without tensor units or with MTPLX_SEGMENTED_KV_DECODE=lse. The memory gain is the same; the speed is not measured on those Macs yet. --- mtplx/segmented_kv.py | 80 +++++++++-- tests/test_sdpa_segmented_hd128.py | 1 + tests/test_segmented_kv_lse_route.py | 206 +++++++++++++++++++++++++++ tests/test_segmented_kv_support.py | 39 ++++- 4 files changed, 315 insertions(+), 11 deletions(-) create mode 100644 tests/test_segmented_kv_lse_route.py diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index b32b1fc2c..b1bcc3e38 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -22,7 +22,9 @@ 2026-10-07 phase 1 measurement) and a branch from the same history costs 0. Attention over the segments: ``mtplx/kernels/sdpa_segmented.py`` (exact split-softmax merge, -one launch for up to 12 segments) for the verify/decode window. A cache with a single segment +one launch for up to 12 segments) for the verify/decode window on GPUs with tensor units (M5 +class). On M1 to M4 (or with ``MTPLX_SEGMENTED_KV_DECODE=lse``) the window attends per segment +with its logsumexp and one exact merge instead (``decode_route``, ``attend_segments_lse``). A cache with a single segment (a cold prompt, or a conversation not yet sealed) runs the stock route on that segment's rows, a zero-copy view of the buffer the stock cache would have (counter ``single_q``). Only a call over several segments that no kernel takes (a masked or oversized window, a missing @@ -82,6 +84,23 @@ def segment_route_switched_off() -> list[str]: return reasons +def decode_route() -> str: + """How verify/decode attention reads several segments: ``kernel`` or ``lse``. + + ``kernel``: the fused segment kernel (``mtplx/kernels/sdpa_segmented.py``), which needs the + Metal 4 tensor units of an M5-class GPU. ``lse``: per segment attention with its logsumexp + plus one exact merge (``attend_segments_lse``), which runs on any Mac. Automatic: ``kernel`` + where the tensor units exist, ``lse`` elsewhere. ``MTPLX_SEGMENTED_KV_DECODE=lse`` selects + the ``lse`` route on any Mac (the way to test it on an M5); ``kernel`` cannot be forced onto + a GPU without tensor units. + """ + from .nax_verify import nax_available + + if os.environ.get("MTPLX_SEGMENTED_KV_DECODE", "").strip().lower() == "lse": + return "lse" + return "kernel" if nax_available() else "lse" + + #: Verdict on the loaded model (``evaluate_model_support``); None until a model was checked. _MODEL_SUPPORT: dict[str, Any] | None = None _UNSUPPORTED_LOGGED = False @@ -123,11 +142,11 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No stored; it switches the feature off (``segmented_kv_enabled``) and is reported on /health. """ from .kernels.sdpa_segmented import SUPPORTED_HEAD_DIMS, segments_supported - from .nax_verify import nax_available global _MODEL_SUPPORT, _LSE_NOTE_LOGGED layers = list(layers) flags = list(hooked) if hooked is not None else [True] * len(layers) + route = decode_route() reasons: list[str] = [] shapes: set[tuple[int, int, int]] = set() if not layers: @@ -145,17 +164,21 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No continue shapes.add((int(head_dim), int(q_heads), int(kv_heads))) for head_dim, q_heads, kv_heads in sorted(shapes): - if head_dim not in SUPPORTED_HEAD_DIMS: - reasons.append(f"head_dim_{head_dim}_unsupported") - elif q_heads % kv_heads or not segments_supported(1, q_heads // kv_heads, head_dim): + if q_heads % kv_heads: reasons.append(f"gqa_{q_heads}_{kv_heads}_unsupported") - reasons.extend(segment_route_switched_off()) - if not nax_available(): - reasons.append("nax_unavailable") + elif route == "kernel": + # Only the segment kernel has shape limits; the lse route takes any head dim and GQA. + if head_dim not in SUPPORTED_HEAD_DIMS: + reasons.append(f"head_dim_{head_dim}_unsupported") + elif not segments_supported(1, q_heads // kv_heads, head_dim): + reasons.append(f"gqa_{q_heads}_{kv_heads}_unsupported") + if route == "kernel": + reasons.extend(segment_route_switched_off()) reasons = list(dict.fromkeys(reasons)) verdict = { "supported": not reasons, "reasons": reasons, + "decode_route": route, "head_dims": sorted({s[0] for s in shapes}), "gqa": sorted({s[1] // s[2] for s in shapes if s[2] and s[1] % s[2] == 0}), "supported_head_dims": list(SUPPORTED_HEAD_DIMS), @@ -163,6 +186,14 @@ def evaluate_model_support(layers: Iterable[Any], *, hooked: Sequence[bool] | No } _MODEL_SUPPORT = verdict _log_unsupported(reasons) + if not reasons and route == "lse" and segmented_kv_requested(): + print( + "[mtplx] MTPLX_SEGMENTED_KV is on without the segment kernel (no tensor units, or " + "MTPLX_SEGMENTED_KV_DECODE=lse): attention over several segments runs one exact " + "attention call per segment plus a logsumexp merge. The memory gain is the same; the " + "speed on this Mac is not measured yet", + flush=True, + ) if not reasons and segmented_kv_requested() and prefill_route() == "gather" and not _LSE_NOTE_LOGGED: _LSE_NOTE_LOGGED = True print( @@ -197,6 +228,8 @@ def kernel_probe_shapes() -> list[tuple[int, int]]: if not segmented_kv_enabled(): return [] support = _MODEL_SUPPORT or {} + if support.get("decode_route") == "lse": + return [] # no segment kernel runs on this route, nothing to probe pairs = [(int(d), int(g)) for d in support.get("head_dims", []) for g in support.get("gqa", [])] return [(d, g) for d, g in pairs if segments_supported(1, g, d)] @@ -840,11 +873,14 @@ def decode_segments_attention( segments_supported, ) - if segment_route_switched_off() or lane_disabled("segmented_kv_sdpa"): + lse_route = decode_route() == "lse" + if not lse_route and (segment_route_switched_off() or lane_disabled("segmented_kv_sdpa")): return None q_len = int(queries.shape[2]) if not 1 <= q_len <= VERIFY_WINDOW_MAX: return None + if lse_route: + return _decode_segments_attention_lse(queries, cache, scale=scale, mask=mask) if cache.segment_count == 1 and cache.offset < int(packed_threshold): # One segment below the packed threshold: the stock route over a zero-copy view, # exactly what the stock cache does there (several segments always take the kernel, @@ -889,6 +925,31 @@ def decode_segments_attention( return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=2) +def _decode_segments_attention_lse( + queries: mx.array, cache: SegmentedKVCache, *, scale: float, mask: Any +) -> mx.array | None: + """The ``lse`` decode route (Macs without tensor units): per segment (out, lse) plus one merge. + + Same contract as the kernel route: a verify window of 1 to ``VERIFY_WINDOW_MAX`` rows whose + rows are the cache's last ones, a causal or no mask, any number of segments, so no sub-windows + and no ``MAX_FUSED_SEGMENTS`` limit. One segment returns None: the stock route on a zero-copy + view of its rows is the fastest and exact there. Each segment runs the MLX kernel with + logsumexp output when the build has it and the exact float32 reference otherwise + (``segment_sdpa_lse``); the merge is exact, so the result equals the dense attention up to + the rounding of the per-segment outputs to the query dtype. + """ + if cache.segment_count == 1: + return None + if not (mask is None or (isinstance(mask, str) and mask == "causal")): + return None + segments = cache.attention_segments() + if not segments or int(segments[-1][2]) < int(queries.shape[2]): + return None + out = attend_segments_lse(queries, cache, scale=scale) + count_route(f"lse_q{int(queries.shape[2])}") + return out + + # ---- install -------------------------------------------------------------------------------- @@ -1089,6 +1150,7 @@ def segmented_kv_health(entries: Iterable[Any]) -> dict[str, Any] | None: "ssd": segmented_ssd_enabled(), "prefill_route": prefill_route(), "lse_api": lse_api_found(), + "decode_route": decode_route(), "entries": per_entry, "max_segments": max((e["segments"] for e in per_entry), default=0), "sealed_bytes_unique": unique_state_bytes(all_states), diff --git a/tests/test_sdpa_segmented_hd128.py b/tests/test_sdpa_segmented_hd128.py index 2367397e6..5039bc87e 100644 --- a/tests/test_sdpa_segmented_hd128.py +++ b/tests/test_sdpa_segmented_hd128.py @@ -257,6 +257,7 @@ def test_sub_windows_respect_both_kernel_limits_at_every_gqa(monkeypatch, gqa, e mx.set_default_device(mx.cpu) try: monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "1") + monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: True) # this test is about the kernel route hk, d = 2, 128 cache = _low_gqa_cache(hk, d, window) rows: list[int] = [] diff --git a/tests/test_segmented_kv_lse_route.py b/tests/test_segmented_kv_lse_route.py new file mode 100644 index 000000000..327f380a3 --- /dev/null +++ b/tests/test_segmented_kv_lse_route.py @@ -0,0 +1,206 @@ +"""Segmented KV on Macs without tensor units: the ``lse`` decode route (per segment attention +plus an exact logsumexp merge, ``MTPLX_SEGMENTED_KV_DECODE=lse``). + +The route must equal the dense attention over the gathered rows (exactness), run through the real +attention layer like the stock cache (greedy parity: same argmax through a projection, same +routes, no gather) and leave the kernel route alone. Runs on any GPU, including an M5 that +rehearses the other Macs. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import mlx.core as mx +import pytest +from mlx.utils import tree_map +from mlx_lm.models.cache import KVCache +from mlx_lm.models.qwen3_next import Qwen3NextAttention + +import mtplx.segmented_kv as module +from mtplx.attention_split import configure_split_full_attention +from mtplx.segmented_kv import SegmentedKVCache, decode_segments_attention + + +@pytest.fixture +def cpu(): + """The exactness tests run on the CPU: MLX float32 GEMM on a tensor-unit GPU rounds like TF32.""" + previous = mx.default_device() + mx.set_default_device(mx.cpu) + yield + mx.set_default_device(previous) + + +@pytest.fixture(autouse=True) +def _lse_route(monkeypatch): + monkeypatch.setenv("MTPLX_SEGMENTED_KV_DECODE", "lse") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MERGE_MAX_ROWS", "0") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_MAX_SEGMENTS", "1000") + monkeypatch.setenv("MTPLX_SEGMENTED_KV_PREFILL", "lse") + module.route_counts.clear() + + +def _dense(q, k, v, scale, n_rows): + """Reference: float32 softmax attention, the new rows (last Q) causal over the whole history.""" + b, hq, q_len, d = q.shape + hk = k.shape[1] + k = mx.repeat(k.astype(mx.float32), hq // hk, axis=1) + v = mx.repeat(v.astype(mx.float32), hq // hk, axis=1) + scores = (q.astype(mx.float32) @ k.transpose(0, 1, 3, 2)) * scale + visible = mx.arange(n_rows)[None, :] <= (n_rows - q_len + mx.arange(q_len))[:, None] + scores = mx.where(visible[None, None], scores, -mx.inf) + return (mx.softmax(scores, axis=-1) @ v).astype(q.dtype) + + +def _cache(hk, d, sizes, dtype, seed=0): + """A segmented cache with sealed segments of ``sizes[:-1]`` rows and a tail of ``sizes[-1]``.""" + mx.random.seed(seed) + cache = SegmentedKVCache() + keys, values = [], [] + for index, n in enumerate(sizes): + k = (mx.random.normal((1, hk, n, d)) * 1.5).astype(dtype) + v = mx.random.normal((1, hk, n, d)).astype(dtype) + cache.append_rows(k, v) + keys.append(k) + values.append(v) + if index < len(sizes) - 1: + cache.seal() + return cache, mx.concatenate(keys, axis=2), mx.concatenate(values, axis=2) + + +@pytest.mark.parametrize("dtype, tol", [(mx.float32, 2e-5), (mx.bfloat16, 2e-2)]) +@pytest.mark.parametrize("hq, hk, d", [(24, 4, 256), (8, 2, 128), (4, 4, 64)]) +@pytest.mark.parametrize("q_len", [1, 4, 9, 32]) +@pytest.mark.parametrize("sizes", [[300, 40], [1000, 513, 70, 33]]) +def test_the_lse_route_equals_the_dense_attention(cpu, dtype, tol, hq, hk, d, q_len, sizes) -> None: + cache, k, v = _cache(hk, d, sizes, dtype) + q = mx.random.normal((1, hq, q_len, d)).astype(dtype) + scale = 1.0 / math.sqrt(d) + out = decode_segments_attention(q, cache, scale=scale, mask="causal", packed_threshold=8192) + assert out is not None and out.shape == q.shape and out.dtype == dtype + expected = _dense(q, k, v, scale, sum(sizes)) + assert float(mx.abs(out.astype(mx.float32) - expected.astype(mx.float32)).max().item()) <= tol + assert module.route_counts.get(f"lse_q{q_len}") == 1 + + +def test_rows_without_a_visible_key_in_a_segment_merge_without_nan(cpu) -> None: + """The merge ignores segments whose logsumexp is -inf (the shared -inf-safe merge): a tail shorter + than the window masks its first rows entirely, and those rows take the earlier segments only.""" + cache, k, v = _cache(2, 64, [200, 2], mx.float32) + q = mx.random.normal((1, 4, 4, 64)) + out = module.attend_segments_lse(q, cache, scale=0.125) + assert bool(mx.all(mx.isfinite(out)).item()) + expected = mx.softmax((q[:, :, :2] @ mx.repeat(k[:, :, :200], 2, axis=1).transpose(0, 1, 3, 2)) * 0.125, axis=-1) @ mx.repeat(v[:, :, :200], 2, axis=1) + assert float(mx.abs(out[:, :, :2] - expected).max().item()) <= 2e-5 + + +def test_a_row_that_is_masked_everywhere_merges_to_zeros(cpu) -> None: + out = module.merge_segment_outputs( + [mx.ones((1, 1, 2, 4)), mx.ones((1, 1, 2, 4))], [mx.full((1, 1, 2), -mx.inf), mx.full((1, 1, 2), -mx.inf)] + ) + assert bool(mx.all(out == 0).item()) + + +def test_a_single_segment_runs_the_stock_route(cpu) -> None: + cache, _, _ = _cache(2, 64, [500], mx.float32) + q = mx.random.normal((1, 4, 4, 64)) + assert decode_segments_attention(q, cache, scale=0.125, mask="causal", packed_threshold=8192) is None + + +def test_a_mask_or_an_oversized_window_goes_to_the_stock_ladder(cpu) -> None: + cache, _, _ = _cache(2, 64, [500, 40], mx.float32) + q = mx.random.normal((1, 4, 4, 64)) + mask = mx.ones((4, 540), dtype=mx.bool_) + assert decode_segments_attention(q, cache, scale=0.125, mask=mask, packed_threshold=8192) is None + wide = mx.random.normal((1, 4, 33, 64)) + assert decode_segments_attention(wide, cache, scale=0.125, mask="causal", packed_threshold=8192) is None + + +def test_more_than_the_fused_kernel_limit_of_segments_stays_on_the_route(cpu) -> None: + sizes = [64] * 30 + [40] # the kernel takes 12 per launch and 20 launches; here nothing limits it + cache, k, v = _cache(2, 64, sizes, mx.float32) + q = mx.random.normal((1, 4, 4, 64)) + out = decode_segments_attention(q, cache, scale=0.125, mask="causal", packed_threshold=8192) + assert float(mx.abs(out - _dense(q, k, v, 0.125, sum(sizes))).max().item()) <= 2e-5 + + +def test_the_kernel_is_never_touched_on_the_lse_route(cpu, monkeypatch) -> None: + def boom(**_kwargs): + raise AssertionError("segment kernel called") + + monkeypatch.setattr("mtplx.kernels.sdpa_segmented.sdpa_nax_flash_dsplit_segments", boom) + cache, _, _ = _cache(2, 64, [300, 40], mx.float32) + q = mx.random.normal((1, 4, 4, 64)) + assert decode_segments_attention(q, cache, scale=0.125, mask=None, packed_threshold=8192) is not None + + +# ---- through the real attention layer: greedy parity with the stock cache --------------------- + +HIDDEN = 128 +VOCAB = 512 + + +def _model(): + args = SimpleNamespace( + hidden_size=HIDDEN, num_attention_heads=24, num_key_value_heads=4, head_dim=256, + attention_bias=False, rms_norm_eps=1e-6, partial_rotary_factor=0.25, + rope_theta=10000.0, rope_scaling=None, max_position_embeddings=65536, + ) + mx.random.seed(0) + attn = Qwen3NextAttention(args) + attn.update(tree_map(lambda p: p.astype(mx.bfloat16), attn.parameters())) + mx.eval(attn.parameters()) + layer = SimpleNamespace(is_linear=False, self_attn=attn) + return SimpleNamespace(model=SimpleNamespace(layers=[layer])), attn + + +def _x(seed: int, n: int) -> mx.array: + return (mx.random.normal((1, n, HIDDEN), key=mx.random.key(seed)) * 0.5).astype(mx.bfloat16) + + +def _run_turns(attn, cache): + outs = [] + for seed, n in [(1, 2304), (2, 2304)]: + outs.append(attn(_x(seed, n), mask="causal", cache=cache)) + if isinstance(cache, SegmentedKVCache): + state = cache.state + cache = SegmentedKVCache() + cache.state = state + outs.append(attn(_x(3, 300), mask="causal", cache=cache)) + for seed in (4, 5, 6): + outs.append(attn(_x(seed, 4), mask="causal", cache=cache)) + cache.trim(2) + outs.append(attn(_x(7, 9), mask="causal", cache=cache)) + cache.trim(5) + outs.append(attn(_x(8, 1), mask=None, cache=cache)) + mx.eval(outs) + return outs + + +@pytest.mark.skipif(not mx.metal.is_available(), reason="needs the Metal GPU") +def test_the_attention_layer_agrees_with_the_stock_cache_and_never_gathers() -> None: + model, attn = _model() + configure_split_full_attention(model) + proj = mx.random.normal((attn.o_proj.weight.shape[0], VOCAB), key=mx.random.key(9)).astype(mx.bfloat16) + stock = _run_turns(attn, KVCache()) + module.route_counts.clear() + seg = _run_turns(attn, SegmentedKVCache()) + assert not any(k.startswith("gather") for k in module.route_counts), module.route_counts + assert module.route_counts.get("lse_q4", 0) >= 3 and module.route_counts.get("lse_q9", 0) >= 1 + decided_rows = total_rows = 0 + for a, b in zip(stock, seg): + assert bool(mx.all(mx.isfinite(b)).item()) + # greedy parity: wherever the top-1 margin exceeds 4x the logit difference the cache can + # explain, the argmax token through a fixed projection must be the same + la, lb = a.astype(mx.float32) @ proj.astype(mx.float32), b.astype(mx.float32) @ proj.astype(mx.float32) + top2 = mx.sort(la, axis=-1)[..., -2:] + decided = (top2[..., 1] - top2[..., 0]) > 4 * mx.abs(la - lb).max(axis=-1) + same = mx.argmax(la, axis=-1) == mx.argmax(lb, axis=-1) + assert bool(mx.all(same | ~decided).item()) + decided_rows += int(decided.sum().item()) + total_rows += int(decided.size) + top = float(mx.maximum(mx.abs(a).max(), mx.abs(b).max()).item()) + err = float(mx.abs(a.astype(mx.float32) - b.astype(mx.float32)).max().item()) + assert err / 2.0 ** (math.floor(math.log2(top)) - 7) <= 8 + assert decided_rows > 0.5 * total_rows # the check is not vacuous diff --git a/tests/test_segmented_kv_support.py b/tests/test_segmented_kv_support.py index a6e7e404f..909a54fce 100644 --- a/tests/test_segmented_kv_support.py +++ b/tests/test_segmented_kv_support.py @@ -94,9 +94,44 @@ def test_an_attention_switch_set_to_0_is_a_kill_switch(monkeypatch, name, reason assert segmented_kv_enabled() is False -def test_missing_nax_is_a_reason(monkeypatch) -> None: +def test_missing_nax_selects_the_lse_route_instead_of_switching_segments_off(monkeypatch) -> None: monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: False) - assert "nax_unavailable" in evaluate_model_support([_attn()])["reasons"] + verdict = evaluate_model_support([_attn()]) + assert verdict["supported"] and verdict["reasons"] == [] and verdict["decode_route"] == "lse" + assert segmented_kv_enabled() is True + assert module.kernel_probe_shapes() == [] # no segment kernel runs, nothing to probe + + +def test_the_lse_route_takes_shapes_the_segment_kernel_does_not(monkeypatch) -> None: + monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: False) + assert evaluate_model_support([_attn(head_dim=64, q=16, kv=4), _attn(head_dim=80, q=40, kv=1)])["supported"] + assert not evaluate_model_support([_attn(head_dim=64, q=15, kv=4)])["supported"] # GQA must divide + + +def test_the_attention_kill_switches_only_apply_to_the_kernel_route(monkeypatch) -> None: + monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: False) + monkeypatch.setenv("MTPLX_NAX_FLASH_ROUTE", "0") + assert evaluate_model_support([_attn()])["supported"] + + +def test_the_lse_route_can_be_selected_on_a_tensor_unit_gpu(monkeypatch) -> None: + assert module.decode_route() == "kernel" + monkeypatch.setenv("MTPLX_SEGMENTED_KV_DECODE", "lse") + assert module.decode_route() == "lse" + assert evaluate_model_support([_attn(head_dim=64, q=16, kv=4)])["supported"] + + +def test_the_kernel_route_cannot_be_forced_onto_a_gpu_without_tensor_units(monkeypatch) -> None: + monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: False) + monkeypatch.setenv("MTPLX_SEGMENTED_KV_DECODE", "kernel") + assert module.decode_route() == "lse" + + +def test_the_rehearsal_switch_selects_the_lse_route(monkeypatch) -> None: + """MTPLX_FORCE_GPU_FAMILY_FALLBACK=1 (an M5 acting as an M1 to M4) reaches the route choice.""" + monkeypatch.undo() # drop the nax_available patch of the autouse fixture + monkeypatch.setenv("MTPLX_FORCE_GPU_FAMILY_FALLBACK", "1") + assert module.decode_route() == "lse" def test_an_unsupported_model_keeps_stock_caches_and_never_builds_a_segmented_one() -> None: From 33d9df748defa2e7365dde2998e708d7c4f7c2ac Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Tue, 13 Oct 2026 15:53:54 +0200 Subject: [PATCH 36/37] Changelog: segmented KV on Macs without tensor units --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index f4a251dea..833fc00ee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ All notable user-facing changes to MTPLX. The format is based on - 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 From 4fb7480e5574793fac51f74f8a9128fefcc7fa7e Mon Sep 17 00:00:00 2001 From: Jeroen van Menen Date: Tue, 13 Oct 2026 16:47:02 +0200 Subject: [PATCH 37/37] Segmented KV lse decode route: round the merged output once, to the query dtype The reference route and the merge now keep per-segment outputs in float32 and round once (segmented KV lse prefill route); the decode route on Macs without tensor units picks that up. Document it in the decode route's docstring, and add the equivalent exactness tests for decode and verify windows: the result equals the float32 attention rounded once and is no further from stock SDPA than that exact result is. --- mtplx/segmented_kv.py | 5 ++-- tests/test_segmented_kv_lse_route.py | 38 ++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/mtplx/segmented_kv.py b/mtplx/segmented_kv.py index b1bcc3e38..2c477239f 100644 --- a/mtplx/segmented_kv.py +++ b/mtplx/segmented_kv.py @@ -935,8 +935,9 @@ def _decode_segments_attention_lse( and no ``MAX_FUSED_SEGMENTS`` limit. One segment returns None: the stock route on a zero-copy view of its rows is the fastest and exact there. Each segment runs the MLX kernel with logsumexp output when the build has it and the exact float32 reference otherwise - (``segment_sdpa_lse``); the merge is exact, so the result equals the dense attention up to - the rounding of the per-segment outputs to the query dtype. + (``segment_sdpa_lse``). The reference keeps every segment's output in float32 and the merge + rounds once, so the result is the dense float32 attention rounded to the query dtype, as + close to the stock kernel as that one is to the exact result. """ if cache.segment_count == 1: return None diff --git a/tests/test_segmented_kv_lse_route.py b/tests/test_segmented_kv_lse_route.py index 327f380a3..1c9046ceb 100644 --- a/tests/test_segmented_kv_lse_route.py +++ b/tests/test_segmented_kv_lse_route.py @@ -84,6 +84,44 @@ def test_the_lse_route_equals_the_dense_attention(cpu, dtype, tol, hq, hk, d, q_ assert module.route_counts.get(f"lse_q{q_len}") == 1 +def _differing(a, b) -> float: + """Fraction of elements that differ between two bfloat16 arrays.""" + return float(mx.mean(a.astype(mx.float32) != b.astype(mx.float32)).item()) + + +@pytest.mark.parametrize("q_len", [1, 4, 9]) +@pytest.mark.parametrize("sizes", [[3000, 200], [2000, 1024, 600, 40]]) +def test_the_lse_route_rounds_once_like_the_stock_attention(cpu, q_len, sizes) -> None: + """bfloat16 is what the model runs: the route must be the float32 attention rounded once. + + A tolerance of one ulp does not see a second rounding (each segment rounded to bfloat16 before + the merge), which puts a quarter or more of the elements one ulp off; on real Qwen3.8-27B + tensors it showed as about one element in three differing from the stock kernel, against one + in six for the exact result. + """ + cache, k, v = _cache(4, 256, sizes, mx.bfloat16) + q = mx.random.normal((1, 24, q_len, 256)).astype(mx.bfloat16) + scale = 1.0 / math.sqrt(256) + out = decode_segments_attention(q, cache, scale=scale, mask="causal", packed_threshold=8192) + exact = _dense(q, k, v, scale, sum(sizes)) + assert out.dtype == mx.bfloat16 + assert _differing(out, exact) <= 0.005 + + +@pytest.mark.parametrize("q_len", [1, 4]) +def test_the_lse_route_is_as_close_to_stock_sdpa_as_the_exact_result(cpu, q_len) -> None: + """Against the stock fused SDPA on the gathered rows, the route may differ only where the + exact result itself differs from it (the stock kernel's own rounding).""" + sizes = [3000, 1000, 200] + cache, k, v = _cache(4, 256, sizes, mx.bfloat16) + q = mx.random.normal((1, 24, q_len, 256)).astype(mx.bfloat16) + scale = 1.0 / math.sqrt(256) + out = decode_segments_attention(q, cache, scale=scale, mask="causal", packed_threshold=8192) + stock = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask="causal" if q_len > 1 else None) + exact = _dense(q, k, v, scale, sum(sizes)) + assert _differing(out, stock) <= _differing(exact, stock) + 0.005 + + def test_rows_without_a_visible_key_in_a_segment_merge_without_nan(cpu) -> None: """The merge ignores segments whose logsumexp is -inf (the shared -inf-safe merge): a tail shorter than the window masks its first rows entirely, and those rows take the earlier segments only."""