diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5039769..833fc00ee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,23 @@ All notable user-facing changes to MTPLX. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow [Semantic Versioning](https://semver.org/). +## [Unreleased] + +### Added + +- **Segmented KV cache, behind `MTPLX_SEGMENTED_KV=1` (default off).** The full-attention layers keep a conversation as sealed, exact-size KV segments plus one growing tail instead of one buffer that the session bank's snapshot aliases. A follow-up turn writes only into a new tail, so MLX's copy-on-write no longer copies the history: on Qwen3.8 27B a stock snapshot costs +4.95 GiB at 80K and +3.1 GiB at 50K at the first write of the next turn, with segments +0.125 GiB. + - Verify windows of 1 to 32 rows attend straight over the segments with one fused split-softmax kernel launch for up to 12 segments (`mtplx/kernels/sdpa_segmented.py`); a window the kernel does not take whole runs as sub-windows over the segments. A cold prompt is one segment and runs the stock route on a zero-copy view of its rows. Prefill over existing segments uses the fused SDPA with logsumexp when the MLX build has it and gathers the history otherwise (one log line says so). + - With the switch on the full-attention layers are segment caches from the first token, for models the segment kernel serves (head_dim 128 or 256, the split-attention hook installed, a GPU with NAX TensorOps), under any profile: segments run their own kernel and do not need `MTPLX_GQA_PACKED_SDPA` or `MTPLX_NAX_FLASH_ROUTE`, which only switch them off when set to 0. The kernel self-check runs whenever the switch is on and tests the segment kernel at the model's head_dim before the first cache is built; a failure keeps the stock cache. On Qwen3.6-35B-A3B under the `sustained` profile (50K plus 4 turns, neither switch set) the peak above the turn start drops from 2.7 to 3.3 GiB to 0.45 to 0.76 GiB and ms per verify from 34 to 41 to 25 to 30, with no history gather. The hook serves the gated Qwen3.5/3.6/3.8 attention, Qwen3's q/k-norm attention and the Llama-style attention without norms (Llama, Mistral); on Mistral 7B v0.3 (AR path, 22K plus 4 turns) the peak above the turn start drops from 3.2 to 7.0 GiB to 0.7 to 1.3 GiB at the same ms per step, with greedy output identical. Any other model, and a model loaded without that check, keeps the stock cache; `/health` says why. + - The session bank stores the list of segments as the snapshot and counts a shared segment once; the admission gate prices no history copy for it. Small recent segments merge in geometric tiers, also when bank snapshots still hold them, which keeps a conversation at 3 to 5 segments over 12 follow-up turns at 50K. + - A conversation stays on segments past the dense-decode ceiling (157,286 tokens on a 64 GB Mac with the auto ceiling). The repage to the paged cache built its pages from a gather of every segment: on Qwen3.8 27B a follow-up turn of 161,835 tokens with 84 new ones was priced 11.6 GiB for that copy and refused with 10.1 GiB free; with segments it is priced for the new rows only, about 1.1 GiB. Quantized KV and an explicit `MTPLX_SUSTAINED_PREFILL_LAYOUT` keep the repage. + - Segmented entries persist to the SSD tier in the existing block format, without gathering the history (`MTPLX_SEGMENTED_KV_SSD=0` skips them): shared prefixes dedupe, entries restore with the switch off and stock entries restore with it on. A follow-up spill hashes only the new segment. + - Measured on Qwen3.8 27B against the stock cache with compiled verify: 3 to 12 % fewer ms per verify at 8K to 32K, 0 to 4 % at 50K, and 1 to 4.5 GiB less peak above the turn start at every context. Not covered: the compiled verify (segmented caches verify eager; `MTPLX_VERIFY_ASYNC_CHUNK_LAYERS=8` still helps), paged and quantized KV. `/health` gets a `segmented_kv` block (segments per entry, merges, route and SSD counters), absent when the switch is off. + - **Macs without the tensor units (M1 to M4) can use segments too.** The segment kernel needs the Metal 4 tensor units of an M5-class GPU, and on M1 to M4 the support check kept the stock cache (`nax_unavailable`). Without them (or with `MTPLX_SEGMENTED_KV_DECODE=lse`, or `MTPLX_FORCE_GPU_FAMILY_FALLBACK=1` to try it on an M5) decode attends per segment in float32, each with its logsumexp, and merges the pieces exactly; prefill and the single-segment case are unchanged, and the M5 default stays the kernel. The memory gain is the same; the speed is not measured on those Macs yet. `/health` shows `decode_route` and `lse_api`. On an M5 with the fallback forced (Qwen3.8 27B, 20K and 50K plus 4 turns), the peak above the turn start drops 2.1 to 5.1 GiB per follow-up turn and decode runs at 0.58 to 0.82 of the stock cache's tok/s there. + +### Fixed + +- **A conversation restored from the SSD tier gets the cache layout the request asked for**, the same one a restore from memory gets. The SSD restore built the runtime's default layout, which for a long context is the paged KV cache (`VllmMetalPagedKVCache`) instead of the dense cache: about 5x slower prefill and slower decode after every restore from disk. Measured on Qwen3.8 27B (with the segmented KV cache of a later change, which had the same omission), the first token after a restore from SSD at 50K and 80K came after 5 to 8 s instead of 16 to 48 s and decode went from 14 to 18 to 25 to 30 tok/s; a run without segments showed the same 14 to 17 tok/s after its SSD restore. + ## [2.12.2] - 2026-10-03 ### Fixed diff --git a/mtplx/attention_split.py b/mtplx/attention_split.py index 289e1d2ca..33b9a10da 100644 --- a/mtplx/attention_split.py +++ b/mtplx/attention_split.py @@ -19,6 +19,16 @@ note_unowned_rotary_origin, rope_offset_of, ) +from .segmented_kv import ( + VERIFY_WINDOW_MAX, + SegmentedKVCache, + attend_segments_lse, + decode_segments_attention, + gathered_view, + segmented_kv_requested, +) +from .segmented_kv import count_route as segmented_count +from .segmented_kv import prefill_route as segmented_prefill_route def _env_enabled(name: str, *, default: bool = False) -> bool: @@ -336,6 +346,39 @@ def _attention_has_gated_q_proj(attn: Any) -> bool: return int(weight.shape[0]) == expected +def _attention_has_plain_q_proj(attn: Any) -> bool: + """Plain q/k-norm attention (Qwen3): q_proj has one output per head dimension, no gate.""" + q_proj = getattr(attn, "q_proj", None) + q_norm = getattr(attn, "q_norm", None) + weight = getattr(q_proj, "weight", None) + norm_weight = getattr(q_norm, "weight", None) + if weight is None or norm_weight is None: + return False + num_heads = int(getattr(attn, "num_attention_heads", getattr(attn, "n_heads", 0))) + return int(weight.shape[0]) == num_heads * int(norm_weight.shape[0]) + + +def _attention_has_unnormed_q_proj(attn: Any) -> bool: + """Llama-style attention (Llama, Mistral): q/k/v projections and rope, no q/k norm, no gate.""" + if getattr(attn, "q_norm", None) is not None or getattr(attn, "k_norm", None) is not None: + return False + if any(getattr(attn, name, None) is None for name in ("k_proj", "v_proj", "o_proj", "rope")): + return False + weight = getattr(getattr(attn, "q_proj", None), "weight", None) + head_dim = getattr(attn, "head_dim", None) + if not head_dim and getattr(attn, "scale", None): + head_dim = round(float(attn.scale) ** -2) + num_heads = int(getattr(attn, "num_attention_heads", getattr(attn, "n_heads", 0))) + if weight is None or not head_dim or not num_heads: + return False + return int(weight.shape[0]) == num_heads * int(head_dim) + + +def _attention_has_ungated_q_proj(attn: Any) -> bool: + """Ungated attention the hook body serves on a segmented cache (Qwen3, Llama, Mistral).""" + return _attention_has_plain_q_proj(attn) or _attention_has_unnormed_q_proj(attn) + + def _install_split_attention_hook(attn: Any) -> bool: cls = type(attn) if getattr(cls, "_mtplx_split_full_attention_installed", False): @@ -350,7 +393,11 @@ def split_call( cache: Any | None = None, ) -> mx.array: route_off = not getattr(self, "_mtplx_split_full_attention_enabled", False) - if route_off and not cache_owns_rotary_origin(cache): + if ( + route_off + and not cache_owns_rotary_origin(cache) + and not isinstance(cache, SegmentedKVCache) + ): # No MTPLX attention route asked for, and the cache has no rotary # origin of its own: the stock forward, untouched. A cache that # owns one (an image request on a compiled route) rotates at it @@ -371,27 +418,44 @@ def split_call( # standalone sigmoid eagerly and a fused one inside the compiled # verifier's trace, which differ in the last bit of some bfloat16 # and float32 values (mtplx/attention_math.py). - if not _attention_has_gated_q_proj(self): + # A segmented cache always takes the body (the segment route + # below), under any profile: the stock forward would gather the + # history on every call. + gated_q_proj = _attention_has_gated_q_proj(self) + if not gated_q_proj and not ( + isinstance(cache, SegmentedKVCache) and _attention_has_ungated_q_proj(self) + ): + # Only the gated (Qwen3-Next style) attention goes through this body, and, for + # MTPLX_SEGMENTED_KV, an ungated attention (Qwen3 with q/k norm, Llama and + # Mistral without) on a segmented cache: the stock forward would gather the + # history on every call. return original_call(self, x, mask=mask, cache=cache) from mlx_lm.models.base import scaled_dot_product_attention + n_heads = getattr(self, "num_attention_heads", None) or self.n_heads + n_kv_heads = getattr(self, "num_key_value_heads", None) or self.n_kv_heads B, L, _ = x.shape q_proj_output = self.q_proj(x) - queries, gate = mx.split( - q_proj_output.reshape(B, L, self.num_attention_heads, -1), - 2, - axis=-1, - ) - gate = gate.reshape(B, L, -1) + if gated_q_proj: + queries, gate = mx.split( + q_proj_output.reshape(B, L, n_heads, -1), + 2, + axis=-1, + ) + gate = gate.reshape(B, L, -1) + else: + queries, gate = q_proj_output.reshape(B, L, n_heads, -1), None keys = self.k_proj(x) values = self.v_proj(x) - queries = self.q_norm(queries).transpose(0, 2, 1, 3) - keys = self.k_norm( - keys.reshape(B, L, self.num_key_value_heads, -1) - ).transpose(0, 2, 1, 3) - values = values.reshape(B, L, self.num_key_value_heads, -1).transpose( + keys = keys.reshape(B, L, n_kv_heads, -1) + if getattr(self, "q_norm", None) is not None: + queries = self.q_norm(queries) + keys = self.k_norm(keys) + queries = queries.transpose(0, 2, 1, 3) + keys = keys.transpose(0, 2, 1, 3) + values = values.reshape(B, L, n_kv_heads, -1).transpose( 0, 2, 1, @@ -425,11 +489,51 @@ def split_call( and int(B) == 1 and can_slice_mask ) + segmented_output = None + segmented_route = "" if cache is not None: rope_offset = _cache_rope_offset(cache) queries = self.rope(queries, offset=rope_offset) keys = self.rope(keys, offset=rope_offset) - if blockwise_enabled or vllm_metal_paged_enabled: + if isinstance(cache, SegmentedKVCache): + # Sealed segments plus a tail (MTPLX_SEGMENTED_KV): the new rows go to + # the tail. A prefill chunk over several segments takes the logsumexp + # route and a verify window attends straight over the segments; any + # other call runs the ladder below on one view of the rows (zero-copy + # for a single segment, gathered otherwise). + cache.append_rows(keys, values) + if ( + int(queries.shape[2]) > VERIFY_WINDOW_MAX + and cache.segment_count > 1 + and segmented_prefill_route() == "lse" + and (mask is None or (isinstance(mask, str) and mask == "causal")) + ): + segmented_output = attend_segments_lse( + queries, cache, scale=self.scale + ) + segmented_route = "segments_lse" + segmented_count("prefill_lse") + else: + segmented_output = decode_segments_attention( + queries, + cache, + scale=self.scale, + mask=mask, + packed_threshold=int( + getattr(self, "_mtplx_gqa_packed_sdpa_threshold", 8192) + ), + ) + segmented_route = "nax_flash_segments" + if segmented_output is None: + q_rows = int(queries.shape[2]) + segmented_count( + f"{'single' if cache.segment_count == 1 else 'gather'}_q" + f"{q_rows if q_rows <= VERIFY_WINDOW_MAX else 'prefill'}" + ) + view = gathered_view(cache) + keys, values = view.keys, view.values + cache = view + elif blockwise_enabled or vllm_metal_paged_enabled: cache.update_without_fetch(keys, values) else: keys, values = cache.update_and_fetch(keys, values) @@ -559,7 +663,10 @@ def split_call( ) route = "sdpa" fallback = "" - if should_use_vllm_metal_paged: + if segmented_output is not None: + output = segmented_output + route = segmented_route + elif should_use_vllm_metal_paged: impl_override = ( "fast_sdpa_gather" if getattr(self, "_mtplx_vllm_metal_exact_gather_layer", False) @@ -799,6 +906,8 @@ def split_call( ) _note_kv_attention(self, cache, queries, output, mask, route, fallback) output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + if gate is None: + return self.o_proj(output) return self.o_proj(attention_gate(output, gate)) cls.__call__ = split_call @@ -883,6 +992,7 @@ def configure_split_full_attention( } full_layers = list(_full_attention_layers(model)) full_layer_count = len(full_layers) + hooked_layers: list[bool] = [] for full_idx, attn in enumerate(full_layers): exact_gather_layer = bool( vllm_metal_paged @@ -894,7 +1004,13 @@ def configure_split_full_attention( ) ) ) - stats["installed"] += int(_install_split_attention_hook(attn)) + newly_hooked = bool(_install_split_attention_hook(attn)) + # The hook is installed once per attention class: later instances report False. + hooked_layers.append( + (newly_hooked or bool(getattr(type(attn), "_mtplx_split_full_attention_installed", False))) + and (_attention_has_gated_q_proj(attn) or _attention_has_ungated_q_proj(attn)) + ) + stats["installed"] += int(newly_hooked) attn._mtplx_split_full_attention_enabled = bool( active or sdpa_2pass or vllm_metal_paged or gqa_packed ) @@ -916,4 +1032,11 @@ def configure_split_full_attention( attn._mtplx_split_full_attention_calls = 0 stats["layers"] += 1 stats["exact_gather_layers"] += int(exact_gather_layer) + if segmented_kv_requested(): + # MTPLX_SEGMENTED_KV: decide once, from the hooked layers, whether this model can use + # segments without ever falling back to a gather (mtplx/segmented_kv.py). + from .segmented_kv import evaluate_model_support + + verdict = evaluate_model_support(full_layers, hooked=hooked_layers) + stats["segmented_kv_supported"] = bool(verdict["supported"]) return stats diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index 52503e5bc..06337dd6d 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,22 @@ 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 _count_ssd(route: str) -> None: + from mtplx.segmented_kv import count_route + + count_route(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") + return module is not None and isinstance(value, module.SegmentedKVState) + + class TreeCodec: """Flatten JSON-safe trees plus MLX arrays into raw tensor blobs.""" @@ -171,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 @@ -249,6 +272,23 @@ 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. 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": [ + 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 +379,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( @@ -360,22 +413,35 @@ 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 + else: + hit = None + seg_rows = value if hasattr(value, "known_digest") and self.known_blob is not None else 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) + hit = self.known_blob(known) if known 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 + _count_ssd("ssd_blocks_reused") + 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: + _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/cache_state.py b/mtplx/cache_state.py index 6e7a2546a..b6ca84b1c 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 { @@ -5015,7 +5021,14 @@ def restore_cache( restore_meta_state: bool = True, clone_states: bool = True, ) -> None: + 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/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/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/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/kernels/sdpa_segmented.py b/mtplx/kernels/sdpa_segmented.py new file mode 100644 index 000000000..ceefc5a3a --- /dev/null +++ b/mtplx/kernels/sdpa_segmented.py @@ -0,0 +1,352 @@ +"""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. + +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 + +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 + +#: 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 +#: 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 + +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, 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 + 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 ( + ("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]"), + ): + 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, 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 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") + 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 not in SUPPORTED_HEAD_DIMS: + 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 > MAX_M_ROWS: + return _bail("m_rows_gt_32") + 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") + 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 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 + ) + + 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), + ("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 in SUPPORTED_HEAD_DIMS and 1 <= q_len <= MAX_Q_LEN and gqa_factor * q_len <= MAX_M_ROWS + 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 new file mode 100644 index 000000000..2c477239f --- /dev/null +++ b/mtplx/segmented_kv.py @@ -0,0 +1,1178 @@ +"""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 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 +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``. + +GatedDeltaNet (recurrent) layers keep their state objects as they are. + +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 +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 _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 _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 + + +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 +_LSE_NOTE_LOGGED = False + + +def segmented_kv_enabled() -> bool: + """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 not None and bool(_MODEL_SUPPORT.get("supported", False)) + + +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, _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]: + """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 + + 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: + 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) + 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): + 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 q_heads % kv_heads: + reasons.append(f"gqa_{q_heads}_{kv_heads}_unsupported") + 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), + "full_attention_layers": len(layers), + } + _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( + "[mtplx] MTPLX_SEGMENTED_KV is on, but this MLX build has no logsumexp output on the fused " + "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, + ) + 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 {} + 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)] + + +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: + 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(route: str) -> None: + """One more call served by ``route`` (``route_counts``, reported on /health).""" + route_counts[route] = route_counts.get(route, 0) + 1 + + +_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. + + 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 = sdpa_with_lse(q, q, q, scale=1.0) + mx.eval(result) + ok = isinstance(result, (tuple, list)) and len(result) == 2 + except Exception: # noqa: BLE001 + 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: + """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__ = ("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) + # 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: + return int(self.keys.shape[2]) + + @property + def nbytes(self) -> int: + return int(self.keys.nbytes) + int(self.values.nbytes) + + +@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. 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) + + @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 _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 _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: + seen.setdefault(ref.segment.id, ref.segment.nbytes) + return sum(seen.values()) + + +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: + 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 + + # ---- 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 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 = _unique_segment_bytes(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).""" + 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: + 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() + 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 + count_route("seal") + if rows == tail.capacity: + keys, values = tail.keys, tail.values + mx.eval(keys, values) + else: + count_route("seal_copy") + keys = _own_copy(tail.keys[..., :rows, :]) + values = _own_copy(tail.values[..., :rows, :]) + self._tail = None + self._tail_n = 0 + 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. + + 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: + 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. + # 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: + raise ValueError(f"unsupported segmented KV state: {type(value).__name__}") + self._sealed = refs + + # ---- merge policy --------------------------------------------------------------------- + + 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 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. + """ + 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 = _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] + small = int(older.n) + int(newer.n) <= cap_rows + geometric = int(older.n) <= tier * int(newer.n) + over = len(self._sealed) > limit and small + if not (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) + self._sealed[-2:] = [SegRef(KVSegment(keys, values, sealed=True), int(older.n) + int(newer.n))] + count_route("merge") + + +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 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 + 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) + # 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) + 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 (``sdpa_with_lse``).""" + k, v = _rows_view(keys, n), _rows_view(values, n) + 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]) + + +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(): + 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_route("sdpa_lse_kernel_unsupported") + else: + count_route("sdpa_lse_kernel") + return out + count_route("sdpa_lse_reference") + 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], *, 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] 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 + for out, lse in zip(outs, lses): + # 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(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, dtype=queries.dtype) + + +# ---- 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, + *, + scale: float, + mask: Any, + packed_threshold: int, +) -> mx.array | None: + """Verify/decode attention straight over the segments, or None for the stock ladder. + + 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 + 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, + ) + + 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, + # the alternative would be a gather). + 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: + 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_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, + # 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] + if tail_n < q_len: + return None + 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_route(f"chunked_q{q_len}") + 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 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 + 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 -------------------------------------------------------------------------------- + + +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(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 (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``). + """ + 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_route("restored_stock_as_segment") + return state + if isinstance(state, SegmentedKVState): + return state.to_arrays() + return state + + +def install_segmented_attention_kv_cache(cache: list[Any], *, step: int | None = None) -> dict[str, int | str]: + """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} + 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 + rows = int(getattr(entry, "offset", 0) or 0) + if ( + _is_trimmable(entry) + and _is_stock_plain_kv(entry) + and rows == 0 + and getattr(entry, "keys", None) is None + ): + cache[idx] = SegmentedKVCache(step=int(step or getattr(entry, "step", 256))) + stats["entries"] = int(stats["entries"]) + 1 + else: + stats["skipped"] = int(stats["skipped"]) + 1 + 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 _env_on("MTPLX_SEGMENTED_KV_SSD", default=True) + + +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 _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 + ): + 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) + + +# ---- /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_requested(): + return None + if not segmented_kv_enabled(): + # 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]] = [] + all_states: list[SegmentedKVState] = [] + 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) + all_states.append(state) + 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, + "requested": True, + "model_support": model_support(), + "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), + "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 ----------------------------------------------------------------------------- + + +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..df0693797 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 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 + # 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 458b38a55..0f6811a53 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -35,6 +35,15 @@ 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, + segmented_kv_enabled, + segmented_kv_health, + segmented_ssd_enabled, + unique_state_bytes, + walk_segmented_states, +) def _policy_uses_committed_history(policy: str | None) -> bool: @@ -390,6 +399,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 +410,29 @@ 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.""" + 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: return _tree_nbytes(snapshot.states) + _tree_nbytes(snapshot.meta_states) @@ -1266,7 +1302,12 @@ 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()) + 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: @@ -2361,6 +2402,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) @@ -2911,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, @@ -3113,6 +3163,14 @@ def _enqueue_cold_entry( cold["enabled"] = False cold["skip_reason"] = "no_cold_tier" return + 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" + return put_entry = getattr(self.cold_tier, "put_entry", None) if not callable(put_entry): if cold is not None: @@ -3278,6 +3336,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) and not segmented_ssd_enabled(): + return if self.cold_enqueue_dispatch is None: self.eviction_log.append( { @@ -3340,6 +3400,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) and not segmented_ssd_enabled(): + return False try: snapshot = snapshot_cache_lazy_hybrid(entry.cache_ref) # Either form of the draft history goes to disk: without it a @@ -3419,7 +3481,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 +3586,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_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_sdpa_segmented_hd128.py b/tests/test_sdpa_segmented_hd128.py new file mode 100644 index 000000000..5039bc87e --- /dev/null +++ b/tests/test_sdpa_segmented_hd128.py @@ -0,0 +1,305 @@ +"""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 + + +@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 + + 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") + 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] = [] + + 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 diff --git a/tests/test_segmented_kv.py b/tests/test_segmented_kv.py new file mode 100644 index 000000000..b5c961630 --- /dev/null +++ b/tests/test_segmented_kv.py @@ -0,0 +1,632 @@ +"""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 mlx.core as mx +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, + 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, + route_counts, + 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") + + +@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) + 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_n == 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 + 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 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._sealed[1].segment.capacity == 40 + + +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) + + +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_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)]) + 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_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] + 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 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): + 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() + merges = route_counts.get("merge", 0) + 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 route_counts.get("merge", 0) > merges + + +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 + + +@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)) + 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) + 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() + cache, q = _lse_cache() + attend_segments_lse(q, cache, scale=0.25) + 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: + 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[..., 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") + + 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 {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) + 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" + + +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 + + +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 diff --git a/tests/test_segmented_kv_attention.py b/tests/test_segmented_kv_attention.py new file mode 100644 index 000000000..7925ae3f4 --- /dev/null +++ b/tests/test_segmented_kv_attention.py @@ -0,0 +1,235 @@ +"""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_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.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) + 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") +@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") + 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 + + +@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 + + +@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, 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 + 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 new file mode 100644 index 000000000..ab3788976 --- /dev/null +++ b/tests/test_segmented_kv_bank.py @@ -0,0 +1,292 @@ +"""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 +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 + +LAYERS, HEADS, DIM = 4, 4, 256 +BASE, TURN = 4096, 512 +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 + + 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() + assert bank.clear(session_id="s") == 1 + del entry + gc.collect() + mx.clear_cache() + assert held - mx.get_active_memory() >= BASE * ROW_BYTES * 0.9 + + +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 + + 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 + + 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) + _, 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 + + +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)) + _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"} <= 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_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 diff --git a/tests/test_segmented_kv_lse_route.py b/tests/test_segmented_kv_lse_route.py new file mode 100644 index 000000000..1c9046ceb --- /dev/null +++ b/tests/test_segmented_kv_lse_route.py @@ -0,0 +1,244 @@ +"""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 _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.""" + 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_one_path.py b/tests/test_segmented_kv_one_path.py new file mode 100644 index 000000000..4c3587e09 --- /dev/null +++ b/tests/test_segmented_kv_one_path.py @@ -0,0 +1,149 @@ +"""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 + +import mtplx.segmented_kv as segmented_kv_module +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") + # 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): + 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_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_a_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_stock_caches_with_rows_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 new file mode 100644 index 000000000..f9123101b --- /dev/null +++ b/tests/test_segmented_kv_ssd.py @@ -0,0 +1,356 @@ +"""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 + +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 +from mtplx.segmented_kv import ( + SegmentedKVCache, + SegmentedKVState, + SegmentedRows, + segmented_kv_enabled, +) + +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) + + +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") + 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() + + +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 + + +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): 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: + 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) + 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 new file mode 100644 index 000000000..909a54fce --- /dev/null +++ b/tests/test_segmented_kv_support.py @@ -0,0 +1,514 @@ +"""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() + # 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 + 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"), + ([], 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 + + +@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_selects_the_lse_route_instead_of_switching_segments_off(monkeypatch) -> None: + monkeypatch.setattr("mtplx.nax_verify.nax_available", lambda: False) + 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: + 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) + 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_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 mlx_lm.models.qwen3_next import Qwen3NextAttention + + from mtplx.attention_split import configure_split_full_attention + + 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 + + +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) + 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 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)])) + stats = configure_split_full_attention(model) + 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 + + +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)])) + 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 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