From 3c6d9f1a4a43353d72f27593322069266c21bb00 Mon Sep 17 00:00:00 2001 From: Jack Mazac <114446521+jackmaxwil@users.noreply.github.com> Date: Thu, 5 Nov 2026 04:57:53 +0000 Subject: [PATCH 1/8] ramp: keep DeepSeek-V4 FP8 residents as their source bytes (opt-in mxfp8) bind_deepseek_v4_flash_non_vq_weights(resident_format="mxfp8") binds every FP8 resident that lands on a plain nn.Linear, or on wo_a's MultiLinear, as its shipped e4m3 bytes and ue8m0 exponents in a weights-only mxfp8 module (nn.QuantizedLinear, or mlx-lm's QuantizedMultiLinear for wo_a) instead of decoding it to bf16. Every weight dequantizes bit-identically to the bf16 decode, at 8.25 instead of 16 bits per weight. Activations are not quantized. The default stays "bf16", so nothing changes until the Mac kill test passes. load_authenticated_dsv4_composite passes the option through, and the report lists the repacked tensors and the .scales they created, without counting the scales as logical tensors. keep.convert.fp8_block.repack_fp8_block_as_mxfp8 does the byte copy and refuses anything that would need rounding: scales that are not exact powers of two, a NaN scale byte, or a 32-wide group that would straddle two source blocks. Building the same weights through mx.quantize(mode= "mxfp8") is not exact (it changed 608 of 98,304 values in a check), which is why the copy goes through this function. benchmarks/bench_dsv4_mxfp8_residents.py is the kill test: bit-exactness against today's decode, then M=1 latency against bf16 and affine-8 g64, with wo_a timed in its grouped form. Kill if any bit differs, if mxfp8 is under 1.5x bf16, or if it is over 1.1x affine-8. On CPU MLX 0.31.2 the exactness half passes at all eight V4-Flash resident shapes; the speed half has to run on the Mac. Estimate from the roadmap's byte table: on V4-Flash this would shrink residents from about 15.7 GB to about 9.9 GB in memory and cut decode reads by a similar amount. --- benchmarks/bench_dsv4_mxfp8_residents.py | 415 +++++++++++++++++++ src/keep/convert/fp8_block.py | 125 +++++- src/ramp/models/deepseek_v4_flash_adapter.py | 191 ++++++++- src/ramp/models/dsv4_composite_loader.py | 16 +- tests/test_bench_dsv4_mxfp8_residents.py | 168 ++++++++ tests/test_deepseek_v4_flash_adapter.py | 139 +++++++ tests/test_fp8_block_dequant.py | 142 +++++++ 7 files changed, 1189 insertions(+), 7 deletions(-) create mode 100644 benchmarks/bench_dsv4_mxfp8_residents.py create mode 100644 tests/test_bench_dsv4_mxfp8_residents.py diff --git a/benchmarks/bench_dsv4_mxfp8_residents.py b/benchmarks/bench_dsv4_mxfp8_residents.py new file mode 100644 index 0000000..cbfcae5 --- /dev/null +++ b/benchmarks/bench_dsv4_mxfp8_residents.py @@ -0,0 +1,415 @@ +"""Kill test: keep DeepSeek-V4 FP8 residents as their source bytes (mxfp8). + +Two checks in one fresh process: + +1. **Exactness.** Each chosen resident, byte-copied into MLX ``mxfp8`` by + :func:`keep.convert.fp8_block.repack_fp8_block_as_mxfp8`, must dequantize + bit-identically to :func:`keep.convert.fp8_block.dequantize_fp8_block`, + which is what RAMP binds in bf16 today. +2. **Decode-shape speed (M=1).** Time ``quantized_matmul(mode="mxfp8")`` + against the bf16 matmul RAMP runs today and an affine-8 g64 baseline (the + roadmap's S1 alternative). + +Verdict, from the ranked plan: **KILL** if any bit differs, if mxfp8 is less +than 1.5x the bf16 speed, or if it is more than 10% slower than affine-8. +Speed is judged on the summed per-call medians across the chosen tensors, the +way one decode step reads them. + +Inputs: a resident package or the source checkpoint (``--checkpoint``, a +directory with ``model.safetensors.index.json``), or ``--synthetic`` FP8 +tensors at the real DeepSeek-V4-Flash shapes, which needs no download. +Weights are cycled across ``--copies`` buffers so each timed call reads DRAM, +not a cache. Speed only means something on an Apple GPU; on a CPU device the +script still runs the exactness half and marks the speed rows invalid. + +Example (M4 Max, no model download needed):: + + uv run python benchmarks/bench_dsv4_mxfp8_residents.py --synthetic \\ + --output-json artifacts/benchmarks/dsv4-mxfp8-residents-m4max.json +""" + +from __future__ import annotations + +import argparse +import datetime as _dt +import json +import os +import platform +import statistics +import sys +import time +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import mlx.core as mx +import numpy as np + +from keep.convert.fp8_block import dequantize_fp8_block, repack_fp8_block_as_mxfp8 +from keep.io.source_safetensors import read_safetensors_file_header + +REPO_ROOT = Path(__file__).resolve().parents[1] +HEAVY_JOB_LOCK = REPO_ROOT / ".keep-heavy-job.lock" +RECORD_TYPE = "dsv4_mxfp8_resident_kill_test_v1" +FP8_BLOCK = (128, 128) + +#: DeepSeek-V4-Flash-0731 resident shapes (out, in) as they ship, from the +#: model defaults: hidden 4096, 64 heads x 512, q_lora_rank 1024, o_groups 8 x +#: o_lora_rank 1024, shared expert width 2048. ``wo_a`` ships flat and runs as +#: an 8-group MultiLinear; see :data:`GROUPED_TENSORS`. +V4_FLASH_RESIDENT_SHAPES: dict[str, tuple[int, int]] = { + "attn.wq_a": (1024, 4096), + "attn.wq_b": (32768, 1024), + "attn.wkv": (512, 4096), + "attn.wo_a": (8192, 4096), + "attn.wo_b": (4096, 8192), + "ffn.shared_experts.w1": (2048, 4096), + "ffn.shared_experts.w2": (4096, 2048), + "ffn.shared_experts.w3": (2048, 4096), +} + +#: Tensors timed as a grouped (batched) matmul, the way RAMP calls them: +#: ``wo_a`` is ``(o_groups * o_lora_rank, in)`` on disk and +#: ``(o_groups, o_lora_rank, in)`` at runtime, fed one row per group. +GROUPED_TENSORS: dict[str, int] = {"attn.wo_a": 8} + +MIN_SPEEDUP_OVER_BF16 = 1.5 +MAX_SLOWDOWN_VS_AFFINE8 = 1.10 +_FORBIDDEN_ENV = ("GLM_MLX_WIRED_LIMIT_GB", "GLM_SINGLE_HOST_MLX_WIRED_LIMIT_GB") + + +def synthetic_fp8_pair( + rng: np.random.Generator, shape: tuple[int, int] +) -> tuple[np.ndarray, np.ndarray]: + """Random finite e4m3 codes plus ue8m0 block exponents at ``shape``.""" + + rows, cols = shape + codes = rng.integers(0, 256, size=(rows, cols), dtype=np.uint8) + codes = np.where((codes & 0x7F) == 0x7F, 0x7E, codes).astype(np.uint8) + blocks = (-(-rows // FP8_BLOCK[0]), -(-cols // FP8_BLOCK[1])) + scales = rng.integers(115, 126, size=blocks, dtype=np.uint8) + return codes, scales + + +class _CheckpointReader: + """Raw-byte reader over a safetensors index; one memmap per shard.""" + + def __init__(self, root: Path) -> None: + index = json.loads((root / "model.safetensors.index.json").read_text()) + self.root = root + self.weight_map: dict[str, str] = dict(index["weight_map"]) + self._shards: dict[str, tuple[np.memmap, int, Mapping[str, Any]]] = {} + + def raw(self, name: str) -> tuple[np.ndarray, str, tuple[int, ...]]: + shard = self.weight_map[name] + if shard not in self._shards: + path = self.root / shard + header = read_safetensors_file_header(path) + self._shards[shard] = ( + np.memmap(path, dtype=np.uint8, mode="r"), + header.header_size, + header.tensors, + ) + mm, header_size, tensors = self._shards[shard] + tensor = tensors[name] + start, end = tensor.data_offsets + base = 8 + header_size + return np.array(mm[base + start : base + end]), tensor.dtype, tuple(tensor.shape) + + def fp8_pair(self, prefix: str) -> tuple[np.ndarray, np.ndarray]: + codes, code_dtype, shape = self.raw(f"{prefix}.weight") + scales, scale_dtype, scale_shape = self.raw(f"{prefix}.scale") + if code_dtype != "F8_E4M3" or scale_dtype != "F8_E8M0": + raise ValueError( + f"{prefix} is {code_dtype}/{scale_dtype}, not an F8_E4M3 resident" + ) + return codes.reshape(shape), scales.reshape(scale_shape) + + +def check_exactness(codes: np.ndarray, scales: np.ndarray) -> dict[str, Any]: + """Compare the mxfp8 byte copy against today's bf16 decode, bit for bit.""" + + packed, group_scales = repack_fp8_block_as_mxfp8(codes, scales, block_size=FP8_BLOCK) + reference = mx.array(dequantize_fp8_block(codes, scales, block_size=FP8_BLOCK)).astype( + mx.bfloat16 + ) + restored = mx.dequantize( + mx.array(packed), + mx.array(group_scales), + group_size=32, + bits=8, + mode="mxfp8", + dtype=mx.bfloat16, + ) + mismatches = int(mx.sum(restored.view(mx.uint16) != reference.view(mx.uint16)).item()) + return { + "elements": int(codes.size), + "bit_mismatches": mismatches, + "bf16_bytes": int(codes.size) * 2, + "mxfp8_bytes": int(packed.nbytes + group_scales.nbytes), + } + + +def _median_call_ms( + fn: Callable[[int], mx.array], *, iterations: int, warmup: int +) -> float: + for i in range(warmup): + mx.eval(fn(i)) + started = time.perf_counter() + for i in range(iterations): + mx.eval(fn(i)) + return (time.perf_counter() - started) * 1_000.0 / iterations + + +def time_variants( + codes: np.ndarray, + scales: np.ndarray, + *, + copies: int, + iterations: int, + warmup: int, + rounds: int, + seed: int, + groups: int = 1, +) -> dict[str, Any]: + """Per-call M=1 latency for bf16, mxfp8 and affine-8 g64, interleaved. + + ``groups > 1`` times the grouped form: rows split into ``groups`` equal + blocks, each multiplied by its own activation row, as ``wo_a`` runs. + """ + + rows, cols = codes.shape + if groups < 1 or rows % groups: + raise ValueError(f"{rows} rows do not split into {groups} groups") + packed, group_scales = repack_fp8_block_as_mxfp8(codes, scales, block_size=FP8_BLOCK) + dense = mx.array(dequantize_fp8_block(codes, scales, block_size=FP8_BLOCK)).astype( + mx.bfloat16 + ) + if groups > 1: + packed = packed.reshape(groups, rows // groups, -1) + group_scales = group_scales.reshape(groups, rows // groups, -1) + dense = dense.reshape(groups, rows // groups, cols) + affine_w, affine_s, affine_b = mx.quantize(dense, group_size=64, bits=8) + # Distinct buffers per copy so consecutive calls cannot hit a cache. + bf16_copies = [dense + mx.zeros_like(dense) for _ in range(copies)] + packed_copies = [mx.array(packed) for _ in range(copies)] + scale_copies = [mx.array(group_scales) for _ in range(copies)] + affine_copies = [ + (affine_w + mx.zeros_like(affine_w), affine_s + 0, affine_b + 0) for _ in range(copies) + ] + mx.eval(bf16_copies, packed_copies, scale_copies, affine_copies) + x_shape = (1, groups, 1, cols) if groups > 1 else (1, cols) + x = mx.random.normal(x_shape, key=mx.random.key(seed)).astype(mx.bfloat16) + + def bf16(i: int) -> mx.array: + return x @ bf16_copies[i % copies].swapaxes(-1, -2) + + def mxfp8(i: int) -> mx.array: + j = i % copies + return mx.quantized_matmul( + x, packed_copies[j], scale_copies[j], transpose=True, group_size=32, bits=8, mode="mxfp8" + ) + + def affine8(i: int) -> mx.array: + w, s, b = affine_copies[i % copies] + return mx.quantized_matmul(x, w, s, b, transpose=True, group_size=64, bits=8) + + variants = {"bf16": bf16, "mxfp8": mxfp8, "affine8_g64": affine8} + samples: dict[str, list[float]] = {name: [] for name in variants} + order = list(variants) + for round_index in range(rounds): + # Rotate the order so no variant always runs first or last. + rotated = order[round_index % len(order) :] + order[: round_index % len(order)] + for name in rotated: + samples[name].append( + _median_call_ms(variants[name], iterations=iterations, warmup=warmup) + ) + weight_bytes = { + "bf16": rows * cols * 2, + "mxfp8": int(packed.nbytes + group_scales.nbytes), + "affine8_g64": int(affine_w.nbytes + affine_s.nbytes + affine_b.nbytes), + } + out: dict[str, Any] = {} + for name, values in samples.items(): + median_ms = statistics.median(values) + out[name] = { + "median_ms_per_call": median_ms, + "round_ms": values, + "weight_bytes": weight_bytes[name], + "effective_gb_per_s": weight_bytes[name] / (median_ms / 1_000.0) / 1e9, + } + del bf16_copies, packed_copies, scale_copies, affine_copies + return out + + +def verdict( + exactness: Mapping[str, Mapping[str, Any]], + timings: Mapping[str, Mapping[str, Mapping[str, Any]]] | None, + *, + speed_valid: bool, +) -> dict[str, Any]: + """Apply the kill criteria. Exactness alone can kill; speed needs a GPU.""" + + mismatches = sum(int(row["bit_mismatches"]) for row in exactness.values()) + result: dict[str, Any] = { + "bit_mismatches": mismatches, + "exact": mismatches == 0, + "speed_valid": bool(speed_valid and timings), + } + reasons: list[str] = [] + if mismatches: + reasons.append(f"{mismatches} dequantized weights differ from the bf16 decode") + if timings: + totals = { + variant: sum(float(rows[variant]["median_ms_per_call"]) for rows in timings.values()) + for variant in ("bf16", "mxfp8", "affine8_g64") + } + speedup = totals["bf16"] / totals["mxfp8"] + vs_affine = totals["mxfp8"] / totals["affine8_g64"] + result.update( + { + "summed_median_ms": totals, + "mxfp8_speedup_over_bf16": speedup, + "mxfp8_time_over_affine8": vs_affine, + } + ) + if speed_valid: + if speedup < MIN_SPEEDUP_OVER_BF16: + reasons.append( + f"mxfp8 is {speedup:.3f}x bf16, below the {MIN_SPEEDUP_OVER_BF16}x bar" + ) + if vs_affine > MAX_SLOWDOWN_VS_AFFINE8: + reasons.append( + f"mxfp8 takes {vs_affine:.3f}x affine-8's time, above {MAX_SLOWDOWN_VS_AFFINE8}x" + ) + if reasons: + result["verdict"] = "KILL" + elif result["speed_valid"]: + result["verdict"] = "PASS" + else: + result["verdict"] = "EXACT_ONLY" # exactness passed; rerun on an Apple GPU for speed + result["reasons"] = reasons + return result + + +def run( + *, + tensors: Sequence[str], + layer: int, + checkpoint: Path | None, + synthetic_shapes: Mapping[str, tuple[int, int]] | None, + copies: int, + iterations: int, + warmup: int, + rounds: int, + seed: int, + exactness_only: bool, +) -> dict[str, Any]: + if (checkpoint is None) == (synthetic_shapes is None): + raise ValueError("pass exactly one of checkpoint or synthetic_shapes") + rng = np.random.default_rng(seed) + reader = _CheckpointReader(checkpoint) if checkpoint is not None else None + on_gpu = mx.default_device() == mx.gpu + exactness: dict[str, dict[str, Any]] = {} + timings: dict[str, dict[str, Any]] = {} + shapes: dict[str, list[int]] = {} + for tensor in tensors: + if reader is not None: + codes, scales = reader.fp8_pair(f"layers.{layer}.{tensor}") + else: + codes, scales = synthetic_fp8_pair(rng, synthetic_shapes[tensor]) + shapes[tensor] = list(codes.shape) + exactness[tensor] = check_exactness(codes, scales) + if not exactness_only: + timings[tensor] = time_variants( + codes, + scales, + copies=copies, + iterations=iterations, + warmup=warmup, + rounds=rounds, + seed=seed, + groups=GROUPED_TENSORS.get(tensor, 1), + ) + mx.clear_cache() + return { + "record_type": RECORD_TYPE, + "created_utc": _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="seconds"), + "source": "synthetic" if reader is None else str(checkpoint), + "layer": None if reader is None else layer, + "shapes": shapes, + "device": str(mx.default_device()), + "device_info": {k: (v if isinstance(v, (int, float, str)) else str(v)) for k, v in mx.device_info().items()} + if on_gpu + else {}, + "mlx_version": mx.__version__, + "platform": platform.platform(), + "settings": { + "copies": copies, + "iterations": iterations, + "warmup": warmup, + "rounds": rounds, + "seed": seed, + }, + "exactness": exactness, + "timings": timings, + "verdict": verdict(exactness, timings or None, speed_valid=on_gpu), + } + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--checkpoint", type=Path, help="resident package or source checkpoint dir") + source.add_argument("--synthetic", action="store_true", help="random FP8 at V4-Flash shapes") + parser.add_argument("--layer", type=int, default=20) + parser.add_argument("--tensors", default=",".join(V4_FLASH_RESIDENT_SHAPES)) + parser.add_argument("--copies", type=int, default=8) + parser.add_argument("--iterations", type=int, default=200) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rounds", type=int, default=5) + parser.add_argument("--seed", type=int, default=20261105) + parser.add_argument("--exactness-only", action="store_true") + parser.add_argument("--no-lock", action="store_true", help="skip the heavy-job lock") + parser.add_argument("--output-json", type=Path) + args = parser.parse_args(argv) + + forbidden = [name for name in _FORBIDDEN_ENV if name in os.environ] + if forbidden: + parser.error("unset custom MLX wired limits first: " + ", ".join(forbidden)) + tensors = [name.strip() for name in args.tensors.split(",") if name.strip()] + unknown = [name for name in tensors if args.synthetic and name not in V4_FLASH_RESIDENT_SHAPES] + if unknown: + parser.error("no synthetic shape for: " + ", ".join(unknown)) + + def _go() -> dict[str, Any]: + return run( + tensors=tensors, + layer=args.layer, + checkpoint=args.checkpoint, + synthetic_shapes=V4_FLASH_RESIDENT_SHAPES if args.synthetic else None, + copies=args.copies, + iterations=args.iterations, + warmup=args.warmup, + rounds=args.rounds, + seed=args.seed, + exactness_only=args.exactness_only, + ) + + if args.no_lock: + result = _go() + else: + from keep.quality.dsv4_teacher_agreement import heavy_job_lock + + with heavy_job_lock(HEAVY_JOB_LOCK, holder="bench_dsv4_mxfp8_residents"): + result = _go() + text = json.dumps(result, indent=2, sort_keys=True) + if args.output_json is not None: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(text + "\n") + print(text) + return 0 if result["verdict"]["verdict"] != "KILL" else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/keep/convert/fp8_block.py b/src/keep/convert/fp8_block.py index a0f0b83..031945f 100644 --- a/src/keep/convert/fp8_block.py +++ b/src/keep/convert/fp8_block.py @@ -19,13 +19,27 @@ Pure numpy so it runs headless; the MLX fast path can come later if decode throughput matters. + +:func:`repack_fp8_block_as_mxfp8` is the lossless alternative to decoding: it +re-lays the same e4m3 bytes and ue8m0 exponents out in MLX's ``mxfp8`` layout +(group 32 along the input axis), so a resident can stay at 8.25 bits per weight +instead of expanding to bf16, with every dequantized value bit-identical. """ from __future__ import annotations import numpy as np -__all__ = ["decode_e4m3", "decode_ue8m0", "dequantize_fp8_block"] +__all__ = [ + "MXFP8_GROUP_SIZE", + "decode_e4m3", + "decode_ue8m0", + "dequantize_fp8_block", + "repack_fp8_block_as_mxfp8", +] + +#: MLX ``mxfp8`` supports exactly one group size: 32 elements per E8M0 scale. +MXFP8_GROUP_SIZE = 32 _E4M3_TABLE: np.ndarray | None = None @@ -130,3 +144,112 @@ def dequantize_fp8_block( c0, c1 = j * block_cols, min((j + 1) * block_cols, cols) decoded[r0:r1, c0:c1] *= row_scales[j] return decoded + + +def _power_of_two_scales_as_ue8m0(scales: np.ndarray) -> np.ndarray: + """Encode literal power-of-two multipliers as ue8m0 bytes, or refuse. + + Only exact powers of two in ``2**-127 .. 2**127`` have an E8M0 encoding. + Anything else would need rounding, and rounding is exactly what a lossless + repack must never do, so it raises instead. + """ + + values = scales.astype(np.float64) + if not np.all(np.isfinite(values)) or np.any(values <= 0.0): + raise ValueError("power-of-two block scales must be finite and positive") + mantissa, exponent = np.frexp(values) + if np.any(mantissa != 0.5): + raise ValueError( + "block scales are not all exact powers of two; an mxfp8 repack " + "would have to round them, so it is refused" + ) + unbiased = exponent.astype(np.int64) - 1 + if np.any(unbiased < -127) or np.any(unbiased > 127): + raise ValueError("power-of-two block scales are outside the E8M0 range") + return (unbiased + 127).astype(np.uint8) + + +def repack_fp8_block_as_mxfp8( + codes: np.ndarray, + scales: np.ndarray, + *, + block_size: tuple[int, int] = (128, 128), + group_size: int = MXFP8_GROUP_SIZE, +) -> tuple[np.ndarray, np.ndarray]: + """Re-lay block-scaled e4m3 weights out as MLX ``mxfp8`` without changing a bit. + + Returns ``(packed, group_scales)``: + + - ``packed``: ``uint32`` of shape ``(rows, cols // 4)``, the unchanged e4m3 + code bytes read little-endian four to a word, which is the element order + ``mx.dequantize(..., mode="mxfp8")`` expects. + - ``group_scales``: ``uint8`` of shape ``(rows, cols // 32)``, each block's + ue8m0 exponent repeated over the rows and the 32-wide column groups the + block covers. + + Dequantizing the result reproduces :func:`dequantize_fp8_block` exactly: + e4m3 times a power of two is representable in float32 and bf16 alike. + Building the same weights with ``mx.quantize(..., mode="mxfp8")`` instead is + *not* exact -- it re-derives each group's exponent by rounding -- so this + function is the only supported way to produce an mxfp8 resident. + + Requires ``cols % 32 == 0`` and ``block_size[1] % 32 == 0`` so every + 32-wide group sits inside one source block. ue8m0 byte 255 (NaN) is + refused rather than bound. + """ + + if not isinstance(codes, np.ndarray): + raise ValueError(f"codes must be a numpy ndarray, got {type(codes).__name__}") + if not isinstance(scales, np.ndarray): + raise ValueError(f"scales must be a numpy ndarray, got {type(scales).__name__}") + if codes.dtype != np.uint8: + raise ValueError(f"e4m3 codes must be uint8, got {codes.dtype}") + if codes.ndim != 2: + raise ValueError(f"codes must be 2-D, got shape {codes.shape}") + if group_size != MXFP8_GROUP_SIZE: + raise ValueError( + f"mxfp8 has a fixed group size of {MXFP8_GROUP_SIZE}, got {group_size!r}" + ) + block_rows, block_cols = block_size + for label, value in (("block_size[0]", block_rows), ("block_size[1]", block_cols)): + if ( + isinstance(value, bool) + or not isinstance(value, (int, np.integer)) + or value <= 0 + ): + raise ValueError(f"{label} must be a positive int, got {value!r}") + rows, cols = codes.shape + if cols % group_size != 0: + raise ValueError( + f"input dimension {cols} is not a multiple of the mxfp8 group size {group_size}" + ) + if block_cols % group_size != 0: + raise ValueError( + f"block width {block_cols} is not a multiple of the mxfp8 group size " + f"{group_size}, so a group would straddle two source scales" + ) + expected = (-(-rows // block_rows), -(-cols // block_cols)) + if scales.shape != expected: + raise ValueError( + f"scales shape {scales.shape} does not match " + f"{expected} blocks for weight {codes.shape} " + f"at block size {block_size}" + ) + if scales.dtype == np.uint8: + if np.any(scales == 255): + raise ValueError("ue8m0 byte 255 is NaN; refusing to repack a NaN block scale") + exponents = scales + elif np.issubdtype(scales.dtype, np.floating): + exponents = _power_of_two_scales_as_ue8m0(scales) + else: + raise ValueError( + f"scales must be uint8 (ue8m0-encoded) or a floating dtype holding " + f"powers of two, got dtype {scales.dtype}" + ) + + # ``astype`` copies, so the result never aliases the caller's code buffer. + packed = np.ascontiguousarray(codes).view(" int: @@ -2822,6 +2833,9 @@ def to_dict(self) -> dict[str, object]: "skipped_out_of_scope_tensors": list(self.skipped_out_of_scope_tensors), "skipped_unmatched_tensors": list(self.skipped_unmatched_tensors), "missing_model_parameters": list(self.missing_model_parameters), + "resident_format": self.resident_format, + "mxfp8_repacked_tensors": list(self.mxfp8_repacked_tensors), + "mxfp8_scale_parameters": list(self.mxfp8_scale_parameters), } @@ -2985,6 +2999,113 @@ def _decode_source_tensor( return mx.array(values), "direct" +RESIDENT_FORMATS = ("bf16", "mxfp8") + + +def _repack_source_tensor_mxfp8( + dtype: str, + raw: bytes, + shape: tuple[int, ...], + scale: tuple[str, bytes, tuple[int, ...]] | None, +) -> tuple[mx.array, mx.array]: + """Byte-copy an FP8 resident into MLX mxfp8 ``(weight, scales)`` arrays. + + Same dtype guards as the decode path. The values this binds dequantize to + exactly what :func:`_decode_source_tensor` would have bound in bf16. + """ + + if scale is None or dtype != _FP8_CODE_DTYPE: + raise ValueError( + f"mxfp8 repack needs {_FP8_CODE_DTYPE} codes with a .scale companion, " + f"got {dtype} {'with' if scale is not None else 'without'} a scale" + ) + scale_dtype, scale_raw, scale_shape = scale + if scale_dtype != _FP8_SCALE_DTYPE: + raise ValueError(f"block scales must be {_FP8_SCALE_DTYPE}, got {scale_dtype}") + codes = np.frombuffer(raw, dtype=np.uint8).reshape(shape) + scales = np.frombuffer(scale_raw, dtype=np.uint8).reshape(scale_shape) + packed, group_scales = repack_fp8_block_as_mxfp8(codes, scales, block_size=_FP8_BLOCK) + return mx.array(packed), mx.array(group_scales) + + +def _mxfp8_eligible_module_paths( + model: nn.Module, + plan: Sequence[tuple[str, str, str | None]], +) -> dict[str, str]: + """Map ``target .weight`` -> module path for residents that can stay mxfp8. + + Eligible means: an FP8 source (it has a ``.scale``), a plain ``nn.Linear`` + target or the ``wo_a`` ``MultiLinear`` (whose flat source rows regroup by + a pure reshape), and an input width divisible by the mxfp8 group. Anything + else stays on the bf16 decode path. + """ + + modules = dict(model.named_modules()) + eligible: dict[str, str] = {} + if _FP8_BLOCK[1] % MXFP8_GROUP_SIZE != 0: + return eligible + for source, target, scale_name in plan: + if scale_name is None or not target.endswith(".weight"): + continue + module_path = target[: -len(".weight")] + module = modules.get(module_path) + if type(module) is nn.Linear: + if module.weight.ndim != 2: + continue + elif type(module) is MultiLinear and _RESHAPE_ALLOWED_RE.search(source): + if module.weight.ndim != 3: + continue + else: + continue + if module.weight.shape[-1] % MXFP8_GROUP_SIZE != 0: + continue + eligible[target] = module_path + return eligible + + +def _swap_linears_for_mxfp8(model: nn.Module, module_paths: Iterable[str]) -> None: + """Replace the named linear modules with weights-only mxfp8 ones. + + ``nn.Linear`` becomes ``nn.QuantizedLinear`` and ``MultiLinear`` becomes + mlx-lm's ``QuantizedMultiLinear``. ``quantize_input=False`` is explicit: + activations stay in floating point, so the bound weights reproduce the + bf16 teacher's weights exactly. The placeholder quantization of the + skeleton weights is lazy and is overwritten by the bind before anything + evaluates it. + """ + + wanted = set(module_paths) + if not wanted: + return + nn.quantize( + model, + group_size=MXFP8_GROUP_SIZE, + bits=8, + mode="mxfp8", + quantize_input=False, + class_predicate=lambda path, module: path in wanted, + ) + modules = dict(model.named_modules()) + not_swapped = sorted( + path + for path in wanted + if type(modules.get(path)) not in (nn.QuantizedLinear, QuantizedMultiLinear) + or getattr(modules.get(path), "mode", None) != "mxfp8" + ) + if not_swapped: + raise RuntimeError( + "mxfp8 module swap left plain modules behind: " + ", ".join(not_swapped[:10]) + ) + + +def _linear_input_dims(module: nn.Module) -> int: + """Input width of a plain or quantized linear layer.""" + + if isinstance(module, nn.QuantizedLinear): + return int(module.weight.shape[1]) * 32 // int(module.bits) + return int(module.weight.shape[1]) + + def bind_deepseek_v4_flash_non_vq_weights( model: DeepseekV4FlashVQModel, checkpoint_dir: str | Path, @@ -2993,6 +3114,7 @@ def bind_deepseek_v4_flash_non_vq_weights( include_mtp: bool = True, strict: bool = True, compute_dtype: mx.Dtype = mx.bfloat16, + resident_format: str = "bf16", ) -> DeepseekV4FlashNonVQBindReport: """Bind every non-routed tensor, shard by shard, off the safetensors index. @@ -3001,12 +3123,25 @@ def bind_deepseek_v4_flash_non_vq_weights( and ``F32`` tensors load directly; ``I64`` routing tables keep their native dtype. Routed experts are skipped -- they are the VQ path. + ``resident_format="mxfp8"`` (opt-in) keeps every FP8 resident that lands on + a plain ``nn.Linear`` or on ``wo_a``'s ``MultiLinear`` as its source bytes + instead: the module becomes a weights-only ``nn.QuantizedLinear`` (or + mlx-lm ``QuantizedMultiLinear``) in ``mode="mxfp8"``, holding the unchanged + e4m3 codes and their ue8m0 exponents (8.25 bits per weight instead of 16), + and dequantizes bit-identically to the bf16 decode. Non-FP8 tensors and + inputs not divisible by 32 keep the default path. The default, ``"bf16"``, + is unchanged. + ``layers`` restricts the backbone layers touched (a layer-0 smoke test does not need 163 GB of I/O). ``strict`` is only meaningful for a full bind: it is forced off when ``layers`` or ``include_mtp=False`` deliberately leaves parameters unbound. """ + if resident_format not in RESIDENT_FORMATS: + raise ValueError( + f"resident_format must be one of {RESIDENT_FORMATS}, got {resident_format!r}" + ) checkpoint_path = Path(checkpoint_dir) weight_map = _load_weight_map(checkpoint_path) partial_scope = layers is not None or not include_mtp @@ -3078,6 +3213,14 @@ def bind_deepseek_v4_flash_non_vq_weights( + ", ".join(unmatched[:20]) ) + mxfp8_targets: dict[str, str] = {} + if resident_format == "mxfp8": + mxfp8_targets = _mxfp8_eligible_module_paths(model, plan) + _swap_linears_for_mxfp8(model, mxfp8_targets.values()) + target_params = {key: value for key, value in tree_flatten(model.parameters())} + mxfp8_repacked: list[str] = [] + mxfp8_scales: list[str] = [] + # Group by shard so each file is opened (and its header parsed) once. by_shard: dict[str, list[tuple[str, str, str | None]]] = {} for source, target, scale_name in plan: @@ -3116,6 +3259,43 @@ def bind_deepseek_v4_flash_non_vq_weights( scale_tensor.shape, ) + if target in mxfp8_targets: + packed, group_scales = _repack_source_tensor_mxfp8( + tensor.dtype, raw, tensor.shape, scale_payload + ) + scales_target = f"{mxfp8_targets[target]}.scales" + regrouped = False + for name, value in ((target, packed), (scales_target, group_scales)): + expected = target_params.get(name) + if ( + expected is not None + and tuple(value.shape) != tuple(expected.shape) + and _RESHAPE_ALLOWED_RE.search(source) + and value.size == expected.size + ): + # ``wo_a``: flat source rows regroup into + # MultiLinear's (groups, rows, cols) layout. Rows + # carry their own expanded scales, so this is a + # pure relabel of whole rows. + value = value.reshape(expected.shape) + regrouped = True + if expected is None or tuple(value.shape) != tuple(expected.shape): + raise ValueError( + f"{source} repacks to {name} shape {tuple(value.shape)} " + f"but the mxfp8 module expects " + f"{None if expected is None else tuple(expected.shape)}" + ) + batch.append((name, value)) + if regrouped: + reshaped.append(source) + mxfp8_repacked.append(source) + mxfp8_scales.append(scales_target) + if target.startswith("mtp_drafter."): + bound_mtp.append(target) + else: + bound_backbone.append(target) + continue + value, kind = _decode_source_tensor( tensor.dtype, raw, @@ -3167,7 +3347,7 @@ def bind_deepseek_v4_flash_non_vq_weights( # config -- for a bind that touched one layer. mx.eval([value for _, value in batch]) - bound_all = set(bound_backbone) | set(bound_mtp) + bound_all = set(bound_backbone) | set(bound_mtp) | set(mxfp8_scales) missing = tuple(sorted(set(target_params) - bound_all)) report = DeepseekV4FlashNonVQBindReport( bound_backbone_parameters=tuple(sorted(bound_backbone)), @@ -3182,6 +3362,9 @@ def bind_deepseek_v4_flash_non_vq_weights( skipped_out_of_scope_tensors=tuple(sorted(skipped_scope)), skipped_unmatched_tensors=tuple(sorted(unmatched)), missing_model_parameters=missing, + resident_format=resident_format, + mxfp8_repacked_tensors=tuple(sorted(mxfp8_repacked)), + mxfp8_scale_parameters=tuple(sorted(mxfp8_scales)), ) if strict and missing: raise ValueError( diff --git a/src/ramp/models/dsv4_composite_loader.py b/src/ramp/models/dsv4_composite_loader.py index e0fb054..9ddc0d4 100644 --- a/src/ramp/models/dsv4_composite_loader.py +++ b/src/ramp/models/dsv4_composite_loader.py @@ -636,13 +636,25 @@ def validate_dsv4_composite_inputs( def load_authenticated_dsv4_composite( inputs: Dsv4ValidatedCompositeInputs, + *, + resident_format: str = "bf16", ) -> Dsv4LoadedComposite: - """Construct and bind exactly the authenticated resident+routed model.""" + """Construct and bind exactly the authenticated resident+routed model. + + ``resident_format="mxfp8"`` keeps FP8 residents as their source bytes + (weights-only mxfp8, bit-identical weights) instead of decoding them to + bf16; see :func:`bind_deepseek_v4_flash_non_vq_weights`. The artifact and + its identity are the same either way -- only the in-memory layout differs. + """ args = deepseek_v4_flash_args_from_config(inputs.config) model = DeepseekV4FlashVQModel(args, with_mtp=True) non_vq = bind_deepseek_v4_flash_non_vq_weights( - model, inputs.resident_dir, include_mtp=True, strict=True + model, + inputs.resident_dir, + include_mtp=True, + strict=True, + resident_format=resident_format, ) backbone = bind_deepseek_v4_flash_vq_experts( model, diff --git a/tests/test_bench_dsv4_mxfp8_residents.py b/tests/test_bench_dsv4_mxfp8_residents.py new file mode 100644 index 0000000..ec8eb40 --- /dev/null +++ b/tests/test_bench_dsv4_mxfp8_residents.py @@ -0,0 +1,168 @@ +"""The mxfp8-resident kill test: its exactness half and its verdict rules.""" + +from __future__ import annotations + +import json +import struct +from pathlib import Path + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") + +from benchmarks import bench_dsv4_mxfp8_residents as bench # noqa: E402 + + +def _row(bf16: float, mxfp8: float, affine8: float) -> dict: + return { + "bf16": {"median_ms_per_call": bf16}, + "mxfp8": {"median_ms_per_call": mxfp8}, + "affine8_g64": {"median_ms_per_call": affine8}, + } + + +_EXACT = {"t": {"bit_mismatches": 0}} + + +def test_verdict_passes_only_when_exact_and_fast_on_a_gpu(): + timings = {"t": _row(bf16=2.0, mxfp8=1.0, affine8=1.0)} + assert bench.verdict(_EXACT, timings, speed_valid=True)["verdict"] == "PASS" + # Same numbers from a CPU run prove exactness only. + assert bench.verdict(_EXACT, timings, speed_valid=False)["verdict"] == "EXACT_ONLY" + assert bench.verdict(_EXACT, None, speed_valid=True)["verdict"] == "EXACT_ONLY" + + +@pytest.mark.parametrize( + ("bf16", "mxfp8", "affine8", "needle"), + [ + (1.4, 1.0, 1.0, "below the 1.5x bar"), + (2.0, 1.2, 1.0, "affine-8"), + ], +) +def test_verdict_kills_on_either_speed_criterion(bf16, mxfp8, affine8, needle): + result = bench.verdict(_EXACT, {"t": _row(bf16, mxfp8, affine8)}, speed_valid=True) + assert result["verdict"] == "KILL" + assert any(needle in reason for reason in result["reasons"]) + + +def test_verdict_kills_on_any_bit_mismatch_even_without_timings(): + result = bench.verdict({"t": {"bit_mismatches": 1}}, None, speed_valid=False) + assert result["verdict"] == "KILL" and result["exact"] is False + + +def test_speed_is_judged_on_the_summed_decode_step_not_per_tensor(): + timings = { + "fast": _row(bf16=3.0, mxfp8=1.0, affine8=1.0), + "slow": _row(bf16=1.0, mxfp8=1.0, affine8=1.0), + } + result = bench.verdict(_EXACT, timings, speed_valid=True) + assert result["mxfp8_speedup_over_bf16"] == pytest.approx(2.0) + assert result["verdict"] == "PASS" + + +def test_synthetic_run_is_exact_and_reports_every_variant(): + result = bench.run( + tensors=["a", "b"], + layer=0, + checkpoint=None, + synthetic_shapes={"a": (64, 256), "b": (130, 192)}, # ragged row and column blocks + copies=2, + iterations=2, + warmup=1, + rounds=2, + seed=0, + exactness_only=False, + ) + assert result["record_type"] == bench.RECORD_TYPE + assert result["verdict"]["exact"] is True + for tensor in ("a", "b"): + assert set(result["timings"][tensor]) == {"bf16", "mxfp8", "affine8_g64"} + assert result["exactness"][tensor]["mxfp8_bytes"] < result["exactness"][tensor]["bf16_bytes"] + if mx.default_device() != mx.gpu: + assert result["verdict"]["verdict"] == "EXACT_ONLY" + + +def _write_fp8_checkpoint(root: Path, tensors: dict[str, tuple[np.ndarray, np.ndarray]]) -> None: + header: dict[str, object] = {} + blobs: list[bytes] = [] + offset = 0 + for prefix, (codes, scales) in sorted(tensors.items()): + for suffix, dtype, array in ( + ("weight", "F8_E4M3", codes), + ("scale", "F8_E8M0", scales), + ): + raw = array.tobytes() + header[f"{prefix}.{suffix}"] = { + "dtype": dtype, + "shape": list(array.shape), + "data_offsets": [offset, offset + len(raw)], + } + offset += len(raw) + blobs.append(raw) + raw_header = json.dumps(header, separators=(",", ":")).encode() + raw_header += b" " * ((-len(raw_header)) % 8) + shard = "model-00001-of-00001.safetensors" + with (root / shard).open("wb") as handle: + handle.write(struct.pack(" 0 + + +# --------------------------------------------------------------------------- +# Opt-in mxfp8 residents: the FP8 bytes stay as they shipped +# --------------------------------------------------------------------------- + + +def _named(model) -> dict: + return dict(model.named_modules()) + + +def test_mxfp8_resident_bind_keeps_the_source_bytes_bit_identically(tmp_path): + from ramp.models.deepseek_v4_flash_adapter import _linear_input_dims + + args = _mtp_args() + _write_synthetic_checkpoint(tmp_path, args, with_mtp=True) + reference = DeepseekV4FlashVQModel(args) + ref_report = bind_deepseek_v4_flash_non_vq_weights(reference, tmp_path, strict=True) + model = DeepseekV4FlashVQModel(args) + report = bind_deepseek_v4_flash_non_vq_weights( + model, tmp_path, strict=True, resident_format="mxfp8" + ) + + assert ref_report.resident_format == "bf16" and ref_report.mxfp8_repacked_tensors == () + assert report.resident_format == "mxfp8" + assert report.missing_model_parameters == () + # Logical accounting is unchanged: the .scales are bookkeeping, not tensors. + assert report.bound_count == ref_report.bound_count + assert set(report.bound_backbone_parameters) == set(ref_report.bound_backbone_parameters) + assert len(report.mxfp8_scale_parameters) == len(report.mxfp8_repacked_tensors) + + repacked = set(report.mxfp8_repacked_tensors) + # Plain linears with a 32-divisible input width stay FP8 ... + for name in ( + "layers.0.attn.wq_a.weight", + "layers.0.attn.wkv.weight", + "layers.0.ffn.shared_experts.w1.weight", + "layers.0.ffn.shared_experts.w2.weight", + "layers.0.ffn.shared_experts.w3.weight", + "mtp.0.main_proj.weight", + ): + assert name in repacked, name + # ... while inputs too narrow for a 32-wide group keep the bf16 decode + # (wo_a's per-group input is 8 wide in this fixture; see the next test). + assert "layers.0.attn.wo_a.weight" not in repacked + assert "layers.0.attn.wo_a.weight" in report.fp8_block_decoded_tensors + assert "layers.0.attn.wq_b.weight" not in repacked # input width 8 + assert set(report.fp8_block_decoded_tensors).isdisjoint(repacked) + + ref_modules, modules = _named(reference), _named(model) + for scales_name in report.mxfp8_scale_parameters: + path = scales_name[: -len(".scales")] + quantized, plain = modules[path], ref_modules[path] + assert type(quantized) is nn.QuantizedLinear, path + assert (quantized.mode, quantized.group_size, quantized.bits) == ("mxfp8", 32, 8) + assert quantized.weight.dtype == mx.uint32 and quantized.scales.dtype == mx.uint8 + restored = mx.dequantize( + quantized.weight, + quantized.scales, + group_size=32, + bits=8, + mode="mxfp8", + dtype=mx.bfloat16, + ) + assert restored.shape == plain.weight.shape + assert np.array_equal( + np.array(restored.view(mx.uint16)), np.array(plain.weight.view(mx.uint16)) + ), path + assert _linear_input_dims(quantized) == _linear_input_dims(plain) + + x = mx.array(np.random.default_rng(5).normal(size=(2, plain.weight.shape[1])), dtype=mx.float32) + got, want = np.array(quantized(x)), np.array(x @ plain.weight.astype(mx.float32).T) + bound = 1e-5 * (np.abs(np.array(x)) @ np.abs(np.array(plain.weight.astype(mx.float32))).T) + assert np.all(np.abs(got - want) <= bound + 1e-6), path + + # Everything not repacked is bound exactly as the default path binds it. + ref_params = dict(tree_flatten(reference.parameters())) + params = dict(tree_flatten(model.parameters())) + repacked_modules = {name[: -len(".scales")] for name in report.mxfp8_scale_parameters} + def as_numpy(array): + # numpy has no bfloat16; widening to float32 is exact. + if mx.issubdtype(array.dtype, mx.floating): + array = array.astype(mx.float32) + return np.array(array) + + for name, value in ref_params.items(): + if name.rsplit(".", 1)[0] in repacked_modules: + continue + assert np.array_equal(as_numpy(params[name]), as_numpy(value), equal_nan=True), name + + +def test_mxfp8_is_opt_in_and_the_format_is_validated(tmp_path): + args = _tiny_args(num_hidden_layers=1, num_hash_layers=0, compress_ratios=[0]) + _write_synthetic_checkpoint(tmp_path, args) + model = DeepseekV4FlashVQModel(args, with_mtp=False) + with pytest.raises(ValueError, match="resident_format"): + bind_deepseek_v4_flash_non_vq_weights( + model, tmp_path, include_mtp=False, resident_format="fp8" + ) + report = bind_deepseek_v4_flash_non_vq_weights(model, tmp_path, include_mtp=False) + assert report.resident_format == "bf16" + assert not any( + isinstance(module, nn.QuantizedLinear) for _, module in model.named_modules() + ) + + +def test_mxfp8_bind_regroups_wo_a_into_a_quantized_multilinear(tmp_path): + from mlx_lm.models.mla import QuantizedMultiLinear + + # 8 heads x 8 dims over 2 groups gives wo_a a 32-wide per-group input. + args = _tiny_args( + num_hidden_layers=1, num_hash_layers=0, compress_ratios=[0], num_attention_heads=8 + ) + _write_synthetic_checkpoint(tmp_path, args) + reference = DeepseekV4FlashVQModel(args, with_mtp=False) + bind_deepseek_v4_flash_non_vq_weights(reference, tmp_path, include_mtp=False) + model = DeepseekV4FlashVQModel(args, with_mtp=False) + report = bind_deepseek_v4_flash_non_vq_weights( + model, tmp_path, include_mtp=False, strict=True, resident_format="mxfp8" + ) + + assert "layers.0.attn.wo_a.weight" in report.mxfp8_repacked_tensors + assert "layers.0.attn.wo_a.weight" in report.reshaped_tensors + assert report.missing_model_parameters == () + wo_a, plain = model.model.layers[0].attn.wo_a, reference.model.layers[0].attn.wo_a + assert type(wo_a) is QuantizedMultiLinear and wo_a.mode == "mxfp8" + restored = mx.dequantize( + wo_a.weight, wo_a.scales, group_size=32, bits=8, mode="mxfp8", dtype=mx.bfloat16 + ) + assert restored.shape == plain.weight.shape + assert np.array_equal( + np.array(restored.view(mx.uint16)), np.array(plain.weight.view(mx.uint16)) + ) + groups, _, width = plain.weight.shape + x = mx.array(np.random.default_rng(6).normal(size=(1, groups, 3, width)), dtype=mx.float32) + dense = plain.weight.astype(mx.float32) + got, want = np.array(wo_a(x)), np.array(x @ dense.swapaxes(-1, -2)) + bound = 1e-5 * (np.abs(np.array(x)) @ np.abs(np.array(dense)).swapaxes(-1, -2)) + assert np.all(np.abs(got - want) <= bound + 1e-6) diff --git a/tests/test_fp8_block_dequant.py b/tests/test_fp8_block_dequant.py index f765f61..33f8b32 100644 --- a/tests/test_fp8_block_dequant.py +++ b/tests/test_fp8_block_dequant.py @@ -9,6 +9,7 @@ decode_e4m3, decode_ue8m0, dequantize_fp8_block, + repack_fp8_block_as_mxfp8, ) @@ -170,3 +171,144 @@ def test_block_dequant_rejects_non_ndarray_inputs(): dequantize_fp8_block([[0x38]], scales) # type: ignore[arg-type] with pytest.raises(ValueError): dequantize_fp8_block(codes, [[1.0]]) # type: ignore[arg-type] + + +# -------------------------------------------------------------------------- +# Lossless mxfp8 repack (residents kept as their source bytes) +# -------------------------------------------------------------------------- + + +def _finite_fp8_pair(rng, rows, cols, block=(128, 128)): + codes = rng.integers(0, 256, size=(rows, cols)).astype(np.uint8) + codes = np.where((codes & 0x7F) == 0x7F, 0x7E, codes).astype(np.uint8) + blocks = (-(-rows // block[0]), -(-cols // block[1])) + scales = rng.integers(110, 140, size=blocks).astype(np.uint8) + return codes, scales + + +def test_mxfp8_repack_layout_is_a_pure_byte_relabel(): + rng = np.random.default_rng(0) + codes, scales = _finite_fp8_pair(rng, 256, 384) + packed, group_scales = repack_fp8_block_as_mxfp8(codes, scales) + + assert packed.dtype == np.uint32 and packed.shape == (256, 96) + assert group_scales.dtype == np.uint8 and group_scales.shape == (256, 12) + # The code bytes are untouched, four to a little-endian word. + np.testing.assert_array_equal(packed.astype(" Date: Thu, 5 Nov 2026 05:03:23 +0000 Subject: [PATCH 2/8] bench: gate DSV4 headline rows on memory coming back in (row schema v2) The roadmap found the quiet gate one-sided: a row could fault gigabytes of its own working set back in and still stamp itself clean on zero pageouts and swapouts. One post-fix row passed with 245,255 page-ins (3.83 GB). collect_metric_snapshot now also reports pageins, swapins and compressor decompressions (totals and deltas). New headline rows are schema_version 2 and carry an inbound_gate. A v2 row is clean only if its outbound deltas are zero, as before, and it stayed under 4,096 page-ins, 4,096 decompressions (64 MiB each at 16 KiB pages) and zero swap-ins. vm_stat is system-wide, so zero page-ins is not attainable: the cleanest Task 4 row saw 144. The rows that faulted their working set back saw 7,296 to 304,047. A row may not declare a laxer gate than the harness's. Version 1 rows, including the committed Task 4 evidence, replay under their original rule; the quiet record's keys are unchanged. Also: - --resident-format {bf16,mxfp8} on bind-proof, run-row and series binds residents through the new opt-in mxfp8 path. Rows and the bind proof record the format, and a summary refuses rows that mix formats. - --rewarm-after-quiet (opt-in) re-runs the untimed warmup after the quiet window, so the timed window measures steady state instead of pages macOS compressed while the harness sat idle. --- benchmarks/bench_dsv4_mtp_headline.py | 96 +++++++++++-- src/ramp/benchmark/metrics.py | 14 ++ tests/test_bench_dsv4_mtp_headline_inbound.py | 136 ++++++++++++++++++ 3 files changed, 237 insertions(+), 9 deletions(-) create mode 100644 tests/test_bench_dsv4_mtp_headline_inbound.py diff --git a/benchmarks/bench_dsv4_mtp_headline.py b/benchmarks/bench_dsv4_mtp_headline.py index 9ed7d0c..ecb9b07 100644 --- a/benchmarks/bench_dsv4_mtp_headline.py +++ b/benchmarks/bench_dsv4_mtp_headline.py @@ -53,6 +53,17 @@ "GLM_SINGLE_HOST_MLX_WIRED_LIMIT_GB", ) EXPECTED_VERIFY_IMPLEMENTATION = "nax_e8p_m32n64" +#: Rows written by this harness. Version 1 rows (the committed Task 4 +#: evidence) are judged by their original outbound-only rule; version 2 rows +#: must also stay under :data:`INBOUND_GATE`. +ROW_SCHEMA_VERSION = 2 +#: Inbound-paging ceiling for a clean version 2 row, in ``vm_stat`` pages +#: (16 KiB on Apple silicon, so 4,096 pages is 64 MiB). ``vm_stat`` counts are +#: system-wide, so zero is not attainable: the cleanest Task 4 row still saw +#: 144 page-ins. The rows that faulted their working set back in saw 7,296 to +#: 304,047. Pages returning from swap are never acceptable. +INBOUND_GATE = {"max_pageins": 4096, "max_swapins": 0, "max_decompressions": 4096} +RESIDENT_FORMATS = ("bf16", "mxfp8") BENCHMARK_SOURCE_PATH = Path(__file__).resolve() RUNTIME_SOURCE_PATH = (REPO_ROOT / "src/keep/quality/dsv4_mtp_runtime.py").resolve() @@ -199,6 +210,38 @@ def _prompt(pack: Path, prompt_tokens: int) -> tuple[Any, mx.array, dict[str, An def _row_clean(row: Mapping[str, Any]) -> bool: + """Whether a row is admissible evidence under its own schema's rule.""" + + version = row.get("schema_version", 1) + if type(version) is not int or version not in (1, ROW_SCHEMA_VERSION): + raise ValueError(f"unsupported headline row schema_version {version!r}") + outbound_clean = _row_clean_outbound(row) + if version == 1: + return outbound_clean + # Validate the inbound record even when the row is already dirty, so a + # malformed v2 row cannot hide behind an outbound failure. + inbound_clean = _row_clean_inbound(row) + return outbound_clean and inbound_clean + + +def _row_clean_inbound(row: Mapping[str, Any]) -> bool: + gate = row.get("inbound_gate") + if not isinstance(gate, Mapping) or dict(gate) != INBOUND_GATE: + raise ValueError("headline inbound gate is missing or differs from the harness") + deltas = {} + for counter in ("pageins", "swapins", "decompressions"): + value = row.get(f"{counter}_delta") + if type(value) is not int or value < 0: + raise ValueError(f"headline {counter}_delta must be a nonnegative integer") + deltas[counter] = value + return ( + deltas["pageins"] <= INBOUND_GATE["max_pageins"] + and deltas["swapins"] <= INBOUND_GATE["max_swapins"] + and deltas["decompressions"] <= INBOUND_GATE["max_decompressions"] + ) + + +def _row_clean_outbound(row: Mapping[str, Any]) -> bool: quiet = row.get("quiet") if not isinstance(quiet, Mapping): raise TypeError("headline quiet record is missing or malformed") @@ -530,6 +573,11 @@ def _summarize_dsv4_mtp_rows( spec_signatures: set[str] = set() compact_stats: dict[int, dict[str, Any]] = {} prompt_authority = _validate_prompt_pack(rows[0]) + resident_formats = {str(row.get("resident_format", "bf16")) for row in rows} + if len(resident_formats) != 1 or not resident_formats <= set(RESIDENT_FORMATS): + raise ValueError( + f"headline rows mix or misstate resident formats: {sorted(resident_formats)}" + ) for row in rows: if row.get("record_type") != "dsv4_mtp_headline_row_v1": raise ValueError("unexpected DSV4 MTP headline row type") @@ -626,6 +674,7 @@ def _summarize_dsv4_mtp_rows( "clean_pairs": len(clean_pairs), "dirty_pairs": len(pairs) - len(clean_pairs), "min_clean_pairs": min_clean_pairs, + "resident_format": next(iter(resident_formats)), "orders": dict(sorted(orders.items())), "artifact_identity": representative_row["artifact_identity"], "artifact_identity_sha256": identity_sha, @@ -770,7 +819,9 @@ def _load(args: argparse.Namespace): payload_receipt=receipt, expected_payload_receipt_sha256=args.content_receipt_sha256, ) - loaded = load_authenticated_dsv4_composite(inputs) + loaded = load_authenticated_dsv4_composite( + inputs, resident_format=getattr(args, "resident_format", "bf16") + ) mx.eval(loaded.model.parameters()) return loaded @@ -800,6 +851,8 @@ def _bind_proof(args: argparse.Namespace) -> dict[str, Any]: "harness_identity": _current_harness_identity(), "harness_identity_sha256": _canonical_sha256(_current_harness_identity()), "resident_bound_count": non_vq.bound_count, + "resident_format": non_vq.resident_format, + "resident_mxfp8_repacked_tensors": len(non_vq.mxfp8_repacked_tensors), "resident_backbone_bound_count": non_vq.backbone_bound_count, "resident_mtp_bound_count": non_vq.mtp_bound_count, "resident_bound_names_sha256": _canonical_sha256( @@ -856,6 +909,15 @@ def _run_row(args: argparse.Namespace) -> dict[str, Any]: window_seconds=args.quiet_window_seconds, max_attempts=args.quiet_max_attempts, ) + rewarm = bool(getattr(args, "rewarm_after_quiet", False)) + if rewarm: + # The quiet window is idle time in which macOS may compress the + # working set the warmup just heated. Touch it again, untimed, so the + # timed window measures steady state rather than decompression. + if args.role == "baseline": + generate_dsv4_autoregressive(loaded.model, prompt, max_new_tokens=3) + else: + generate_dsv4_speculative(loaded.model, prompt, max_new_tokens=4) thermal = _thermal_snapshot() before_vm = collect_vm_stat_counts() reset_mlx_peak_memory() @@ -877,17 +939,15 @@ def _run_row(args: argparse.Namespace) -> dict[str, Any]: } ) elapsed = time.perf_counter() - started - after_vm = collect_vm_stat_counts() + # One snapshot carries the outbound and inbound deltas for the window. metrics = collect_metric_snapshot(previous_vm_stat_counts=before_vm) - pageins_delta = ( - None - if before_vm is None or after_vm is None - else int(after_vm.get("pageins", 0)) - int(before_vm.get("pageins", 0)) - ) tokens = list(result.token_ids) row = { "record_type": "dsv4_mtp_headline_row_v1", - "schema_version": 1, + "schema_version": ROW_SCHEMA_VERSION, + "inbound_gate": dict(INBOUND_GATE), + "resident_format": getattr(args, "resident_format", "bf16"), + "rewarm_after_quiet": rewarm, "fresh_process": True, "pid": os.getpid(), "ppid": os.getppid(), @@ -911,7 +971,6 @@ def _run_row(args: argparse.Namespace) -> dict[str, Any]: "speculative_stats": spec_stats, "quiet": quiet, "thermal": thermal, - "pageins_delta": pageins_delta, "wired_limit_environment_absent": True, **metrics, } @@ -980,6 +1039,9 @@ def _series(args: argparse.Namespace) -> None: str(args.quiet_window_seconds), "--quiet-max-attempts", str(args.quiet_max_attempts), + "--resident-format", + args.resident_format, + *(["--rewarm-after-quiet"] if args.rewarm_after_quiet else []), ] env = dict(os.environ) env["KEEP_TASK4_LOCK_HELD"] = "1" @@ -1031,6 +1093,20 @@ def _common_artifact_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE) parser.add_argument("--resident-dir", type=Path, default=DEFAULT_RESIDENT) parser.add_argument("--vq-dir", type=Path, default=DEFAULT_VQ) + parser.add_argument( + "--resident-format", + choices=RESIDENT_FORMATS, + default="bf16", + help="bind FP8 residents decoded to bf16 (default) or as their mxfp8 source bytes", + ) + + +def _rewarm_arg(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--rewarm-after-quiet", + action="store_true", + help="re-run the untimed warmup after the quiet window, before timing", + ) def main() -> None: @@ -1055,6 +1131,7 @@ def main() -> None: row.add_argument("--role", choices=("baseline", "speculative"), required=True) row.add_argument("--quiet-window-seconds", type=float, default=5.0) row.add_argument("--quiet-max-attempts", type=int, default=3) + _rewarm_arg(row) row.add_argument("--output", type=Path, required=True) row.set_defaults(func=_run_row, heavy=True) @@ -1085,6 +1162,7 @@ def main() -> None: series.add_argument("--min-clean-pairs", type=int, default=3) series.add_argument("--quiet-window-seconds", type=float, default=5.0) series.add_argument("--quiet-max-attempts", type=int, default=3) + _rewarm_arg(series) series.add_argument("--output-dir", type=Path, required=True) series.set_defaults(func=_series, heavy=True) diff --git a/src/ramp/benchmark/metrics.py b/src/ramp/benchmark/metrics.py index ea13ab7..6f79181 100644 --- a/src/ramp/benchmark/metrics.py +++ b/src/ramp/benchmark/metrics.py @@ -13,6 +13,12 @@ MetricSnapshot = dict[str, int | None] VMStatCounts = dict[str, int] +#: ``vm_stat`` counters for memory coming *back*: file/anonymous page-ins, +#: swap-ins, and pages pulled out of the memory compressor. A row can stay at +#: zero pageouts/swapouts while faulting gigabytes of its own working set back +#: in through these, so a quiet gate that ignores them is one-sided. +INBOUND_VM_COUNTERS = ("pageins", "swapins", "decompressions") + _VM_STAT_RE = re.compile(r"^\s*(?P[^:]+):\s+(?P[0-9]+)\.?\s*$") @@ -189,4 +195,12 @@ def collect_metric_snapshot( "swapouts_total": None if vm_counts is None else vm_counts.get("swapouts"), "pageouts_delta": _delta(vm_counts, previous_vm_stat_counts, "pageouts"), "swapouts_delta": _delta(vm_counts, previous_vm_stat_counts, "swapouts"), + **{ + f"{counter}_total": None if vm_counts is None else vm_counts.get(counter) + for counter in INBOUND_VM_COUNTERS + }, + **{ + f"{counter}_delta": _delta(vm_counts, previous_vm_stat_counts, counter) + for counter in INBOUND_VM_COUNTERS + }, } diff --git a/tests/test_bench_dsv4_mtp_headline_inbound.py b/tests/test_bench_dsv4_mtp_headline_inbound.py new file mode 100644 index 0000000..07c8f1a --- /dev/null +++ b/tests/test_bench_dsv4_mtp_headline_inbound.py @@ -0,0 +1,136 @@ +"""Version 2 headline rows also gate on memory coming back in. + +The roadmap found the quiet gate one-sided: a row could fault gigabytes of its +working set back in (page-ins, swap-ins, compressor decompressions) and still +stamp itself clean on zero pageouts and swapouts. Version 2 rows close that; +version 1 rows (the committed Task 4 evidence) keep their original rule. +""" + +from __future__ import annotations + +import pytest + +import ramp.benchmark.metrics as metrics +from benchmarks import bench_dsv4_mtp_headline as benchmark + + +def _quiet(): + return { + "enabled": True, + "available": True, + "quiet": True, + "attempts": 1, + "window_seconds": 5.0, + "pageouts_delta": 0, + "swapouts_delta": 0, + } + + +def _v2_row(**changes): + row = { + "schema_version": 2, + "quiet": _quiet(), + "pageouts_delta": 0, + "swapouts_delta": 0, + "inbound_gate": dict(benchmark.INBOUND_GATE), + "pageins_delta": 144, + "swapins_delta": 0, + "decompressions_delta": 12, + } + row.update(changes) + return row + + +def test_v2_row_under_the_inbound_ceiling_is_clean(): + assert benchmark._row_clean(_v2_row()) is True + + +@pytest.mark.parametrize( + "changes", + [ + # 245,255 page-ins (3.83 GB) passed the old gate in the post-fix series. + {"pageins_delta": 245_255}, + {"pageins_delta": benchmark.INBOUND_GATE["max_pageins"] + 1}, + {"swapins_delta": 1}, + {"decompressions_delta": benchmark.INBOUND_GATE["max_decompressions"] + 1}, + ], +) +def test_v2_row_that_faulted_memory_back_in_is_dirty(changes): + assert benchmark._row_clean(_v2_row(**changes)) is False + + +def test_the_same_inbound_traffic_does_not_change_a_v1_verdict(): + """Committed version 1 evidence replays under the rule it was gated by.""" + + row = _v2_row(schema_version=1, pageins_delta=245_255) + del row["inbound_gate"], row["swapins_delta"], row["decompressions_delta"] + assert benchmark._row_clean(row) is True + del row["schema_version"] # rows that predate the field are version 1 + assert benchmark._row_clean(row) is True + + +def test_v2_outbound_failure_still_dirties_the_row(): + assert benchmark._row_clean(_v2_row(pageouts_delta=1)) is False + + +@pytest.mark.parametrize( + "changes", + [ + {"pageins_delta": None}, + {"swapins_delta": -1}, + {"decompressions_delta": 1.0}, + {"inbound_gate": None}, + # A row may not declare a laxer ceiling than the harness enforces. + {"inbound_gate": {**benchmark.INBOUND_GATE, "max_pageins": 10**9}}, + # Malformed inbound fields are caught even when outbound already failed. + {"pageouts_delta": 3, "swapins_delta": None}, + ], +) +def test_v2_inbound_record_is_validated_not_trusted(changes): + with pytest.raises(ValueError, match="inbound|_delta"): + benchmark._row_clean(_v2_row(**changes)) + + +@pytest.mark.parametrize("version", [0, 3, "2", True]) +def test_unknown_row_schema_versions_are_refused(version): + with pytest.raises(ValueError, match="schema_version"): + benchmark._row_clean(_v2_row(schema_version=version)) + + +def test_metric_snapshot_reports_inbound_counters(monkeypatch): + monkeypatch.setattr( + metrics, + "collect_vm_stat_counts", + lambda: { + "pageouts": 5, + "swapouts": 6, + "pageins": 1_000, + "swapins": 7, + "decompressions": 50, + }, + ) + monkeypatch.setattr(metrics, "_mlx_memory_value", lambda _name: 0) + snapshot = metrics.collect_metric_snapshot( + previous_vm_stat_counts={ + "pageouts": 5, + "swapouts": 6, + "pageins": 856, + "swapins": 7, + "decompressions": 38, + } + ) + assert snapshot["pageins_delta"] == 144 + assert snapshot["swapins_delta"] == 0 + assert snapshot["decompressions_delta"] == 12 + assert snapshot["pageins_total"] == 1_000 + assert set(metrics.INBOUND_VM_COUNTERS) == {"pageins", "swapins", "decompressions"} + + +def test_vm_stat_parser_reads_the_inbound_lines(): + counts = metrics.parse_vm_stat_counts( + "Mach Virtual Memory Statistics: (page size of 16384 bytes)\n" + "Pageins: 123.\n" + "Swapins: 4.\n" + "Decompressions: 56.\n" + ) + assert counts == {"pageins": 123, "swapins": 4, "decompressions": 56} From bd7fd54d2aaff24b03210f9a81be28672c1d4545 Mon Sep 17 00:00:00 2001 From: Jack Mazac <114446521+jackmaxwil@users.noreply.github.com> Date: Thu, 5 Nov 2026 05:10:39 +0000 Subject: [PATCH 3/8] keep: REAP expert saliency from the real gate, collected in the calibration pass keep.quality.reap scores each routed expert by REAP's criterion, the mean over the slots routed to it of gate weight times the L2 norm of its output before gating. The inputs come from what KEEP's teacher runner already computes: the model's own gate (bias-corrected selection, sqrtsoftplus weights, renormalisation) and the per-route expert outputs. Upstream REAP's code instead takes softmax of raw logits, picks top-k without the correction bias, and runs every expert on every token. Here the cost is a norm and five scatter-adds per layer, folded into float64 once per layer like the imatrix. - layer_major_prefill takes an optional ReapAccumulator; the forward stays bit-identical (tested alone and alongside the imatrix). - Calibration sessions now also write reap__* arrays. Capture validation checks them and requires their route counts to match the imatrix's. The imatrix finalizer ignores them, so existing runs still finalize. - finalize_dsv4_reap_stats sums them, and refuses if any session lacks them rather than silently under-scoring what that session routed to. - produce_dsv4_teacher_cache.py finalize-reap writes reap-stats.npz. reap-select picks a uniform survivor set and reports per-layer Jaccard of the pruned set against route frequency. That is the plan's kill test: Jaccard >= 0.9 means REAP adds nothing over the counts KEEP already had. - expert_saliency offers reap, frequency, weighted_frequency and ean_mean, plus optional shrinkage of rare experts toward the layer mean. - super_expert_candidates applies REAP's outlier rule to max |f|, so negative outliers are not invisible as they are upstream. - prune_checkpoint_plan maps checkpoint names to a uniform prune: it renumbers kept experts and slices gate weight, bias and bias_vl (V4.1's image-token bias, which must not be left at the old width). It refuses hash-routed blocks (tid2eid), which need a remapping policy. A test checks that the sliced gate picks the same experts with the same renormalised weights as the full gate with pruned experts masked out of selection. --- benchmarks/produce_dsv4_teacher_cache.py | 87 +++++ src/keep/quality/dsv4_teacher_runner.py | 95 ++++- src/keep/quality/reap.py | 475 +++++++++++++++++++++++ tests/test_dsv4_teacher_runner.py | 117 ++++++ tests/test_reap.py | 333 ++++++++++++++++ 5 files changed, 1105 insertions(+), 2 deletions(-) create mode 100644 src/keep/quality/reap.py create mode 100644 tests/test_reap.py diff --git a/benchmarks/produce_dsv4_teacher_cache.py b/benchmarks/produce_dsv4_teacher_cache.py index 53e718e..75549b6 100644 --- a/benchmarks/produce_dsv4_teacher_cache.py +++ b/benchmarks/produce_dsv4_teacher_cache.py @@ -237,6 +237,67 @@ def _cmd_finalize(opts: argparse.Namespace) -> int: return 0 +def _cmd_finalize_reap(opts: argparse.Namespace) -> int: + from keep.quality.dsv4_teacher_runner import finalize_dsv4_reap_stats + + report = finalize_dsv4_reap_stats( + Path(opts.out_dir).expanduser(), + output=Path(opts.output).expanduser() if opts.output else None, + ) + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + + +def _cmd_reap_select(opts: argparse.Namespace) -> int: + """Pick survivors and run the REAP-vs-baseline kill test on the same stats.""" + + import numpy as np + + from keep.quality.reap import ( + ReapStats, + expert_saliency, + pruned_set_jaccard, + select_survivors, + super_expert_candidates, + write_survivors, + ) + + stats = ReapStats.load(Path(opts.stats).expanduser()) + chosen = select_survivors( + expert_saliency(stats, opts.criterion, prior_slots=opts.prior_slots), opts.keep + ) + baseline = select_survivors(expert_saliency(stats, opts.compare), opts.keep) + jaccard = pruned_set_jaccard(chosen, baseline, num_experts=stats.num_experts) + report = { + "record_type": "dsv4_reap_select", + "criterion": opts.criterion, + "prior_slots": opts.prior_slots, + "keep": opts.keep, + "num_experts": stats.num_experts, + "sessions": len(stats.sessions), + "compare": opts.compare, + "pruned_set_jaccard_vs_compare": { + str(int(layer)): round(float(value), 4) + for layer, value in zip(stats.layers, jaccard) + }, + "median_jaccard": float(np.median(jaccard)), + # Kill test from the plan: if REAP's pruned set matches frequency's + # at Jaccard >= 0.9, REAP adds nothing over the counts KEEP already has. + "reap_adds_little": bool(np.median(jaccard) >= 0.9), + "never_routed_experts": int((stats.route_count == 0).sum()), + "super_expert_candidates": super_expert_candidates(stats), + } + if opts.output: + write_survivors( + Path(opts.output).expanduser(), + chosen, + layers=stats.layers.tolist(), + meta={k: v for k, v in report.items() if k != "record_type"}, + ) + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + + def _cmd_plan(opts: argparse.Namespace) -> int: from keep.quality.dsv4_teacher_runner import ( build_dsv4_expert_span_index, @@ -347,6 +408,32 @@ def build_parser() -> argparse.ArgumentParser: finalize.add_argument("--prompt-set", default="dsv4_teich_calibration_v1") finalize.set_defaults(func=_cmd_finalize) + finalize_reap = verbs.add_parser( + "finalize-reap", help="sum calibration sessions' REAP expert statistics" + ) + finalize_reap.add_argument("--out-dir", required=True) + finalize_reap.add_argument("--output", help="default: /reap-stats.npz") + finalize_reap.set_defaults(func=_cmd_finalize_reap) + + select = verbs.add_parser( + "reap-select", help="choose surviving experts and compare criteria" + ) + select.add_argument("--stats", required=True, help="reap-stats.npz from finalize-reap") + select.add_argument("--keep", type=int, required=True, help="experts kept per layer") + select.add_argument( + "--criterion", + default="reap", + choices=("reap", "frequency", "weighted_frequency", "ean_mean"), + ) + select.add_argument( + "--compare", + default="frequency", + choices=("reap", "frequency", "weighted_frequency", "ean_mean"), + ) + select.add_argument("--prior-slots", type=float, default=0.0) + select.add_argument("--output", help="write the survivor table as JSON") + select.set_defaults(func=_cmd_reap_select) + plan = verbs.add_parser("plan", help="print the streaming plan and byte budget") plan.add_argument("--checkpoint", default=str(DEFAULT_CHECKPOINT)) plan.add_argument("--pack", default=str(DEFAULT_PACK)) diff --git a/src/keep/quality/dsv4_teacher_runner.py b/src/keep/quality/dsv4_teacher_runner.py index 9979193..aa12438 100644 --- a/src/keep/quality/dsv4_teacher_runner.py +++ b/src/keep/quality/dsv4_teacher_runner.py @@ -92,6 +92,7 @@ "build_dsv4_expert_span_index", "calibration_session_arrays", "finalize_dsv4_calibration_imatrix", + "finalize_dsv4_reap_stats", "layer_major_prefill", "load_dsv4_streaming_workload", "load_dsv4_teich_pack", @@ -1162,6 +1163,7 @@ def layer_major_prefill( on_layer: Callable[[int, int], None] | None = None, release_after_layer: bool = True, dspark_taps: dict[int, list[Any]] | None = None, + reap: Any | None = None, ) -> list[Any]: """Prefill the whole backbone layer-outer, chunk-inner. @@ -1182,6 +1184,12 @@ def layer_major_prefill( :meth:`DeepseekV4FlashBackbone.__call__`. Layer-major order is why this has to be a dict of lists rather than one concatenated array: the taps for layer 40 are complete before layer 41 has run at all. + + ``reap`` is an optional :class:`keep.quality.reap.ReapAccumulator`. It + takes the model's own gate weights and each routed expert's output before + gating, so it scores experts against the routing the model actually uses. + Like the imatrix it reads tensors the forward already computes and leaves + the forward bit-identical. """ import mlx.core as mx @@ -1284,7 +1292,7 @@ def down_hook(activated, expert_ids, order, _capture=capture): ) mask = masks[chunk_index] - if accumulator is None: + if accumulator is None and reap is None: hidden[chunk_index] = layer(h, mask, layer_cache, piece) if layer_index in tap_layers: tap = _dspark_tap(hidden[chunk_index]) @@ -1316,6 +1324,17 @@ def down_hook(activated, expert_ids, order, _capture=capture): flat_inds = inds.reshape(-1, inds.shape[-1]) flat_scores = scores.reshape(-1, scores.shape[-1]).astype(mx.float32) + if reap is not None: + reap.add_chunk( + layer=layer_index, + expert_ids=flat_inds, + gate_weights=flat_scores, + expert_outputs=routed.reshape(-1, *routed.shape[-2:]), + ) + if accumulator is None: + mx.eval(hidden[chunk_index]) + reap.eval_pending() + continue if collect_hidden: rows = moe_input.reshape(-1, moe_input.shape[-1]).astype(mx.float32) accumulator.add_chunk( @@ -1344,6 +1363,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): ) mx.eval(hidden[chunk_index]) accumulator.eval_pending() + if reap is not None: + reap.eval_pending() switch_mlp.down_hook = None layer.ffn.switch_mlp = None @@ -1352,6 +1373,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): expert_provider.release(layer_index) if accumulator is not None: accumulator.flush_device() + if reap is not None: + reap.flush_device() mx.clear_cache() if on_layer is not None: on_layer(layer_index, n_chunks) @@ -1933,6 +1956,19 @@ def validate_dsv4_calibration_capture( raise ValueError( f"route_count__{space} rows must sum to total_route_count__{space}" ) + if any(str(name).startswith("reap__") for name in capture): + from keep.quality.reap import ReapStats + + reap = ReapStats.from_arrays(capture) # validates shapes and invariants + if reap.layers.tolist() != stored_layers.tolist(): + raise ValueError("calibration REAP layer set drifted") + if reap.num_experts != int(num_experts): + raise ValueError("calibration REAP expert count drifted") + hidden_counts = capture.get("route_count__hidden") + if hidden_counts is not None and not np.array_equal( + np.asarray(hidden_counts), reap.route_count + ): + raise ValueError("REAP route counts disagree with the imatrix route counts") def calibration_session_arrays( @@ -2665,8 +2701,12 @@ def _produce_sessions( started = time.perf_counter() token_ids = mx.array(session.token_ids.astype(np.int64)[None, :]) accumulator = None + reap = None collect: tuple[str, ...] = () if mode == "calibration": + from keep.quality.reap import ReapAccumulator + + reap = ReapAccumulator(layers=layers, num_experts=num_experts) accumulator = _ImatrixAccumulator( layers=layers, num_experts=num_experts, @@ -2696,10 +2736,14 @@ def on_layer(layer_index: int, n_chunks: int, _s=session, _i=index) -> None: collect_spaces=collect, on_layer=on_layer, dspark_taps=taps, + reap=reap, ) if mode == "calibration": - arrays = calibration_session_arrays(accumulator) + arrays = { + **calibration_session_arrays(accumulator), + **reap.stats(session=session.prompt_id).as_arrays(), + } validate_dsv4_calibration_capture( arrays, layers=layers, num_experts=num_experts ) @@ -2797,6 +2841,53 @@ def _atomic_json(path: Path, payload: Any) -> None: # --------------------------------------------------------------------------- +def finalize_dsv4_reap_stats( + out_dir: str | Path, *, output: str | Path | None = None +) -> dict[str, Any]: + """Sum every calibration session's REAP statistics into one file. + + Fails closed if any calibration session lacks them (for example a session + produced before REAP collection existed), because a saliency missing some + sessions would silently under-score whatever those sessions routed to. + """ + + from keep.quality.reap import REAP_ARRAY_PREFIX, ReapStats + + root = Path(out_dir) + session_files = sorted((root / "sessions").glob("*.npz")) + if not session_files: + raise ValueError(f"no session outputs under {root / 'sessions'}") + total: ReapStats | None = None + missing: list[str] = [] + for path in session_files: + with np.load(path, allow_pickle=False) as archive: + capture = {name: archive[name] for name in archive.files} + if _identity_text(capture, "record_type") != "dsv4_teacher_calibration_v1": + raise ValueError(f"{path} is not a calibration capture") + if not any(name.startswith(REAP_ARRAY_PREFIX) for name in capture): + missing.append(path.name) + continue + stats = ReapStats.from_arrays(capture, sessions=(_identity_text(capture, "prompt_id"),)) + total = stats if total is None else total.merged(stats) + if missing: + raise ValueError( + f"{len(missing)} calibration sessions have no REAP statistics " + f"(rerun them): {missing[:5]}" + ) + assert total is not None + target = Path(output) if output is not None else root / "reap-stats.npz" + total.save(target) + return { + "record_type": "dsv4_reap_stats_finalize", + "session_count": len(session_files), + "layers": total.layers.tolist(), + "num_experts": total.num_experts, + "routed_slots": int(total.routed_slots.sum()), + "never_routed_experts": int((total.route_count == 0).sum()), + "output": str(target), + } + + def finalize_dsv4_calibration_imatrix( out_dir: str | Path, *, diff --git a/src/keep/quality/reap.py b/src/keep/quality/reap.py new file mode 100644 index 0000000..8e7f0d0 --- /dev/null +++ b/src/keep/quality/reap.py @@ -0,0 +1,475 @@ +"""REAP expert saliency from the real routed forward, and the prune it drives. + +REAP (Router-weighted Expert Activation Pruning) scores expert ``i`` by the +mean, over the tokens actually routed to it, of ``g * ||f_i(x)||_2``: the gate +weight that combined its output times the L2 norm of that output before +gating. Experts with the lowest score are dropped and the router re-selects +top-k among the survivors. + +This module computes that score from what KEEP's teacher runner already has +in hand -- the model's own gate (bias-corrected selection, the family's own +weight function and renormalisation) and the per-route expert outputs -- so it +costs a norm and a few scatter-adds per layer instead of REAP's dense pass, +which runs every expert on every token. Two differences from the upstream +implementation are deliberate: + +- **The real gate.** Upstream computes ``softmax`` of raw router logits and + picks top-k without the correction bias. For a sqrtsoftplus/noaux router + (DeepSeek-V4, V4.1) that is a different routing distribution; the score + here uses the weights the model actually used. +- **Routed tokens only.** Saliency never needs an unrouted expert's output. + +The statistics also support the baselines a kill test compares against +(route frequency, gate-weighted frequency, mean output norm) and the +"super-expert" outlier screen, which here uses ``max |f|`` so negative +outliers are not invisible. + +Pruning is uniform: one surviving expert count for every layer, because RAMP +binds a single ``n_routed_experts``. +""" + +from __future__ import annotations + +import json +import re +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +import numpy as np + +__all__ = [ + "REAP_ARRAY_PREFIX", + "ReapAccumulator", + "ReapStats", + "SALIENCY_CRITERIA", + "expert_saliency", + "prune_checkpoint_plan", + "pruned_set_jaccard", + "select_survivors", + "super_expert_candidates", + "write_survivors", +] + +REAP_ARRAY_PREFIX = "reap__" +SALIENCY_CRITERIA = ("reap", "frequency", "weighted_frequency", "ean_mean") +Criterion = Literal["reap", "frequency", "weighted_frequency", "ean_mean"] + +_STAT_FIELDS = ( + "route_count", # int64 [L, E]: routed slots per expert + "gate_sum", # float64 [L, E]: sum of gate weights + "gate_norm_sum", # float64 [L, E]: sum of g * ||f|| (the REAP numerator) + "norm_sum", # float64 [L, E]: sum of ||f|| + "norm_max", # float64 [L, E]: max ||f|| over routed slots + "abs_max", # float64 [L, E]: max over routed slots of max_h |f_h| +) + + +@dataclass +class ReapStats: + """Summed REAP statistics for one or more sessions.""" + + layers: np.ndarray # int64 [L] + num_experts: int + route_count: np.ndarray + gate_sum: np.ndarray + gate_norm_sum: np.ndarray + norm_sum: np.ndarray + norm_max: np.ndarray + abs_max: np.ndarray + routed_slots: np.ndarray # int64 [L]: total routed slots per layer + sessions: tuple[str, ...] = field(default_factory=tuple) + + def __post_init__(self) -> None: + self.layers = np.asarray(self.layers, dtype=np.int64) + n = self.layers.size + shape = (n, int(self.num_experts)) + self.route_count = np.asarray(self.route_count, dtype=np.int64) + for name in _STAT_FIELDS[1:]: + setattr(self, name, np.asarray(getattr(self, name), dtype=np.float64)) + self.routed_slots = np.asarray(self.routed_slots, dtype=np.int64) + for name in _STAT_FIELDS: + value = getattr(self, name) + if value.shape != shape: + raise ValueError(f"REAP {name} has shape {value.shape}, expected {shape}") + if not np.all(np.isfinite(value)) or np.any(value < 0): + raise ValueError(f"REAP {name} must be finite and non-negative") + if self.routed_slots.shape != (n,): + raise ValueError("REAP routed_slots must be one count per layer") + if np.any(self.route_count.sum(axis=1) != self.routed_slots): + raise ValueError("REAP route counts do not sum to the routed slots") + if np.any((self.route_count == 0) & (self.gate_norm_sum != 0)): + raise ValueError("REAP has gate mass on an expert with no routed slots") + + # -- persistence ------------------------------------------------------- + + def as_arrays(self) -> dict[str, np.ndarray]: + arrays = { + f"{REAP_ARRAY_PREFIX}layers": self.layers, + f"{REAP_ARRAY_PREFIX}num_experts": np.asarray(self.num_experts, dtype=np.int64), + f"{REAP_ARRAY_PREFIX}routed_slots": self.routed_slots, + } + for name in _STAT_FIELDS: + arrays[f"{REAP_ARRAY_PREFIX}{name}"] = getattr(self, name) + return arrays + + @classmethod + def from_arrays( + cls, arrays: Mapping[str, Any], *, sessions: Sequence[str] = () + ) -> "ReapStats": + key = REAP_ARRAY_PREFIX.__add__ + missing = [ + name + for name in ("layers", "num_experts", "routed_slots", *_STAT_FIELDS) + if key(name) not in arrays + ] + if missing: + raise KeyError("REAP arrays missing: " + ", ".join(missing)) + return cls( + layers=np.asarray(arrays[key("layers")]), + num_experts=int(np.asarray(arrays[key("num_experts")]).reshape(())), + routed_slots=np.asarray(arrays[key("routed_slots")]), + sessions=tuple(sessions), + **{name: np.asarray(arrays[key(name)]) for name in _STAT_FIELDS}, + ) + + def merged(self, other: "ReapStats") -> "ReapStats": + """Sum two disjoint sets of sessions over the same layers and experts.""" + + if not np.array_equal(self.layers, other.layers) or self.num_experts != other.num_experts: + raise ValueError("REAP stats cover different layers or expert counts") + overlap = set(self.sessions) & set(other.sessions) + if overlap: + raise ValueError(f"REAP stats would double-count sessions: {sorted(overlap)[:5]}") + return ReapStats( + layers=self.layers, + num_experts=self.num_experts, + route_count=self.route_count + other.route_count, + gate_sum=self.gate_sum + other.gate_sum, + gate_norm_sum=self.gate_norm_sum + other.gate_norm_sum, + norm_sum=self.norm_sum + other.norm_sum, + norm_max=np.maximum(self.norm_max, other.norm_max), + abs_max=np.maximum(self.abs_max, other.abs_max), + routed_slots=self.routed_slots + other.routed_slots, + sessions=self.sessions + other.sessions, + ) + + def save(self, path: str | Path) -> None: + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + tmp = target.with_name(target.name + ".tmp.npz") + np.savez(tmp, **self.as_arrays(), **{f"{REAP_ARRAY_PREFIX}sessions": np.asarray(self.sessions, dtype=str)}) + tmp.replace(target) + + @classmethod + def load(cls, path: str | Path) -> "ReapStats": + with np.load(path, allow_pickle=False) as archive: + arrays = {name: archive[name] for name in archive.files} + sessions = tuple(str(s) for s in arrays.pop(f"{REAP_ARRAY_PREFIX}sessions", np.asarray([], dtype=str))) + return cls.from_arrays(arrays, sessions=sessions) + + +class ReapAccumulator: + """Device scatter-add of REAP statistics; folds into float64 once per layer. + + Same design as the teacher runner's imatrix accumulator: float32 sums on + device for one layer's chunks, then a float64 host fold, so cross-chunk and + cross-session accumulation carries no float32 drift beyond one layer. + """ + + def __init__(self, *, layers: Sequence[int], num_experts: int) -> None: + self.layers = tuple(int(layer) for layer in layers) + self.num_experts = int(num_experts) + self._row = {layer: row for row, layer in enumerate(self.layers)} + shape = (len(self.layers), self.num_experts) + self._host = {name: np.zeros(shape, dtype=np.float64) for name in _STAT_FIELDS} + self._routed_slots = np.zeros(len(self.layers), dtype=np.int64) + self._device: dict[int, dict[str, Any]] = {} + + def add_chunk( + self, *, layer: int, expert_ids: Any, gate_weights: Any, expert_outputs: Any + ) -> None: + """Add one chunk's routed slots. + + ``expert_ids`` and ``gate_weights`` are ``[rows, top_k]``; + ``expert_outputs`` is ``[rows, top_k, hidden]``, each expert's output + for that slot *before* the gate weight multiplies it. + """ + + import mlx.core as mx + + if layer not in self._row: + raise KeyError(f"layer {layer} is not tracked by this accumulator") + if expert_outputs.ndim != 3 or tuple(expert_outputs.shape[:2]) != tuple(expert_ids.shape): + raise ValueError( + f"expert outputs {tuple(expert_outputs.shape)} do not match routes " + f"{tuple(expert_ids.shape)}" + ) + outputs = expert_outputs.astype(mx.float32) + norms = mx.sqrt(mx.sum(outputs * outputs, axis=-1)).reshape(-1) + abs_max = mx.max(mx.abs(outputs), axis=-1).reshape(-1) + ids = expert_ids.reshape(-1).astype(mx.int32) + gates = gate_weights.reshape(-1).astype(mx.float32) + bucket = self._device.get(layer) + if bucket is None: + zeros = lambda: mx.zeros((self.num_experts,), dtype=mx.float32) # noqa: E731 + bucket = {name: zeros() for name in _STAT_FIELDS} + bucket["slots"] = 0 + self._device[layer] = bucket + bucket["route_count"] = bucket["route_count"].at[ids].add(mx.ones_like(norms)) + bucket["gate_sum"] = bucket["gate_sum"].at[ids].add(gates) + bucket["gate_norm_sum"] = bucket["gate_norm_sum"].at[ids].add(gates * norms) + bucket["norm_sum"] = bucket["norm_sum"].at[ids].add(norms) + bucket["norm_max"] = bucket["norm_max"].at[ids].maximum(norms) + bucket["abs_max"] = bucket["abs_max"].at[ids].maximum(abs_max) + bucket["slots"] += int(ids.size) + + def eval_pending(self) -> None: + import mlx.core as mx + + for bucket in self._device.values(): + mx.eval([bucket[name] for name in _STAT_FIELDS]) + + def flush_device(self) -> None: + import mlx.core as mx + + for layer, bucket in self._device.items(): + mx.eval([bucket[name] for name in _STAT_FIELDS]) + row = self._row[layer] + for name in ("gate_sum", "gate_norm_sum", "norm_sum"): + self._host[name][row] += np.asarray(bucket[name], dtype=np.float64) + # float32 scatter-added ones are exact small integers; rint makes a + # future non-integral weighting fail loudly at the int cast below. + self._host["route_count"][row] += np.rint(np.asarray(bucket["route_count"], dtype=np.float64)) + for name in ("norm_max", "abs_max"): + self._host[name][row] = np.maximum( + self._host[name][row], np.asarray(bucket[name], dtype=np.float64) + ) + self._routed_slots[row] += int(bucket["slots"]) + self._device.clear() + + def stats(self, *, session: str | None = None) -> ReapStats: + self.flush_device() + return ReapStats( + layers=np.asarray(self.layers, dtype=np.int64), + num_experts=self.num_experts, + route_count=self._host["route_count"].astype(np.int64), + gate_sum=self._host["gate_sum"], + gate_norm_sum=self._host["gate_norm_sum"], + norm_sum=self._host["norm_sum"], + norm_max=self._host["norm_max"], + abs_max=self._host["abs_max"], + routed_slots=self._routed_slots.copy(), + sessions=() if session is None else (session,), + ) + + def as_arrays(self) -> dict[str, np.ndarray]: + return self.stats().as_arrays() + + +def expert_saliency( + stats: ReapStats, criterion: Criterion = "reap", *, prior_slots: float = 0.0 +) -> np.ndarray: + """Per-(layer, expert) saliency under ``criterion``; higher means keep. + + - ``reap``: mean of ``g * ||f||`` over routed slots (REAP's score). + - ``frequency``: routed slot count. + - ``weighted_frequency``: sum of gate weights. + - ``ean_mean``: mean output norm over routed slots, ignoring the gate. + + An expert with no routed slots scores 0 under every criterion, so it is + pruned first, as upstream does. ``prior_slots > 0`` shrinks the mean-type + scores toward the layer mean as if each expert had seen that many extra + average slots, which damps the noise of rarely routed experts; 0 (the + default) is REAP exactly. + """ + + if criterion not in SALIENCY_CRITERIA: + raise ValueError(f"criterion must be one of {SALIENCY_CRITERIA}, got {criterion!r}") + if prior_slots < 0: + raise ValueError("prior_slots must be non-negative") + count = stats.route_count.astype(np.float64) + if criterion == "frequency": + return count + if criterion == "weighted_frequency": + return stats.gate_sum.copy() + numerator = stats.gate_norm_sum if criterion == "reap" else stats.norm_sum + with np.errstate(invalid="ignore", divide="ignore"): + if prior_slots == 0: + return np.where(count > 0, numerator / np.maximum(count, 1.0), 0.0) + layer_mean = numerator.sum(axis=1, keepdims=True) / np.maximum( + count.sum(axis=1, keepdims=True), 1.0 + ) + shrunk = (numerator + prior_slots * layer_mean) / (count + prior_slots) + return np.where(count > 0, shrunk, 0.0) + + +def select_survivors(saliency: np.ndarray, keep: int) -> np.ndarray: + """Keep the ``keep`` most salient experts in every layer. + + Returns ``int64 [L, keep]`` expert ids, ascending within each layer, which + is the order the sliced router rows and renumbered experts follow. Ties + break toward the lower expert id, so the result is deterministic. + """ + + saliency = np.asarray(saliency, dtype=np.float64) + if saliency.ndim != 2: + raise ValueError("saliency must be [layers, experts]") + layers, experts = saliency.shape + if not 0 < int(keep) <= experts: + raise ValueError(f"keep must be in 1..{experts}, got {keep}") + if not np.all(np.isfinite(saliency)): + raise ValueError("saliency must be finite") + survivors = np.empty((layers, int(keep)), dtype=np.int64) + ids = np.arange(experts) + for row in range(layers): + order = np.lexsort((ids, -saliency[row])) # by saliency desc, then id asc + survivors[row] = np.sort(order[: int(keep)]) + return survivors + + +def pruned_set_jaccard(a: np.ndarray, b: np.ndarray, *, num_experts: int) -> np.ndarray: + """Per-layer Jaccard index of the *pruned* sets implied by two survivor lists.""" + + a = np.asarray(a) + b = np.asarray(b) + if a.shape != b.shape or a.ndim != 2: + raise ValueError("survivor arrays must share a [layers, keep] shape") + out = np.empty(a.shape[0], dtype=np.float64) + universe = set(range(int(num_experts))) + for row in range(a.shape[0]): + pruned_a = universe - set(a[row].tolist()) + pruned_b = universe - set(b[row].tolist()) + union = pruned_a | pruned_b + out[row] = 1.0 if not union else len(pruned_a & pruned_b) / len(union) + return out + + +def super_expert_candidates( + stats: ReapStats, + *, + percentile: float = 99.5, + fraction_of_max: float = 0.1, + layer_fraction: float = 0.75, +) -> list[tuple[int, int]]: + """``(layer, expert)`` pairs whose largest routed output is an outlier. + + REAP's rule on ``max |f|`` rather than the signed maximum: an expert is a + candidate if its peak exceeds both the ``percentile`` of all peaks and + ``fraction_of_max`` of the global peak, within the first ``layer_fraction`` + of layers. Upstream starts its signed maximum at 0, so negative outliers + never register; this screen sees them. + """ + + peaks = stats.abs_max + if peaks.size == 0: + return [] + threshold = max(float(np.percentile(peaks, percentile)), float(peaks.max()) * fraction_of_max) + limit = int(np.ceil(len(stats.layers) * layer_fraction)) + hits = np.argwhere(peaks[:limit] > threshold) + return [(int(stats.layers[row]), int(expert)) for row, expert in hits] + + +_ROUTER_RE = re.compile(r"^(?P(?:layers|mtp)\.(?P\d+))\.ffn\.gate\.(?Pweight|bias|bias_vl|tid2eid)$") +_EXPERT_RE = re.compile( + r"^(?P(?:layers|mtp)\.(?P\d+))\.ffn\.experts\.(?P\d+)\.(?P.+)$" +) + + +def prune_checkpoint_plan( + names: Iterable[str], + survivors: Mapping[int, Sequence[int]], + *, + prefix: str = "layers", +) -> dict[str, Any]: + """Plan a uniform prune over DeepSeek-V4-style checkpoint tensor names. + + ``survivors`` maps a block index under ``prefix`` (``layers`` for the + backbone, ``mtp`` for drafter stages) to the expert ids it keeps. Returns: + + - ``rename``: ``{source_name: target_name}`` for every kept routed-expert + tensor, renumbered ``0..keep-1`` in ascending source order; + - ``slice_rows``: ``{router_tensor_name: [kept ids]}`` for the gate's + ``weight``, ``bias`` and ``bias_vl`` (the image-token bias V4.1 adds, + which must be sliced with the rest or vision tokens route to the wrong + experts); + - ``drop``: expert tensors of pruned experts; + - ``untouched``: every other name, including blocks not in ``survivors``. + + Hash-routed blocks (``gate.tid2eid``) are refused: their token-to-expert + table names fixed experts and needs a remapping policy this plan does not + invent. DeepSeek-V4.1 has none; V4-Flash has three. + """ + + keep_sets: dict[int, list[int]] = {} + width = None + for block, kept in survivors.items(): + kept_sorted = sorted(int(e) for e in kept) + if len(set(kept_sorted)) != len(kept_sorted): + raise ValueError(f"block {block} keeps a duplicate expert") + if width is None: + width = len(kept_sorted) + elif len(kept_sorted) != width: + raise ValueError("every pruned block must keep the same expert count") + keep_sets[int(block)] = kept_sorted + + rename: dict[str, str] = {} + slice_rows: dict[str, list[int]] = {} + drop: list[str] = [] + untouched: list[str] = [] + seen_experts: dict[int, set[int]] = {block: set() for block in keep_sets} + for name in sorted(names): + router = _ROUTER_RE.match(name) + expert = _EXPERT_RE.match(name) + match = router or expert + if match is None or not match.group("prefix").startswith(prefix + "."): + untouched.append(name) + continue + block = int(match.group("idx")) + if block not in keep_sets: + untouched.append(name) + continue + kept = keep_sets[block] + if router is not None: + if router.group("leaf") == "tid2eid": + raise NotImplementedError( + f"{name}: hash-routed block {block} needs a token-to-expert " + "remapping policy before it can be pruned" + ) + slice_rows[name] = list(kept) + continue + old = int(expert.group("expert")) + seen_experts[block].add(old) + if old in kept: + new = kept.index(old) + rename[name] = f"{expert.group('prefix')}.ffn.experts.{new}.{expert.group('rest')}" + else: + drop.append(name) + for block, kept in keep_sets.items(): + absent = sorted(set(kept) - seen_experts[block]) + if absent: + raise ValueError(f"block {block} keeps experts with no tensors: {absent[:10]}") + return { + "record_type": "keep_reap_prune_plan_v1", + "prefix": prefix, + "keep": width, + "survivors": {str(block): kept for block, kept in sorted(keep_sets.items())}, + "rename": rename, + "slice_rows": slice_rows, + "drop": drop, + "untouched": untouched, + } + + +def write_survivors(path: str | Path, survivors: np.ndarray, *, layers: Sequence[int], meta: Mapping[str, Any]) -> None: + """Persist a survivor table with the settings that produced it.""" + + payload = { + "record_type": "keep_reap_survivors_v1", + "layers": [int(layer) for layer in layers], + "keep": int(survivors.shape[1]), + "survivors": {str(int(layer)): survivors[row].tolist() for row, layer in enumerate(layers)}, + **dict(meta), + } + Path(path).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") diff --git a/tests/test_dsv4_teacher_runner.py b/tests/test_dsv4_teacher_runner.py index 6b725b7..f403fc8 100644 --- a/tests/test_dsv4_teacher_runner.py +++ b/tests/test_dsv4_teacher_runner.py @@ -1781,6 +1781,123 @@ def test_gate_and_up_share_the_hidden_space_vector(self, tmp_path): assert gate["importance_sum"].shape == (args.hidden_size,) +class TestReapCollection: + """REAP saliency rides the calibration pass without touching the forward.""" + + def _prefill(self, *, reap=None, accumulator=None): + args, model = _tiny_model() + experts = _tiny_experts(args) + ids = mx.array(np.random.default_rng(5).integers(0, args.vocab_size, (1, 48))) + out = layer_major_prefill( + model, + ids, + chunk=16, + expert_provider=InMemoryExpertProvider(experts), + accumulator=accumulator, + collect_spaces=IMATRIX_SPACES if accumulator is not None else (), + release_after_layer=False, + reap=reap, + ) + return args, out + + def test_reap_alone_and_with_the_imatrix_leave_the_forward_unchanged(self): + from keep.quality.reap import ReapAccumulator + + args, plain = self._prefill() + layers = range(args.num_hidden_layers) + _, reap_only = self._prefill( + reap=ReapAccumulator(layers=layers, num_experts=args.n_routed_experts) + ) + _, both = self._prefill( + reap=ReapAccumulator(layers=layers, num_experts=args.n_routed_experts), + accumulator=_ImatrixAccumulator( + layers=layers, + num_experts=args.n_routed_experts, + dims={"hidden": args.hidden_size, "down": args.moe_intermediate_size}, + ), + ) + for a, b, c in zip(plain, reap_only, both): + assert mx.array_equal(a, b) and mx.array_equal(a, c) + + def test_reap_sees_the_same_routes_and_gate_weights_as_the_imatrix(self): + from keep.quality.reap import ReapAccumulator + + args, _ = _tiny_model() + layers = list(range(args.num_hidden_layers)) + reap = ReapAccumulator(layers=layers, num_experts=args.n_routed_experts) + imatrix = _ImatrixAccumulator( + layers=layers, + num_experts=args.n_routed_experts, + dims={"hidden": args.hidden_size, "down": args.moe_intermediate_size}, + ) + self._prefill(reap=reap, accumulator=imatrix) + stats = reap.stats() + arrays = imatrix.as_arrays() + np.testing.assert_array_equal(stats.route_count, arrays["route_count__hidden"]) + np.testing.assert_allclose( + stats.gate_sum, arrays["affinity_score_sum__hidden"], rtol=1e-6 + ) + assert np.all(stats.gate_norm_sum[stats.route_count > 0] > 0) + validate_dsv4_calibration_capture( + {**arrays, **stats.as_arrays()}, + layers=layers, + num_experts=args.n_routed_experts, + ) + + def test_capture_validation_refuses_reap_counts_that_disagree(self): + from keep.quality.reap import ReapAccumulator + + args, _ = _tiny_model() + layers = list(range(args.num_hidden_layers)) + reap = ReapAccumulator(layers=layers, num_experts=args.n_routed_experts) + imatrix = _ImatrixAccumulator( + layers=layers, + num_experts=args.n_routed_experts, + dims={"hidden": args.hidden_size, "down": args.moe_intermediate_size}, + ) + self._prefill(reap=reap, accumulator=imatrix) + arrays = {**imatrix.as_arrays(), **reap.stats().as_arrays()} + counts = arrays["reap__route_count"].copy() + nonzero = np.argwhere(counts > 1)[0] + counts[tuple(nonzero)] -= 1 + counts[nonzero[0], (nonzero[1] + 1) % counts.shape[1]] += 1 + arrays["reap__route_count"] = counts + with pytest.raises(ValueError, match="REAP"): + validate_dsv4_calibration_capture( + arrays, layers=layers, num_experts=args.n_routed_experts + ) + + def test_calibration_run_writes_reap_stats_and_finalize_sums_them(self, tmp_path): + from keep.quality.dsv4_teacher_runner import finalize_dsv4_reap_stats + from keep.quality.reap import ReapStats + + args, model = _tiny_model() + experts = _tiny_experts(args) + pack = _write_tiny_pack(tmp_path / "pack.json", vocab=args.vocab_size) + out = tmp_path / "out" + run_dsv4_teacher_production( + mode="calibration", + pack_path=pack, + out_dir=out, + chunk=16, + workload_loader=_tiny_loader(args, model, experts), + ) + report = finalize_dsv4_reap_stats(out) + assert report["session_count"] == 2 + stats = ReapStats.load(report["output"]) + assert len(stats.sessions) == 2 + assert stats.routed_slots.sum() == report["routed_slots"] > 0 + + # A session without REAP arrays (e.g. produced before this existed) + # makes the finalizer refuse rather than under-count. + victim = sorted((out / "sessions").glob("*.npz"))[0] + with np.load(victim, allow_pickle=False) as archive: + kept = {k: archive[k] for k in archive.files if not k.startswith("reap__")} + np.savez(victim, **kept) + with pytest.raises(ValueError, match="no REAP statistics"): + finalize_dsv4_reap_stats(out) + + class TestStreamStats: def test_device_rate_uses_busy_time_not_the_overlapped_span(self): """The span is not a throughput; only ``io_busy_seconds`` is. diff --git a/tests/test_reap.py b/tests/test_reap.py new file mode 100644 index 0000000..36e26e5 --- /dev/null +++ b/tests/test_reap.py @@ -0,0 +1,333 @@ +"""REAP saliency from the real routed forward, survivor choice, and the prune plan.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from keep.quality.reap import ( + ReapStats, + expert_saliency, + prune_checkpoint_plan, + pruned_set_jaccard, + select_survivors, + super_expert_candidates, +) + + +def _stats(route_count, gate_norm_sum, *, abs_max=None, gate_sum=None, norm_sum=None): + route_count = np.asarray(route_count, dtype=np.int64) + gate_norm_sum = np.asarray(gate_norm_sum, dtype=np.float64) + shape = route_count.shape + return ReapStats( + layers=np.arange(shape[0]), + num_experts=shape[1], + route_count=route_count, + gate_sum=np.zeros(shape) if gate_sum is None else gate_sum, + gate_norm_sum=gate_norm_sum, + norm_sum=np.zeros(shape) if norm_sum is None else norm_sum, + norm_max=np.zeros(shape), + abs_max=np.zeros(shape) if abs_max is None else abs_max, + routed_slots=route_count.sum(axis=1), + ) + + +# -- saliency ----------------------------------------------------------------- + + +def test_reap_is_the_mean_gate_weighted_norm_over_routed_slots(): + stats = _stats([[4, 1, 0]], [[2.0, 3.0, 0.0]]) + np.testing.assert_allclose(expert_saliency(stats, "reap"), [[0.5, 3.0, 0.0]]) + # A rarely routed expert with large outputs outranks a frequent quiet one, + # which is exactly where REAP and frequency pruning disagree. + np.testing.assert_array_equal(expert_saliency(stats, "frequency"), [[4, 1, 0]]) + + +def test_never_routed_experts_score_zero_under_every_criterion(): + stats = _stats([[3, 0]], [[1.5, 0.0]], gate_sum=[[0.9, 0.0]], norm_sum=[[6.0, 0.0]]) + for criterion in ("reap", "frequency", "weighted_frequency", "ean_mean"): + assert expert_saliency(stats, criterion)[0, 1] == 0.0 + + +def test_prior_shrinks_rare_experts_toward_the_layer_mean(): + stats = _stats([[1000, 1]], [[1000.0, 50.0]]) # layer mean ~1.05 + raw = expert_saliency(stats, "reap") + shrunk = expert_saliency(stats, "reap", prior_slots=100.0) + assert raw[0, 1] == pytest.approx(50.0) + assert 1.0 < shrunk[0, 1] < 2.0 # one noisy slot no longer dominates + assert shrunk[0, 0] == pytest.approx(raw[0, 0], rel=0.01) + + +def test_unknown_criterion_and_negative_prior_are_refused(): + stats = _stats([[1]], [[1.0]]) + with pytest.raises(ValueError): + expert_saliency(stats, "softmax") + with pytest.raises(ValueError): + expert_saliency(stats, "reap", prior_slots=-1) + + +# -- survivors and comparisons --------------------------------------------------- + + +def test_survivors_are_the_most_salient_in_ascending_id_order(): + saliency = np.array([[0.1, 0.9, 0.5, 0.7], [1.0, 0.0, 0.0, 0.2]]) + np.testing.assert_array_equal(select_survivors(saliency, 2), [[1, 3], [0, 3]]) + + +def test_ties_break_toward_the_lower_expert_id(): + np.testing.assert_array_equal(select_survivors(np.array([[0.5, 0.5, 0.5]]), 2), [[0, 1]]) + + +@pytest.mark.parametrize("keep", [0, 5]) +def test_keep_must_fit_the_expert_count(keep): + with pytest.raises(ValueError): + select_survivors(np.zeros((1, 4)), keep) + + +def test_pruned_set_jaccard(): + a = np.array([[0, 1], [0, 1]]) + b = np.array([[0, 1], [2, 3]]) + np.testing.assert_allclose(pruned_set_jaccard(a, b, num_experts=4), [1.0, 0.0]) + + +def test_super_expert_screen_sees_negative_outliers_through_abs_max(): + # Enough experts that the 99.5th percentile describes the bulk, as it does + # at 40 x 384 (it would be the outliers themselves at 4 x 8). + abs_max = np.full((8, 64), 1.0) + abs_max[1, 5] = 900.0 # e.g. a large negative output, invisible to a signed max + abs_max[7, 2] = 950.0 # last quarter of layers is excluded, as upstream does + stats = _stats(np.ones((8, 64), dtype=np.int64), np.ones((8, 64)), abs_max=abs_max) + assert super_expert_candidates(stats) == [(1, 5)] + + +# -- stats container ---------------------------------------------------------------- + + +def test_stats_validate_their_invariants(): + with pytest.raises(ValueError, match="sum to the routed slots"): + ReapStats( + layers=[0], + num_experts=2, + route_count=[[1, 1]], + gate_sum=np.zeros((1, 2)), + gate_norm_sum=np.zeros((1, 2)), + norm_sum=np.zeros((1, 2)), + norm_max=np.zeros((1, 2)), + abs_max=np.zeros((1, 2)), + routed_slots=[3], + ) + with pytest.raises(ValueError, match="no routed slots"): + _stats([[0, 2]], [[1.0, 1.0]]) + with pytest.raises(ValueError, match="non-negative"): + _stats([[1, 1]], [[-1.0, 1.0]]) + + +def test_stats_round_trip_and_merge(tmp_path): + a = _stats([[2, 1]], [[1.0, 2.0]], abs_max=np.array([[3.0, 1.0]])) + a.sessions = ("s0",) + b = _stats([[1, 3]], [[0.5, 1.0]], abs_max=np.array([[1.0, 4.0]])) + b.sessions = ("s1",) + merged = a.merged(b) + np.testing.assert_array_equal(merged.route_count, [[3, 4]]) + np.testing.assert_allclose(merged.gate_norm_sum, [[1.5, 3.0]]) + np.testing.assert_allclose(merged.abs_max, [[3.0, 4.0]]) + assert merged.sessions == ("s0", "s1") + with pytest.raises(ValueError, match="double-count"): + merged.merged(a) + + path = tmp_path / "reap-stats.npz" + merged.save(path) + loaded = ReapStats.load(path) + assert loaded.sessions == ("s0", "s1") + np.testing.assert_array_equal(loaded.route_count, merged.route_count) + np.testing.assert_allclose(loaded.gate_norm_sum, merged.gate_norm_sum) + + +# -- device accumulator --------------------------------------------------------------- + + +@pytest.mark.mlx +def test_accumulator_matches_a_float64_reference(): + mx = pytest.importorskip("mlx.core") + from keep.quality.reap import ReapAccumulator + + rng = np.random.default_rng(3) + experts, top_k, hidden = 6, 3, 16 + acc = ReapAccumulator(layers=[4, 9], num_experts=experts) + ref = {name: np.zeros((2, experts)) for name in ("count", "g", "gn", "n", "nmax", "amax")} + for layer_row, layer in enumerate((4, 9)): + for rows in (5, 11): # two chunks per layer + ids = rng.integers(0, experts, size=(rows, top_k)) + gates = rng.random((rows, top_k)).astype(np.float32) + outs = rng.normal(size=(rows, top_k, hidden)).astype(np.float32) * 3 + acc.add_chunk( + layer=layer, + expert_ids=mx.array(ids), + gate_weights=mx.array(gates), + expert_outputs=mx.array(outs), + ) + acc.eval_pending() + norms = np.linalg.norm(outs.astype(np.float64), axis=-1) + peaks = np.abs(outs).max(axis=-1) + for e, g, n, p in zip(ids.ravel(), gates.ravel(), norms.ravel(), peaks.ravel()): + ref["count"][layer_row, e] += 1 + ref["g"][layer_row, e] += g + ref["gn"][layer_row, e] += g * n + ref["n"][layer_row, e] += n + ref["nmax"][layer_row, e] = max(ref["nmax"][layer_row, e], n) + ref["amax"][layer_row, e] = max(ref["amax"][layer_row, e], p) + acc.flush_device() + stats = acc.stats(session="sess") + np.testing.assert_array_equal(stats.route_count, ref["count"]) + np.testing.assert_allclose(stats.gate_sum, ref["g"], rtol=2e-6) + np.testing.assert_allclose(stats.gate_norm_sum, ref["gn"], rtol=2e-6) + np.testing.assert_allclose(stats.norm_sum, ref["n"], rtol=2e-6) + np.testing.assert_allclose(stats.norm_max, ref["nmax"], rtol=2e-6) + np.testing.assert_allclose(stats.abs_max, ref["amax"], rtol=1e-7) + assert stats.sessions == ("sess",) + assert stats.routed_slots.tolist() == [48, 48] + + +@pytest.mark.mlx +def test_accumulator_refuses_mismatched_shapes_and_untracked_layers(): + mx = pytest.importorskip("mlx.core") + from keep.quality.reap import ReapAccumulator + + acc = ReapAccumulator(layers=[0], num_experts=4) + ids = mx.zeros((2, 2), dtype=mx.int32) + with pytest.raises(KeyError): + acc.add_chunk(layer=1, expert_ids=ids, gate_weights=ids, expert_outputs=mx.zeros((2, 2, 3))) + with pytest.raises(ValueError, match="do not match"): + acc.add_chunk(layer=0, expert_ids=ids, gate_weights=ids, expert_outputs=mx.zeros((2, 3, 3))) + + +# -- slicing the router is the same as masking it --------------------------------------- + + +@pytest.mark.mlx +def test_sliced_gate_equals_the_full_gate_restricted_to_survivors(): + """The prune applies by slicing router rows; check it against masking. + + Masking pruned experts out of *selection* (correction bias -> -inf) is + what "drop the expert and re-select top-k among survivors" means. The + sliced gate must pick the same experts with the same renormalised + weights, mapped back through the survivor table. + """ + + mx = pytest.importorskip("mlx.core") + from ramp.models.deepseek_v4_flash_adapter import DeepseekV4FlashVQModelArgs, MoEGate + + rng = np.random.default_rng(11) + experts, keep, hidden = 12, 7, 32 + args = DeepseekV4FlashVQModelArgs( + hidden_size=hidden, + n_routed_experts=experts, + num_experts_per_tok=3, + num_hash_layers=0, + num_hidden_layers=1, + compress_ratios=[0], + ) + weight = rng.normal(size=(experts, hidden)).astype(np.float32) + bias = rng.normal(size=(experts,)).astype(np.float32) * 0.1 + survivors = np.sort(rng.choice(experts, size=keep, replace=False)) + + full = MoEGate(args, 0) + full.weight = mx.array(weight) + masked_bias = bias.copy() + masked_bias[np.setdiff1d(np.arange(experts), survivors)] = -1e30 + full.e_score_correction_bias = mx.array(masked_bias) + + sliced_args = DeepseekV4FlashVQModelArgs( + hidden_size=hidden, + n_routed_experts=keep, + num_experts_per_tok=3, + num_hash_layers=0, + num_hidden_layers=1, + compress_ratios=[0], + ) + sliced = MoEGate(sliced_args, 0) + sliced.weight = mx.array(weight[survivors]) + sliced.e_score_correction_bias = mx.array(bias[survivors]) + + x = mx.array(rng.normal(size=(1, 64, hidden)).astype(np.float32)) + full_ids, full_scores = full(x) + sliced_ids, sliced_scores = sliced(x) + mapped = survivors[np.array(sliced_ids)] + np.testing.assert_array_equal(np.sort(mapped, axis=-1), np.sort(np.array(full_ids), axis=-1)) + order_full = np.argsort(np.array(full_ids), axis=-1) + order_sliced = np.argsort(mapped, axis=-1) + np.testing.assert_allclose( + np.take_along_axis(np.array(sliced_scores), order_sliced, axis=-1), + np.take_along_axis(np.array(full_scores), order_full, axis=-1), + rtol=1e-6, + ) + + +# -- checkpoint prune plan ------------------------------------------------------------ + + +def _names(layers=2, experts=4, *, hash_layer=None, with_bias_vl=True): + names = ["embed.weight", "head.weight", "mtp.0.ffn.gate.weight", "mtp.0.ffn.experts.0.w1.weight"] + for layer in range(layers): + p = f"layers.{layer}" + names += [f"{p}.attn.wq_a.weight", f"{p}.ffn.gate.weight", f"{p}.ffn.shared_experts.w1.weight"] + if layer == hash_layer: + names.append(f"{p}.ffn.gate.tid2eid") + else: + names.append(f"{p}.ffn.gate.bias") + if with_bias_vl: + names.append(f"{p}.ffn.gate.bias_vl") + for e in range(experts): + for proj in ("w1", "w2", "w3"): + names += [f"{p}.ffn.experts.{e}.{proj}.weight", f"{p}.ffn.experts.{e}.{proj}.scale"] + return names + + +def test_prune_plan_renumbers_slices_and_drops(): + plan = prune_checkpoint_plan(_names(), {0: [1, 3], 1: [0, 2]}) + assert plan["keep"] == 2 + assert plan["rename"]["layers.0.ffn.experts.3.w2.scale"] == "layers.0.ffn.experts.1.w2.scale" + assert plan["rename"]["layers.1.ffn.experts.0.w1.weight"] == "layers.1.ffn.experts.0.w1.weight" + assert "layers.0.ffn.experts.0.w1.weight" in plan["drop"] + assert len(plan["rename"]) == 2 * 2 * 3 * 2 + assert len(plan["drop"]) == 2 * 2 * 3 * 2 + # bias_vl is sliced with the gate, never left at the old width. + for leaf in ("weight", "bias", "bias_vl"): + assert plan["slice_rows"][f"layers.0.ffn.gate.{leaf}"] == [1, 3] + # Attention, shared experts, embeddings and the drafter are untouched. + for name in ("layers.0.attn.wq_a.weight", "layers.1.ffn.shared_experts.w1.weight", + "embed.weight", "mtp.0.ffn.gate.weight", "mtp.0.ffn.experts.0.w1.weight"): + assert name in plan["untouched"] + + +def test_prune_plan_refuses_hash_layers_and_ragged_or_phantom_survivors(): + with pytest.raises(NotImplementedError, match="hash-routed"): + prune_checkpoint_plan(_names(hash_layer=0), {0: [0, 1], 1: [0, 1]}) + with pytest.raises(ValueError, match="same expert count"): + prune_checkpoint_plan(_names(), {0: [0, 1], 1: [0, 1, 2]}) + with pytest.raises(ValueError, match="no tensors"): + prune_checkpoint_plan(_names(experts=4), {0: [0, 9], 1: [0, 1]}) + with pytest.raises(ValueError, match="duplicate"): + prune_checkpoint_plan(_names(), {0: [1, 1], 1: [0, 1]}) + + +def test_reap_select_command_reports_the_kill_test_and_writes_survivors(tmp_path, capsys): + import json + + from benchmarks import produce_dsv4_teacher_cache as cli + + rng = np.random.default_rng(2) + counts = rng.integers(0, 50, size=(3, 16)) + stats = _stats(counts, counts * rng.random((3, 16)) * 2) + stats.sessions = ("a", "b") + stats_path = tmp_path / "reap-stats.npz" + stats.save(stats_path) + out = tmp_path / "survivors.json" + + assert cli.main(["reap-select", "--stats", str(stats_path), "--keep", "8", "--output", str(out)]) == 0 + report = json.loads(capsys.readouterr().out) + assert report["keep"] == 8 and report["compare"] == "frequency" + assert set(report["pruned_set_jaccard_vs_compare"]) == {"0", "1", "2"} + saved = json.loads(out.read_text()) + assert saved["record_type"] == "keep_reap_survivors_v1" + assert all(len(ids) == 8 and ids == sorted(ids) for ids in saved["survivors"].values()) From 33a5681df826cc7bb66e95f1a6796863b0bfad40 Mon Sep 17 00:00:00 2001 From: Jack Mazac <114446521+jackmaxwil@users.noreply.github.com> Date: Thu, 5 Nov 2026 05:15:05 +0000 Subject: [PATCH 4/8] keep: replay recorded routing through an in-RAM expert hot set (SSD-tier kill test) Keeping every routed expert of a model too large for RAM means a hot set in memory and the rest read from SSD on a miss, on the decode critical path. keep.quality.expert_residency answers whether that is fast enough before an offload engine is built: - RouteRecorder keeps the first 2,048 tokens of routing per layer. In a causal model those prefill routes are also the decode routes. DSV4 calibration sessions now record them (routes__expert_ids int16, routes__gate_weights float16, about 1 MB per session). - simulate_residency replays them through a per-layer hot set under three policies: static (a fixed REAP-ordered set), LRU, and decayed use counts with a half-life (oMLX's V4.1 policy). It can warm-start from REAP saliency. fetch_threshold models a hybrid that drops light misses as a pruned model would and reports the gate mass it drops. simulate_sessions resets residency at each session boundary and pools raw counts. - residency_report applies the plan's rule per hot-set size: KILL if the best policy that fetches every miss still misses more than 5% of routed slots. decode_tokens_per_second turns misses into a tok/s estimate for a given SSD rate. - benchmarks/replay_dsv4_expert_residency.py runs it over a calibration run's sessions plus reap-stats.npz. layer_major_prefill's new hook is generalised to route_observers, so the REAP accumulator and the route recorder ride the same pass. The forward is still bit-identical. --- benchmarks/replay_dsv4_expert_residency.py | 153 ++++++++++ src/keep/quality/dsv4_teacher_runner.py | 43 ++- src/keep/quality/expert_residency.py | 339 +++++++++++++++++++++ tests/test_dsv4_teacher_runner.py | 11 +- tests/test_expert_residency.py | 206 +++++++++++++ 5 files changed, 736 insertions(+), 16 deletions(-) create mode 100644 benchmarks/replay_dsv4_expert_residency.py create mode 100644 src/keep/quality/expert_residency.py create mode 100644 tests/test_expert_residency.py diff --git a/benchmarks/replay_dsv4_expert_residency.py b/benchmarks/replay_dsv4_expert_residency.py new file mode 100644 index 0000000..f62ee60 --- /dev/null +++ b/benchmarks/replay_dsv4_expert_residency.py @@ -0,0 +1,153 @@ +"""Kill test: replay recorded routing through an in-RAM hot set of experts. + +Reads the ``routes__*`` arrays that DSV4 calibration sessions record (the +first tokens of each session, which are also its decode routing) and, if +given, ``reap-stats.npz`` to choose the initial hot set in REAP order. Replays +every hot-set size and policy asked for and applies the plan's rule: at a +given hot-set size, if the best policy that fetches every miss still misses +more than 5% of routed slots, keeping all experts behind an SSD tier at that +size is too slow to be worth building. + +Example (V4-Flash, half the experts resident):: + + uv run python benchmarks/replay_dsv4_expert_residency.py \\ + --calibration-dir ~/keep-artifacts/dsv4-teacher-calibration \\ + --reap-stats ~/keep-artifacts/dsv4-teacher-calibration/reap-stats.npz \\ + --hot 96,128,160 --expert-mb 6.39 --output residency.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from collections.abc import Sequence +from pathlib import Path + +import numpy as np + +from keep.quality.expert_residency import ( + ROUTE_ARRAY_PREFIX, + decode_tokens_per_second, + residency_report, + simulate_sessions, +) + + +def load_route_sessions(calibration_dir: Path) -> tuple[list[tuple[np.ndarray, np.ndarray]], np.ndarray]: + sessions: list[tuple[np.ndarray, np.ndarray]] = [] + layers: np.ndarray | None = None + missing: list[str] = [] + for path in sorted((calibration_dir / "sessions").glob("*.npz")): + with np.load(path, allow_pickle=False) as archive: + if f"{ROUTE_ARRAY_PREFIX}expert_ids" not in archive.files: + missing.append(path.name) + continue + current = archive[f"{ROUTE_ARRAY_PREFIX}layers"].astype(np.int64) + if layers is None: + layers = current + elif not np.array_equal(layers, current): + raise ValueError(f"{path.name} records a different layer set") + sessions.append( + ( + archive[f"{ROUTE_ARRAY_PREFIX}expert_ids"].astype(np.int64), + archive[f"{ROUTE_ARRAY_PREFIX}gate_weights"].astype(np.float32), + ) + ) + if missing: + raise ValueError(f"{len(missing)} sessions have no recorded routes: {missing[:5]}") + if not sessions or layers is None: + raise ValueError(f"no session routes under {calibration_dir / 'sessions'}") + return sessions, layers + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--calibration-dir", type=Path, required=True) + parser.add_argument("--reap-stats", type=Path, help="initial hot set in REAP order") + parser.add_argument("--num-experts", type=int, help="default: from --reap-stats") + parser.add_argument("--hot", default="128", help="comma-separated hot experts per layer") + parser.add_argument("--policies", default="static,lru,decayed") + parser.add_argument("--half-life", type=float, default=64.0) + parser.add_argument("--fetch-threshold", type=float, help="drop misses below this gate weight") + parser.add_argument("--warmup-tokens", type=int, default=0) + parser.add_argument("--expert-mb", type=float, default=8.9856, help="bytes per expert, MB (V4.1 at 2.031 bpw)") + parser.add_argument("--resident-ms", type=float, help="resident-only decode ms/token, for a tok/s estimate") + parser.add_argument("--ssd-gb-s", type=float, default=5.0) + parser.add_argument("--kill-miss-rate", type=float, default=0.05) + parser.add_argument("--output", type=Path) + args = parser.parse_args(argv) + + sessions, layers = load_route_sessions(args.calibration_dir.expanduser()) + initial = None + num_experts = args.num_experts + if args.reap_stats is not None: + from keep.quality.reap import ReapStats, expert_saliency + + stats = ReapStats.load(args.reap_stats.expanduser()) + if not np.array_equal(stats.layers, layers): + parser.error("--reap-stats covers different layers than the recorded routes") + initial = expert_saliency(stats, "reap") + num_experts = num_experts or stats.num_experts + if num_experts is None: + parser.error("pass --num-experts or --reap-stats") + policies = [p.strip() for p in args.policies.split(",") if p.strip()] + if "static" in policies and initial is None: + policies.remove("static") # a static set needs an order to choose it + expert_bytes = args.expert_mb * 1e6 + results = [] + for hot in (int(h) for h in args.hot.split(",")): + for policy in policies: + results.append( + simulate_sessions( + sessions, + num_experts=num_experts, + hot_per_layer=hot, + policy=policy, + half_life=args.half_life, + initial_order=initial, + expert_bytes=expert_bytes, + warmup_tokens=args.warmup_tokens, + ) + ) + if args.fetch_threshold is not None: + results.append( + simulate_sessions( + sessions, + num_experts=num_experts, + hot_per_layer=hot, + policy=policy, + half_life=args.half_life, + initial_order=initial, + expert_bytes=expert_bytes, + warmup_tokens=args.warmup_tokens, + fetch_threshold=args.fetch_threshold, + ) + ) + report = dict(residency_report(results, kill_miss_rate=args.kill_miss_rate)) + report["sessions"] = len(sessions) + report["layers"] = layers.tolist() + report["expert_mb"] = args.expert_mb + if args.resident_ms is not None: + report["tokens_per_second_estimate"] = { + f"{r.policy}@{r.hot_per_layer}" + + ("" if r.dropped == 0 else f"+drop<{args.fetch_threshold}"): round( + decode_tokens_per_second( + resident_ms_per_token=args.resident_ms, + misses_per_token=r.misses_per_token, + expert_bytes=expert_bytes, + ssd_bytes_per_second=args.ssd_gb_s * 1e9, + ), + 2, + ) + for r in results + } + text = json.dumps(report, indent=2, sort_keys=True) + if args.output is not None: + args.output.write_text(text + "\n") + print(text) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/keep/quality/dsv4_teacher_runner.py b/src/keep/quality/dsv4_teacher_runner.py index aa12438..6ebdc8c 100644 --- a/src/keep/quality/dsv4_teacher_runner.py +++ b/src/keep/quality/dsv4_teacher_runner.py @@ -1163,7 +1163,7 @@ def layer_major_prefill( on_layer: Callable[[int, int], None] | None = None, release_after_layer: bool = True, dspark_taps: dict[int, list[Any]] | None = None, - reap: Any | None = None, + route_observers: Sequence[Any] = (), ) -> list[Any]: """Prefill the whole backbone layer-outer, chunk-inner. @@ -1185,11 +1185,14 @@ def layer_major_prefill( to be a dict of lists rather than one concatenated array: the taps for layer 40 are complete before layer 41 has run at all. - ``reap`` is an optional :class:`keep.quality.reap.ReapAccumulator`. It - takes the model's own gate weights and each routed expert's output before - gating, so it scores experts against the routing the model actually uses. - Like the imatrix it reads tensors the forward already computes and leaves - the forward bit-identical. + ``route_observers`` receive every chunk's routing as the model computed + it: ``add_chunk(layer=, expert_ids=, gate_weights=, expert_outputs=)`` + with the model's own gate weights and each routed expert's output before + gating, then ``eval_pending()`` per chunk and ``flush_device()`` per + layer. :class:`keep.quality.reap.ReapAccumulator` and + :class:`keep.quality.expert_residency.RouteRecorder` are the two users. + Like the imatrix, observers read tensors the forward already computes and + leave it bit-identical. """ import mlx.core as mx @@ -1292,7 +1295,7 @@ def down_hook(activated, expert_ids, order, _capture=capture): ) mask = masks[chunk_index] - if accumulator is None and reap is None: + if accumulator is None and not route_observers: hidden[chunk_index] = layer(h, mask, layer_cache, piece) if layer_index in tap_layers: tap = _dspark_tap(hidden[chunk_index]) @@ -1324,8 +1327,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): flat_inds = inds.reshape(-1, inds.shape[-1]) flat_scores = scores.reshape(-1, scores.shape[-1]).astype(mx.float32) - if reap is not None: - reap.add_chunk( + for observer in route_observers: + observer.add_chunk( layer=layer_index, expert_ids=flat_inds, gate_weights=flat_scores, @@ -1333,7 +1336,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): ) if accumulator is None: mx.eval(hidden[chunk_index]) - reap.eval_pending() + for observer in route_observers: + observer.eval_pending() continue if collect_hidden: rows = moe_input.reshape(-1, moe_input.shape[-1]).astype(mx.float32) @@ -1363,8 +1367,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): ) mx.eval(hidden[chunk_index]) accumulator.eval_pending() - if reap is not None: - reap.eval_pending() + for observer in route_observers: + observer.eval_pending() switch_mlp.down_hook = None layer.ffn.switch_mlp = None @@ -1373,8 +1377,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): expert_provider.release(layer_index) if accumulator is not None: accumulator.flush_device() - if reap is not None: - reap.flush_device() + for observer in route_observers: + observer.flush_device() mx.clear_cache() if on_layer is not None: on_layer(layer_index, n_chunks) @@ -2702,11 +2706,19 @@ def _produce_sessions( token_ids = mx.array(session.token_ids.astype(np.int64)[None, :]) accumulator = None reap = None + routes = None + observers: tuple[Any, ...] = () collect: tuple[str, ...] = () if mode == "calibration": + from keep.quality.expert_residency import ( + ROUTE_RECORD_TOKENS, + RouteRecorder, + ) from keep.quality.reap import ReapAccumulator reap = ReapAccumulator(layers=layers, num_experts=num_experts) + routes = RouteRecorder(layers=layers, max_tokens=ROUTE_RECORD_TOKENS) + observers = (reap, routes) accumulator = _ImatrixAccumulator( layers=layers, num_experts=num_experts, @@ -2736,13 +2748,14 @@ def on_layer(layer_index: int, n_chunks: int, _s=session, _i=index) -> None: collect_spaces=collect, on_layer=on_layer, dspark_taps=taps, - reap=reap, + route_observers=observers, ) if mode == "calibration": arrays = { **calibration_session_arrays(accumulator), **reap.stats(session=session.prompt_id).as_arrays(), + **routes.as_arrays(), } validate_dsv4_calibration_capture( arrays, layers=layers, num_experts=num_experts diff --git a/src/keep/quality/expert_residency.py b/src/keep/quality/expert_residency.py new file mode 100644 index 0000000..bd66fd1 --- /dev/null +++ b/src/keep/quality/expert_residency.py @@ -0,0 +1,339 @@ +"""Would a hot set of experts in RAM, with the rest on SSD, keep decode fast? + +The kill test for keeping every routed expert of a model too large for RAM: +replay real routing through a fixed number of resident experts per layer and +count how often a selected expert is not resident. Each miss is an SSD read of +one expert (about 9 MB at 2.031 bpw on DeepSeek-V4.1) on the decode critical +path, so the miss rate converts directly into bytes and milliseconds. + +Routing is a property of the token, not of how it was computed: in a causal +model the experts token ``t`` selects in a prefill pass are the ones it selects +when decoded. So routes recorded during the teacher's calibration prefill +(:class:`RouteRecorder`) are the decode routing sequence for those tokens. + +Policies replayed by :func:`simulate_residency`: + +- ``static``: the resident set never changes (the REAP-ordered hot set a + pruned-plus-SSD artifact would ship). Misses are streamed, not cached. +- ``lru``: evict the least recently used resident. +- ``decayed``: evict the resident with the lowest exponentially decayed use + count (half-life in tokens), the policy oMLX's DeepSeek-V4.1 offload uses. + +``fetch_threshold`` models the hybrid "soft prune": a miss whose gate weight +is below the threshold is dropped (the expert is skipped and the remaining +weights renormalised, exactly as a pruned model would), and only heavier +misses are fetched. The result reports how much gate mass that drops. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from typing import Any, Literal + +import numpy as np + +__all__ = [ + "ROUTE_ARRAY_PREFIX", + "ROUTE_RECORD_TOKENS", + "ResidencyResult", + "RouteRecorder", + "decode_tokens_per_second", + "residency_report", + "simulate_residency", + "simulate_sessions", +] + +ROUTE_ARRAY_PREFIX = "routes__" +#: Tokens of routing recorded per calibration session: enough for a decode- +#: length replay, about 1 MB per session at 43 layers x top-6. +ROUTE_RECORD_TOKENS = 2048 +Policy = Literal["static", "lru", "decayed"] + + +class RouteRecorder: + """Keep the first ``max_tokens`` routes of every layer, in token order. + + Implements the teacher runner's route-observer interface. The layer loop + is layer-outer and chunk-inner, so each layer's rows arrive in token order. + """ + + def __init__(self, *, layers: Sequence[int], max_tokens: int) -> None: + if max_tokens <= 0: + raise ValueError("max_tokens must be positive") + self.layers = tuple(int(layer) for layer in layers) + self.max_tokens = int(max_tokens) + self._ids: dict[int, list[np.ndarray]] = {layer: [] for layer in self.layers} + self._gates: dict[int, list[np.ndarray]] = {layer: [] for layer in self.layers} + self._count = {layer: 0 for layer in self.layers} + + def add_chunk( + self, *, layer: int, expert_ids: Any, gate_weights: Any, expert_outputs: Any = None + ) -> None: + del expert_outputs # routing only + if layer not in self._count: + raise KeyError(f"layer {layer} is not tracked by this recorder") + room = self.max_tokens - self._count[layer] + if room <= 0: + return + ids = np.asarray(expert_ids)[:room] + gates = np.asarray(gate_weights, dtype=np.float32)[:room] + if ids.shape != gates.shape or ids.ndim != 2: + raise ValueError("expert ids and gate weights must share a [rows, top_k] shape") + self._ids[layer].append(ids.astype(np.int16)) + self._gates[layer].append(gates.astype(np.float16)) + self._count[layer] += int(ids.shape[0]) + + def eval_pending(self) -> None: + return None + + def flush_device(self) -> None: + return None + + def as_arrays(self) -> dict[str, np.ndarray]: + counts = {self._count[layer] for layer in self.layers} + if len(counts) != 1: + raise ValueError(f"layers recorded different token counts: {sorted(counts)}") + ids = np.stack([np.concatenate(self._ids[layer]) for layer in self.layers], axis=1) + gates = np.stack([np.concatenate(self._gates[layer]) for layer in self.layers], axis=1) + return { + f"{ROUTE_ARRAY_PREFIX}layers": np.asarray(self.layers, dtype=np.int64), + f"{ROUTE_ARRAY_PREFIX}expert_ids": ids, # int16 [tokens, layers, top_k] + f"{ROUTE_ARRAY_PREFIX}gate_weights": gates, # float16, same shape + } + + +@dataclass(frozen=True) +class ResidencyResult: + policy: str + hot_per_layer: int + num_experts: int + tokens: int + routed_slots: int + misses: int # fetched from SSD + dropped: int # skipped under fetch_threshold + miss_rate: float + misses_per_token: float + dropped_gate_mass_fraction: float + ssd_bytes_per_token: float | None + per_layer_miss_rate: tuple[float, ...] + dropped_gate_mass: float = 0.0 + total_gate_mass: float = 0.0 + per_layer_misses: tuple[int, ...] = () + + def to_dict(self) -> dict[str, Any]: + out = asdict(self) + out["per_layer_miss_rate"] = list(self.per_layer_miss_rate) + return out + + +def simulate_residency( + expert_ids: np.ndarray, + *, + num_experts: int, + hot_per_layer: int, + policy: Policy = "decayed", + half_life: float = 64.0, + initial_order: np.ndarray | None = None, + gate_weights: np.ndarray | None = None, + fetch_threshold: float | None = None, + expert_bytes: float | None = None, + warmup_tokens: int = 0, +) -> ResidencyResult: + """Replay ``expert_ids [tokens, layers, top_k]`` through a per-layer hot set. + + ``initial_order`` is ``[layers, experts]`` saliency (e.g. REAP scores); the + hot set starts as its top ``hot_per_layer`` experts. Without it the set + starts empty and fills on demand. ``warmup_tokens`` are replayed but not + counted. + """ + + ids = np.asarray(expert_ids) + if ids.ndim != 3: + raise ValueError("expert_ids must be [tokens, layers, top_k]") + tokens, layers, top_k = ids.shape + if not 0 < hot_per_layer <= num_experts: + raise ValueError(f"hot_per_layer must be in 1..{num_experts}") + if hot_per_layer < top_k: + raise ValueError("the hot set must hold at least one token's top-k") + if np.any(ids < 0) or np.any(ids >= num_experts): + raise ValueError("expert ids outside the expert range") + if policy not in ("static", "lru", "decayed"): + raise ValueError(f"unknown policy {policy!r}") + if policy == "static" and initial_order is None: + raise ValueError("a static hot set needs initial_order to choose it") + if fetch_threshold is not None and gate_weights is None: + raise ValueError("fetch_threshold needs gate_weights") + gates = None if gate_weights is None else np.asarray(gate_weights, dtype=np.float64) + if gates is not None and gates.shape != ids.shape: + raise ValueError("gate_weights must match expert_ids") + if not 0 <= warmup_tokens < tokens: + raise ValueError("warmup_tokens must leave at least one counted token") + decay = np.log(2.0) / float(half_life) + + misses = np.zeros(layers, dtype=np.int64) + dropped = 0 + dropped_mass = 0.0 + total_mass = 0.0 + for layer in range(layers): + resident = np.zeros(num_experts, dtype=bool) + if initial_order is not None: + order = np.lexsort((np.arange(num_experts), -np.asarray(initial_order)[layer])) + resident[order[:hot_per_layer]] = True + score = np.zeros(num_experts, dtype=np.float64) # decayed count at `stamp` + stamp = np.zeros(num_experts, dtype=np.float64) + last_use = np.full(num_experts, -1.0) + for t in range(tokens): + counted = t >= warmup_tokens + needed = ids[t, layer] + needed_set = set(int(e) for e in needed) + for slot, expert in enumerate(needed): + expert = int(expert) + weight = None if gates is None else float(gates[t, layer, slot]) + if counted and weight is not None: + total_mass += weight + if not resident[expert]: + if fetch_threshold is not None and weight < fetch_threshold: + if counted: + dropped += 1 + dropped_mass += weight + continue + if counted: + misses[layer] += 1 + if policy != "static": + if resident.sum() >= hot_per_layer: + candidates = np.flatnonzero(resident) + candidates = candidates[ + [c not in needed_set for c in candidates.tolist()] + ] + if policy == "lru": + victim = candidates[np.argmin(last_use[candidates])] + else: + current = score[candidates] * np.exp(-decay * (t - stamp[candidates])) + victim = candidates[np.argmin(current)] + resident[victim] = False + resident[expert] = True + # Usage bookkeeping for every selected expert, resident or not. + score[expert] = score[expert] * np.exp(-decay * (t - stamp[expert])) + 1.0 + stamp[expert] = t + last_use[expert] = t + counted_tokens = tokens - warmup_tokens + slots = counted_tokens * layers * top_k + total_misses = int(misses.sum()) + per_layer = tuple(float(m) / (counted_tokens * top_k) for m in misses) + return ResidencyResult( + policy=policy, + hot_per_layer=int(hot_per_layer), + num_experts=int(num_experts), + tokens=int(counted_tokens), + routed_slots=int(slots), + misses=total_misses, + dropped=int(dropped), + miss_rate=total_misses / slots, + misses_per_token=total_misses / counted_tokens, + dropped_gate_mass_fraction=(dropped_mass / total_mass) if total_mass else 0.0, + ssd_bytes_per_token=None if expert_bytes is None else total_misses / counted_tokens * float(expert_bytes), + per_layer_miss_rate=per_layer, + dropped_gate_mass=float(dropped_mass), + total_gate_mass=float(total_mass), + per_layer_misses=tuple(int(m) for m in misses), + ) + + +def simulate_sessions( + sessions: Sequence[tuple[np.ndarray, np.ndarray | None]], + *, + expert_bytes: float | None = None, + **kwargs: Any, +) -> ResidencyResult: + """Replay each session from the initial hot set and pool the counts. + + Residency is reset at every session boundary: unrelated sessions do not + share a warm cache in a real deployment's worst case, and pooling raw + counts weights each session by its tokens. + """ + + if not sessions: + raise ValueError("no sessions to replay") + parts = [ + simulate_residency(ids, gate_weights=gates, expert_bytes=expert_bytes, **kwargs) + for ids, gates in sessions + ] + tokens = sum(r.tokens for r in parts) + slots = sum(r.routed_slots for r in parts) + misses = sum(r.misses for r in parts) + layer_misses = np.sum([r.per_layer_misses for r in parts], axis=0) + top_k = slots // (tokens * len(layer_misses)) + dropped_mass = sum(r.dropped_gate_mass for r in parts) + total_mass = sum(r.total_gate_mass for r in parts) + first = parts[0] + return ResidencyResult( + policy=first.policy, + hot_per_layer=first.hot_per_layer, + num_experts=first.num_experts, + tokens=tokens, + routed_slots=slots, + misses=misses, + dropped=sum(r.dropped for r in parts), + miss_rate=misses / slots, + misses_per_token=misses / tokens, + dropped_gate_mass_fraction=(dropped_mass / total_mass) if total_mass else 0.0, + ssd_bytes_per_token=None if expert_bytes is None else misses / tokens * float(expert_bytes), + per_layer_miss_rate=tuple(float(m) / (tokens * top_k) for m in layer_misses), + dropped_gate_mass=dropped_mass, + total_gate_mass=total_mass, + per_layer_misses=tuple(int(m) for m in layer_misses), + ) + + +def decode_tokens_per_second( + *, + resident_ms_per_token: float, + misses_per_token: float, + expert_bytes: float, + ssd_bytes_per_second: float, + host_sync_ms_per_token: float = 0.0, +) -> float: + """Decode rate once misses are read synchronously on the critical path. + + ``host_sync_ms_per_token`` is the cost of learning which experts a layer + picked before issuing the reads (oMLX measured it as a large share of its + decode cycle); it is paid only by policies that fetch. + """ + + miss_ms = misses_per_token * expert_bytes / ssd_bytes_per_second * 1e3 + return 1e3 / (resident_ms_per_token + miss_ms + host_sync_ms_per_token) + + +def residency_report( + results: Sequence[ResidencyResult], *, kill_miss_rate: float = 0.05 +) -> Mapping[str, Any]: + """Summarise replays and apply the plan's kill rule per hot-set size. + + For each hot-set size, the best policy that *fetches* every miss (no + dropped gate mass) decides: above ``kill_miss_rate`` misses the SSD tier + costs too much at that size. + """ + + by_hot: dict[int, list[ResidencyResult]] = {} + for result in results: + by_hot.setdefault(result.hot_per_layer, []).append(result) + verdicts: dict[str, Any] = {} + for hot, group in sorted(by_hot.items()): + exact = [r for r in group if r.dropped == 0] + if not exact: + continue + best = min(exact, key=lambda r: r.miss_rate) + verdicts[str(hot)] = { + "best_policy": best.policy, + "miss_rate": best.miss_rate, + "misses_per_token": best.misses_per_token, + "ssd_bytes_per_token": best.ssd_bytes_per_token, + "verdict": "KILL" if best.miss_rate > kill_miss_rate else "PASS", + } + return { + "record_type": "keep_expert_residency_replay_v1", + "kill_miss_rate": kill_miss_rate, + "results": [result.to_dict() for result in results], + "by_hot_per_layer": verdicts, + } diff --git a/tests/test_dsv4_teacher_runner.py b/tests/test_dsv4_teacher_runner.py index f403fc8..664d8f2 100644 --- a/tests/test_dsv4_teacher_runner.py +++ b/tests/test_dsv4_teacher_runner.py @@ -1796,7 +1796,7 @@ def _prefill(self, *, reap=None, accumulator=None): accumulator=accumulator, collect_spaces=IMATRIX_SPACES if accumulator is not None else (), release_after_layer=False, - reap=reap, + route_observers=() if reap is None else (reap,), ) return args, out @@ -1888,6 +1888,15 @@ def test_calibration_run_writes_reap_stats_and_finalize_sums_them(self, tmp_path assert len(stats.sessions) == 2 assert stats.routed_slots.sum() == report["routed_slots"] > 0 + # The same pass records each session's routing for the residency replay. + first = sorted((out / "sessions").glob("*.npz"))[0] + with np.load(first, allow_pickle=False) as archive: + route_ids = archive["routes__expert_ids"] + assert route_ids.dtype == np.int16 + assert route_ids.shape[1:] == (args.num_hidden_layers, args.num_experts_per_tok) + assert route_ids.shape[0] > 0 + assert archive["routes__gate_weights"].shape == route_ids.shape + # A session without REAP arrays (e.g. produced before this existed) # makes the finalizer refuse rather than under-count. victim = sorted((out / "sessions").glob("*.npz"))[0] diff --git a/tests/test_expert_residency.py b/tests/test_expert_residency.py new file mode 100644 index 0000000..abd9611 --- /dev/null +++ b/tests/test_expert_residency.py @@ -0,0 +1,206 @@ +"""Replay of routing through an in-RAM hot set: the SSD-tier kill test.""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +from keep.quality.expert_residency import ( + RouteRecorder, + decode_tokens_per_second, + residency_report, + simulate_residency, + simulate_sessions, +) + + +def _routes(tokens, layers, top_k, experts, seed=0, skew=1.2): + """Zipf-ish routing so a hot set is meaningful.""" + + rng = np.random.default_rng(seed) + weights = 1.0 / np.arange(1, experts + 1) ** skew + weights /= weights.sum() + ids = np.empty((tokens, layers, top_k), dtype=np.int64) + for t in range(tokens): + for layer in range(layers): + ids[t, layer] = rng.choice(experts, size=top_k, replace=False, p=weights) + gates = rng.dirichlet(np.ones(top_k), size=(tokens, layers)) + return ids, gates + + +def test_everything_resident_never_misses_once_warm(): + ids, _ = _routes(50, 2, 3, 8) + order = np.zeros((2, 8)) + warm = simulate_residency(ids, num_experts=8, hot_per_layer=8, policy="decayed", initial_order=order) + assert warm.misses == 0 and warm.miss_rate == 0.0 + # A cold start pays exactly one compulsory miss per distinct expert. + cold = simulate_residency(ids, num_experts=8, hot_per_layer=8, policy="decayed") + assert cold.misses == sum(len(np.unique(ids[:, layer])) for layer in range(2)) + + +def test_hand_checked_lru_and_static_counts(): + # One layer, top-1, two slots: A B A C A (experts 0 1 0 2 0) + ids = np.array([0, 1, 0, 2, 0]).reshape(5, 1, 1) + lru = simulate_residency(ids, num_experts=3, hot_per_layer=2, policy="lru") + assert lru.misses == 3 # cold 0, cold 1, then 2 evicts 1 -> 0 stays a hit + static = simulate_residency( + ids, num_experts=3, hot_per_layer=2, policy="static", initial_order=np.array([[3.0, 2.0, 1.0]]) + ) + assert static.misses == 1 # only expert 2 is outside the fixed set {0, 1} + + +def test_decayed_counts_keep_the_frequent_expert_through_one_off_scans(): + # Expert 0 is used often; pairs of one-off experts arrive between its uses. + # LRU evicts 0 on every scan; decayed counts keep it resident. + seq = [0] * 4 + [1, 2, 0, 3, 4, 0, 5, 6, 0] + ids = np.array(seq).reshape(-1, 1, 1) + lru = simulate_residency(ids, num_experts=7, hot_per_layer=2, policy="lru") + decayed = simulate_residency(ids, num_experts=7, hot_per_layer=2, policy="decayed", half_life=64) + assert decayed.misses == 7 # the first 0 and the six one-offs + assert lru.misses == 10 # plus three re-fetches of 0 + + +def test_reap_ordered_warm_start_beats_a_cold_start(): + ids, _ = _routes(300, 3, 4, 32, seed=1) + order = -np.tile(np.arange(32, dtype=float), (3, 1)) # expert 0 most salient + cold = simulate_residency(ids, num_experts=32, hot_per_layer=12, policy="decayed") + warm = simulate_residency( + ids, num_experts=32, hot_per_layer=12, policy="decayed", initial_order=order + ) + assert warm.misses < cold.misses + + +def test_fetch_threshold_drops_light_misses_and_reports_their_mass(): + ids, gates = _routes(200, 2, 4, 32, seed=2) + order = -np.tile(np.arange(32, dtype=float), (2, 1)) + exact = simulate_residency( + ids, num_experts=32, hot_per_layer=8, policy="static", initial_order=order, gate_weights=gates + ) + hybrid = simulate_residency( + ids, + num_experts=32, + hot_per_layer=8, + policy="static", + initial_order=order, + gate_weights=gates, + fetch_threshold=0.2, + ) + assert hybrid.misses + hybrid.dropped == exact.misses + assert hybrid.misses < exact.misses + assert 0.0 < hybrid.dropped_gate_mass_fraction < 1.0 + assert exact.dropped == 0 and exact.dropped_gate_mass_fraction == 0.0 + + +def test_warmup_tokens_are_replayed_but_not_counted(): + ids, _ = _routes(100, 1, 2, 16, seed=3) + result = simulate_residency(ids, num_experts=16, hot_per_layer=4, policy="lru", warmup_tokens=40) + assert result.tokens == 60 and result.routed_slots == 120 + + +def test_sessions_reset_and_pool_by_token_count(): + a, ga = _routes(80, 2, 3, 16, seed=4) + b, gb = _routes(40, 2, 3, 16, seed=5) + pooled = simulate_sessions([(a, ga), (b, gb)], num_experts=16, hot_per_layer=6, policy="lru", expert_bytes=1e6) + one = simulate_residency(a, num_experts=16, hot_per_layer=6, policy="lru") + two = simulate_residency(b, num_experts=16, hot_per_layer=6, policy="lru") + assert pooled.misses == one.misses + two.misses + assert pooled.tokens == 120 + assert pooled.ssd_bytes_per_token == pytest.approx(pooled.misses_per_token * 1e6) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"hot_per_layer": 0}, + {"hot_per_layer": 2}, # fewer than top-k + {"policy": "fifo"}, + {"policy": "static"}, # needs an order + {"fetch_threshold": 0.1}, # needs gates + ], +) +def test_invalid_replays_are_refused(kwargs): + ids, _ = _routes(10, 1, 3, 8) + params = {"num_experts": 8, "hot_per_layer": 4, "policy": "lru", **kwargs} + with pytest.raises(ValueError): + simulate_residency(ids, **params) + + +def test_report_applies_the_kill_rule_per_hot_size(): + ids, _ = _routes(200, 2, 4, 32, seed=6) + small = simulate_residency(ids, num_experts=32, hot_per_layer=6, policy="lru") + full = simulate_residency(ids, num_experts=32, hot_per_layer=32, policy="lru") + report = residency_report([small, full], kill_miss_rate=0.05) + assert report["by_hot_per_layer"]["32"]["verdict"] == "PASS" + assert report["by_hot_per_layer"]["6"]["verdict"] == "KILL" + + +def test_decode_rate_estimate(): + # 30 ms resident, 15 misses x 9 MB at 5 GB/s = 27 ms more + rate = decode_tokens_per_second( + resident_ms_per_token=30.0, misses_per_token=15, expert_bytes=9e6, ssd_bytes_per_second=5e9 + ) + assert rate == pytest.approx(1e3 / 57.0) + + +def test_route_recorder_keeps_the_first_tokens_in_order(): + recorder = RouteRecorder(layers=[3, 7], max_tokens=5) + for layer in (3, 7): + recorder.add_chunk(layer=layer, expert_ids=np.arange(6).reshape(3, 2), gate_weights=np.full((3, 2), 0.5)) + recorder.add_chunk(layer=layer, expert_ids=np.arange(6, 12).reshape(3, 2), gate_weights=np.full((3, 2), 0.25)) + arrays = recorder.as_arrays() + ids = arrays["routes__expert_ids"] + assert ids.shape == (5, 2, 2) and ids.dtype == np.int16 + np.testing.assert_array_equal(ids[:, 0].ravel(), np.arange(10)) + assert arrays["routes__gate_weights"][4, 1, 0] == np.float16(0.25) + + +def test_replay_cli_on_recorded_calibration_sessions(tmp_path, capsys): + from benchmarks import replay_dsv4_expert_residency as cli + from keep.quality.reap import ReapStats + + sessions = tmp_path / "cal" / "sessions" + sessions.mkdir(parents=True) + counts = np.zeros((2, 16), dtype=np.int64) + for index in range(2): + ids, gates = _routes(60, 2, 2, 16, seed=10 + index) + np.savez( + sessions / f"s{index}.npz", + routes__layers=np.array([0, 1]), + routes__expert_ids=ids.astype(np.int16), + routes__gate_weights=gates.astype(np.float16), + ) + for layer in range(2): + counts[layer] += np.bincount(ids[:, layer].ravel(), minlength=16) + stats = ReapStats( + layers=[0, 1], + num_experts=16, + route_count=counts, + gate_sum=counts * 0.5, + gate_norm_sum=counts * 1.0, + norm_sum=counts * 2.0, + norm_max=np.ones((2, 16)), + abs_max=np.ones((2, 16)), + routed_slots=counts.sum(axis=1), + ) + stats.save(tmp_path / "reap-stats.npz") + out = tmp_path / "residency.json" + assert ( + cli.main( + [ + "--calibration-dir", str(tmp_path / "cal"), + "--reap-stats", str(tmp_path / "reap-stats.npz"), + "--hot", "4,16", + "--fetch-threshold", "0.3", + "--resident-ms", "30", + "--output", str(out), + ] + ) + == 0 + ) + report = json.loads(out.read_text()) + assert report["sessions"] == 2 + assert report["by_hot_per_layer"]["16"]["verdict"] == "PASS" + assert {r["policy"] for r in report["results"]} == {"static", "lru", "decayed"} + assert any(key.endswith("+drop<0.3") for key in report["tokens_per_second_estimate"]) From 5baed30da02a24a843c5f5597f17aab21d212681 Mon Sep 17 00:00:00 2001 From: Jack Mazac <114446521+jackmaxwil@users.noreply.github.com> Date: Thu, 5 Nov 2026 05:18:16 +0000 Subject: [PATCH 5/8] ramp: opt-in one-dispatch sinkhorn for DeepSeek-V4 hyper-connections Every hyper-connection turns hc_mult x hc_mult mixing logits into a doubly-stochastic matrix through a row softmax and 19 more rounds of row/column normalisation. The model's compiled path still launches a few kernels per round, because each normalisation is a reduction, and it runs 86 times per decoded V4-Flash token. The roadmap measured 30-70 ms of sinkhorn per token. ramp.kernels.hc_sinkhorn.fused_hc_sinkhorn does the whole sinkhorn in one Metal dispatch, one thread per matrix, with the values in registers, following the compiled path's operation order. exp and summation order are not guaranteed bit-identical, so it is opt-in: RAMP_HC_SINKHORN=metal selects it in HyperConnection (also used by the GLM-5 next adapter), and the default path is untouched. _hc_split_logits_ops computes pre, post and the comb logits exactly as the compiled path does. hc_sinkhorn_reference states the algorithm in float32 NumPy; a test pins it to the compiled path on CPU. Metal-only tests compare the fused kernel with the compiled path and HyperConnection's outputs with the flag on and off; they skip without a Metal GPU. benchmarks/bench_hc_sinkhorn.py is the kill test: KILL if the fused path is under 3x faster at decode shape or differs by more than 2e-6 relative. --- benchmarks/bench_hc_sinkhorn.py | 123 ++++++++++++++ src/ramp/kernels/hc_sinkhorn.py | 163 +++++++++++++++++++ src/ramp/models/deepseek_v4_flash_adapter.py | 55 ++++++- tests/test_hc_sinkhorn.py | 101 ++++++++++++ 4 files changed, 434 insertions(+), 8 deletions(-) create mode 100644 benchmarks/bench_hc_sinkhorn.py create mode 100644 src/ramp/kernels/hc_sinkhorn.py create mode 100644 tests/test_hc_sinkhorn.py diff --git a/benchmarks/bench_hc_sinkhorn.py b/benchmarks/bench_hc_sinkhorn.py new file mode 100644 index 0000000..f56fdb7 --- /dev/null +++ b/benchmarks/bench_hc_sinkhorn.py @@ -0,0 +1,123 @@ +"""Kill test: the one-dispatch hyper-connection sinkhorn against the compiled path. + +Times the model's compiled ``_hc_split_sinkhorn_ops`` against the logits half +plus :func:`ramp.kernels.hc_sinkhorn.fused_hc_sinkhorn`, at decode shape (one +row) and a prefill chunk, and reports milliseconds per DeepSeek-V4-Flash +decode step (86 calls: two hyper-connections in each of 43 blocks). + +KILL if the fused path is less than 3x faster at decode shape, or if it +differs from the compiled path by more than 2e-6 relative. If it passes, set +``RAMP_HC_SINKHORN=metal`` for the decode headline. + + uv run python benchmarks/bench_hc_sinkhorn.py --output-json artifacts/benchmarks/hc-sinkhorn-m4max.json +""" + +from __future__ import annotations + +import argparse +import json +import platform +import sys +import time +from collections.abc import Sequence +from pathlib import Path + +import mlx.core as mx +import numpy as np + +from ramp.kernels.hc_sinkhorn import fused_hc_sinkhorn +from ramp.models.deepseek_v4_flash_adapter import _hc_split_logits_ops, _hc_split_sinkhorn_ops + +CALLS_PER_DECODE_STEP = 86 +MIN_DECODE_SPEEDUP = 3.0 +MAX_RELATIVE_DIFF = 2e-6 + + +def _per_call_ms(fn, iterations: int, warmup: int) -> float: + for _ in range(warmup): + mx.eval(fn()) + started = time.perf_counter() + for _ in range(iterations): + mx.eval(fn()) + return (time.perf_counter() - started) * 1e3 / iterations + + +def measure(rows: int, *, hc: int, iters: int, eps: float, iterations: int, warmup: int, seed: int) -> dict: + rng = np.random.default_rng(seed) + mix = (2 + hc) * hc + mixes = mx.array(rng.normal(size=(1, rows, mix)).astype(np.float32) * 3) + scale = mx.array(rng.uniform(0.5, 2.0, size=(3,)).astype(np.float32)) + base = mx.array(rng.normal(size=(mix,)).astype(np.float32)) + + def compiled(): + return _hc_split_sinkhorn_ops(mixes, scale, base, hc, iters, eps) + + def fused(): + pre, post, logits = _hc_split_logits_ops(mixes, scale, base, hc, eps) + return pre, post, fused_hc_sinkhorn(logits, iters=iters, eps=eps) + + want, got = np.array(compiled()[2]), np.array(fused()[2]) + rel = float(np.max(np.abs(got - want) / np.maximum(np.abs(want), 1e-12))) + t_compiled = _per_call_ms(compiled, iterations, warmup) + t_fused = _per_call_ms(fused, iterations, warmup) + return { + "rows": rows, + "compiled_ms_per_call": t_compiled, + "fused_ms_per_call": t_fused, + "speedup": t_compiled / t_fused, + "compiled_ms_per_decode_step": t_compiled * CALLS_PER_DECODE_STEP, + "fused_ms_per_decode_step": t_fused * CALLS_PER_DECODE_STEP, + "max_relative_diff": rel, + } + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--rows", default="1,1024", help="matrices per call (1 = decode)") + parser.add_argument("--hc-mult", type=int, default=4) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--eps", type=float, default=1e-6) + parser.add_argument("--iterations", type=int, default=500) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--seed", type=int, default=20261105) + parser.add_argument("--output-json", type=Path) + args = parser.parse_args(argv) + + on_gpu = mx.default_device() == mx.gpu + rows = [ + measure( + int(r), + hc=args.hc_mult, + iters=args.iters, + eps=args.eps, + iterations=args.iterations, + warmup=args.warmup, + seed=args.seed, + ) + for r in args.rows.split(",") + ] + decode = next((row for row in rows if row["rows"] == 1), rows[0]) + reasons = [] + if any(row["max_relative_diff"] > MAX_RELATIVE_DIFF for row in rows): + reasons.append(f"fused output differs by more than {MAX_RELATIVE_DIFF} relative") + if decode["speedup"] < MIN_DECODE_SPEEDUP: + reasons.append(f"decode-shape speedup {decode['speedup']:.2f}x is under {MIN_DECODE_SPEEDUP}x") + result = { + "record_type": "ramp_hc_sinkhorn_kill_test_v1", + "device": str(mx.default_device()), + "mlx_version": mx.__version__, + "platform": platform.platform(), + "rows": rows, + "verdict": "KILL" if reasons else ("PASS" if on_gpu else "NOT_ON_GPU"), + "reasons": reasons, + } + text = json.dumps(result, indent=2, sort_keys=True) + if args.output_json is not None: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(text + "\n") + print(text) + return 1 if result["verdict"] == "KILL" else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/ramp/kernels/hc_sinkhorn.py b/src/ramp/kernels/hc_sinkhorn.py new file mode 100644 index 0000000..267e6ce --- /dev/null +++ b/src/ramp/kernels/hc_sinkhorn.py @@ -0,0 +1,163 @@ +"""One-dispatch Sinkhorn normalisation for DeepSeek-V4 hyper-connections. + +Every hyper-connection (two per block: 86 calls per decoded token on +DeepSeek-V4-Flash) turns ``hc_mult x hc_mult`` mixing logits into a doubly +stochastic matrix: a row softmax, a column normalisation, then +``sinkhorn_iters - 1`` alternating row and column normalisations. The +compiled MLX path still launches a few kernels per iteration because each +normalisation is a reduction, so a 4x4 matrix costs dozens of dispatches per +call. The roadmap measured 30 to 70 ms of sinkhorn per decoded token. + +:func:`fused_hc_sinkhorn` does the whole thing in one Metal dispatch, one +thread per matrix, holding the 16 values in registers. It follows the MLX +path's operation order, but exp and summation order are not guaranteed +bit-identical, so it is **opt-in** (``RAMP_HC_SINKHORN=metal``) until the +parity and speed check in ``benchmarks/bench_hc_sinkhorn.py`` passes on the +target Mac. + +:func:`hc_sinkhorn_reference` is a float32 NumPy statement of the same +algorithm; tests pin it against the model's compiled path, and the Metal +kernel against both. +""" + +from __future__ import annotations + +import os +from functools import cache + +import numpy as np + +__all__ = [ + "HC_SINKHORN_ENV", + "fused_hc_sinkhorn", + "fused_hc_sinkhorn_enabled", + "hc_sinkhorn_reference", +] + +HC_SINKHORN_ENV = "RAMP_HC_SINKHORN" +_BACKENDS = ("mlx", "metal") + +_SOURCE = """ + uint row = thread_position_in_grid.x; + if (row >= rows[0]) { + return; + } + const uint n = HC * HC; + const uint base = row * n; + const float e = eps[0]; + float m[HC * HC]; + for (uint k = 0; k < n; ++k) { + m[k] = static_cast(logits[base + k]); + } + // Row softmax (max-subtracted), then + eps. + for (uint i = 0; i < HC; ++i) { + float peak = m[i * HC]; + for (uint j = 1; j < HC; ++j) { + peak = metal::max(peak, m[i * HC + j]); + } + float total = 0.0f; + for (uint j = 0; j < HC; ++j) { + float v = metal::precise::exp(m[i * HC + j] - peak); + m[i * HC + j] = v; + total += v; + } + for (uint j = 0; j < HC; ++j) { + m[i * HC + j] = m[i * HC + j] / total + e; + } + } + // Column normalisation. + for (uint j = 0; j < HC; ++j) { + float total = 0.0f; + for (uint i = 0; i < HC; ++i) { + total += m[i * HC + j]; + } + for (uint i = 0; i < HC; ++i) { + m[i * HC + j] = m[i * HC + j] / (total + e); + } + } + // ITERS - 1 alternating row / column normalisations. + for (uint it = 1; it < ITERS; ++it) { + for (uint i = 0; i < HC; ++i) { + float total = 0.0f; + for (uint j = 0; j < HC; ++j) { + total += m[i * HC + j]; + } + for (uint j = 0; j < HC; ++j) { + m[i * HC + j] = m[i * HC + j] / (total + e); + } + } + for (uint j = 0; j < HC; ++j) { + float total = 0.0f; + for (uint i = 0; i < HC; ++i) { + total += m[i * HC + j]; + } + for (uint i = 0; i < HC; ++i) { + m[i * HC + j] = m[i * HC + j] / (total + e); + } + } + } + for (uint k = 0; k < n; ++k) { + comb[base + k] = m[k]; + } +""" + + +def fused_hc_sinkhorn_enabled() -> bool: + """Whether ``RAMP_HC_SINKHORN=metal`` selects the fused kernel.""" + + value = os.environ.get(HC_SINKHORN_ENV, "mlx").strip().lower() or "mlx" + if value not in _BACKENDS: + raise ValueError(f"{HC_SINKHORN_ENV} must be one of {_BACKENDS}, got {value!r}") + return value == "metal" + + +@cache +def _kernel(): + import mlx.core as mx + + return mx.fast.metal_kernel( + name="ramp_hc_sinkhorn", + input_names=["logits", "eps", "rows"], + output_names=["comb"], + source=_SOURCE, + ) + + +def fused_hc_sinkhorn(logits, *, iters: int, eps: float): + """Doubly-stochastic ``[..., hc, hc]`` mixing matrices in one dispatch. + + ``logits`` are the scaled, biased comb logits (float32); ``iters`` is the + model's ``hc_sinkhorn_iters``. Returns float32 of the same shape. + """ + + import mlx.core as mx + + if logits.ndim < 2 or logits.shape[-1] != logits.shape[-2]: + raise ValueError(f"expected [..., hc, hc] logits, got {tuple(logits.shape)}") + if iters < 1: + raise ValueError("iters must be at least 1") + hc = int(logits.shape[-1]) + rows = int(np.prod(logits.shape[:-2])) if logits.ndim > 2 else 1 + flat = mx.contiguous(logits.astype(mx.float32).reshape(rows, hc * hc)) + (out,) = _kernel()( + inputs=[flat, mx.array([eps], dtype=mx.float32), mx.array([rows], dtype=mx.uint32)], + template=[("HC", hc), ("ITERS", int(iters))], + grid=(rows, 1, 1), + threadgroup=(min(rows, 64), 1, 1), + output_shapes=[(rows, hc * hc)], + output_dtypes=[mx.float32], + ) + return out.reshape(logits.shape) + + +def hc_sinkhorn_reference(logits: np.ndarray, *, iters: int, eps: float) -> np.ndarray: + """The model's sinkhorn in float32 NumPy, step for step.""" + + m = np.asarray(logits, dtype=np.float32) + m = np.exp(m - m.max(axis=-1, keepdims=True), dtype=np.float32) + m = m / m.sum(axis=-1, keepdims=True, dtype=np.float32) + np.float32(eps) + m = m / (m.sum(axis=-2, keepdims=True, dtype=np.float32) + np.float32(eps)) + for _ in range(max(int(iters) - 1, 0)): + m = m / (m.sum(axis=-1, keepdims=True, dtype=np.float32) + np.float32(eps)) + m = m / (m.sum(axis=-2, keepdims=True, dtype=np.float32) + np.float32(eps)) + return m.astype(np.float32) diff --git a/src/ramp/models/deepseek_v4_flash_adapter.py b/src/ramp/models/deepseek_v4_flash_adapter.py index d2bffcb..add7193 100644 --- a/src/ramp/models/deepseek_v4_flash_adapter.py +++ b/src/ramp/models/deepseek_v4_flash_adapter.py @@ -149,6 +149,7 @@ repack_fp8_block_as_mxfp8, ) from keep.io.load import load_quantized_vq_switch_linear +from ramp.kernels.hc_sinkhorn import fused_hc_sinkhorn, fused_hc_sinkhorn_enabled from keep.io.source_safetensors import read_safetensors_file_header from ramp.models.glm4_moe_adapter import QuantizedVQSwitchGLU @@ -807,6 +808,34 @@ def _hc_split_sinkhorn_ops( return pre, post, comb +@mx.compile +def _hc_split_logits_ops( + mixes: mx.array, + scale: mx.array, + base: mx.array, + hc_mult: int, + eps: float, +) -> tuple[mx.array, mx.array, mx.array]: + """``_hc_split_sinkhorn_ops`` up to the comb logits, for the fused sinkhorn. + + ``pre`` and ``post`` are computed exactly as there; only the sinkhorn on + the returned logits moves to :func:`ramp.kernels.hc_sinkhorn.fused_hc_sinkhorn`. + """ + + mixes = mixes.astype(mx.float32) + scale = scale.astype(mx.float32) + base = base.astype(mx.float32) + pre_scale, post_scale, comb_scale = scale[0], scale[1], scale[2] + pre = mx.sigmoid(mixes[..., :hc_mult] * pre_scale + base[:hc_mult]) + eps + post = 2 * mx.sigmoid( + mixes[..., hc_mult : 2 * hc_mult] * post_scale + base[hc_mult : 2 * hc_mult] + ) + logits = mixes[..., 2 * hc_mult :].reshape( + *mixes.shape[:-1], hc_mult, hc_mult + ) * comb_scale + base[2 * hc_mult :].reshape(hc_mult, hc_mult) + return pre, post, logits + + class HyperConnection(nn.Module): def __init__(self, config: DeepseekV4FlashVQModelArgs): super().__init__() @@ -824,14 +853,24 @@ def __call__(self, x: mx.array): y = x.astype(mx.float32) z = mx.fast.rms_norm(y.flatten(-2), None, self.norm_eps) mixes = z @ self.fn.T - pre, post, comb = _hc_split_sinkhorn_ops( - mixes, - self.scale, - self.base, - self.hc_mult, - self.sinkhorn_iters, - self.hc_eps, - ) + if fused_hc_sinkhorn_enabled(): + # Opt-in (RAMP_HC_SINKHORN=metal): one dispatch for the whole + # sinkhorn instead of a few per iteration. Not bit-identical. + pre, post, logits = _hc_split_logits_ops( + mixes, self.scale, self.base, self.hc_mult, self.hc_eps + ) + comb = fused_hc_sinkhorn( + logits, iters=self.sinkhorn_iters, eps=self.hc_eps + ) + else: + pre, post, comb = _hc_split_sinkhorn_ops( + mixes, + self.scale, + self.base, + self.hc_mult, + self.sinkhorn_iters, + self.hc_eps, + ) return (pre[..., None] * y).sum(axis=2).astype(x.dtype), post, comb diff --git a/tests/test_hc_sinkhorn.py b/tests/test_hc_sinkhorn.py new file mode 100644 index 0000000..d5ac829 --- /dev/null +++ b/tests/test_hc_sinkhorn.py @@ -0,0 +1,101 @@ +"""The opt-in one-dispatch hyper-connection sinkhorn and its reference.""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") + +from ramp.kernels.hc_sinkhorn import ( # noqa: E402 + HC_SINKHORN_ENV, + fused_hc_sinkhorn, + fused_hc_sinkhorn_enabled, + hc_sinkhorn_reference, +) +from ramp.models.deepseek_v4_flash_adapter import ( # noqa: E402 + _hc_split_logits_ops, + _hc_split_sinkhorn_ops, +) + +HC, ITERS, EPS = 4, 20, 1e-6 + + +def _inputs(rows, seed=0): + rng = np.random.default_rng(seed) + mix = (2 + HC) * HC + mixes = mx.array(rng.normal(size=(1, rows, mix)).astype(np.float32) * 3) + scale = mx.array(rng.uniform(0.5, 2.0, size=(3,)).astype(np.float32)) + base = mx.array(rng.normal(size=(mix,)).astype(np.float32)) + return mixes, scale, base + + +def _metal_available() -> bool: + metal = getattr(mx, "metal", None) + return bool(metal is not None and metal.is_available()) + + +@pytest.mark.parametrize("rows", [1, 7, 256]) +def test_reference_reproduces_the_models_compiled_sinkhorn(rows): + mixes, scale, base = _inputs(rows) + pre, post, comb = _hc_split_sinkhorn_ops(mixes, scale, base, HC, ITERS, EPS) + pre2, post2, logits = _hc_split_logits_ops(mixes, scale, base, HC, EPS) + # pre/post are computed identically on both paths. + assert mx.array_equal(pre, pre2) and mx.array_equal(post, post2) + want = np.array(comb) + got = hc_sinkhorn_reference(np.array(logits), iters=ITERS, eps=EPS) + np.testing.assert_allclose(got, want, rtol=2e-6, atol=1e-7) + # The last step is a column normalisation, so columns sum to one; rows + # only approach it (20 rounds can leave them ~0.3% off for some inputs). + np.testing.assert_allclose(got.sum(axis=-2), 1.0, atol=1e-5) + + +def test_backend_flag(monkeypatch): + monkeypatch.delenv(HC_SINKHORN_ENV, raising=False) + assert fused_hc_sinkhorn_enabled() is False + monkeypatch.setenv(HC_SINKHORN_ENV, "metal") + assert fused_hc_sinkhorn_enabled() is True + monkeypatch.setenv(HC_SINKHORN_ENV, "cuda") + with pytest.raises(ValueError, match=HC_SINKHORN_ENV): + fused_hc_sinkhorn_enabled() + + +def test_fused_kernel_validates_its_input_shape(): + with pytest.raises(ValueError, match="hc, hc"): + fused_hc_sinkhorn(mx.zeros((3, 4, 5)), iters=ITERS, eps=EPS) + with pytest.raises(ValueError, match="iters"): + fused_hc_sinkhorn(mx.zeros((3, 4, 4)), iters=0, eps=EPS) + + +@pytest.mark.skipif(not _metal_available(), reason="needs a Metal GPU") +@pytest.mark.parametrize("rows", [1, 7, 1024]) +def test_fused_kernel_matches_the_compiled_path(rows): + mixes, scale, base = _inputs(rows, seed=rows) + _, _, comb = _hc_split_sinkhorn_ops(mixes, scale, base, HC, ITERS, EPS) + _, _, logits = _hc_split_logits_ops(mixes, scale, base, HC, EPS) + fused = fused_hc_sinkhorn(logits, iters=ITERS, eps=EPS) + assert fused.shape == comb.shape and fused.dtype == mx.float32 + np.testing.assert_allclose(np.array(fused), np.array(comb), rtol=2e-6, atol=1e-7) + + +@pytest.mark.skipif(not _metal_available(), reason="needs a Metal GPU") +def test_hyper_connection_output_with_the_flag(monkeypatch): + from ramp.models.deepseek_v4_flash_adapter import ( + DeepseekV4FlashVQModelArgs, + HyperConnection, + ) + + args = DeepseekV4FlashVQModelArgs(hidden_size=64, hc_mult=HC, num_hidden_layers=1, compress_ratios=[0]) + hc = HyperConnection(args) + rng = np.random.default_rng(1) + hc.fn = mx.array(rng.normal(size=hc.fn.shape).astype(np.float32) * 0.05) + hc.base = mx.array(rng.normal(size=hc.base.shape).astype(np.float32)) + x = mx.array(rng.normal(size=(1, 3, HC, 64)).astype(np.float32)).astype(mx.bfloat16) + monkeypatch.delenv(HC_SINKHORN_ENV, raising=False) + want = hc(x) + monkeypatch.setenv(HC_SINKHORN_ENV, "metal") + got = hc(x) + for a, b in zip(got, want): + np.testing.assert_allclose( + np.array(a.astype(mx.float32)), np.array(b.astype(mx.float32)), rtol=1e-2, atol=1e-6 + ) From 00580c4bf3adbce2a5e8164a89e0bc1ced0c46a9 Mon Sep 17 00:00:00 2001 From: Jack Mazac <114446521+jackmaxwil@users.noreply.github.com> Date: Thu, 5 Nov 2026 05:23:13 +0000 Subject: [PATCH 6/8] ramp: opt-in single-token E8P kernels for 16-bit codes on GPUs without NAX The shipped DeepSeek-V4-Flash artifact uses 16-bit E8P codes. Every M=1 fast path in QuantizedVQSwitchLinear is gated on 8-bit codes, and the E8P tensor-unit kernels are M5-only (RAMP refuses them on an M4 Max, where 36 of 94 kernel tests are wrong). So on M1-M4 every decoded token's routed projections fall through to the scalar kernel, which spends a whole threadgroup per output element. gather_vqmm_m1.metal already had an E8P branch that nothing dispatched to. With RAMP_E8P_M1=1 (opt-in), M=1 calls with 16-bit codes, unsorted routes and the auto/direct strategy use: - e8p_m1_token_layout for gate/up: the existing M=1 kernel with CODE_BITS=16, one activation row against top-k experts; - e8p_m1_per_route for down: the same kernel source with the activation read offset by route. It is built as a separate kernel from a two-line text substitution, so the production 8-bit kernel is byte-identical. A test proves undoing the substitution restores the file. The flag is off by default. Tests pin the decision table and the dispatch wiring, and Metal-only tests compare both kernels with a dense decode of the same codes. benchmarks/bench_e8p_m1.py is the M4 kill test. It times one routed FFN at decode shape against today's dispatch, sweeps row packing and threadgroup codebook staging, and checks the outputs agree. Verdict: PASS at >= 200 GB/s effective, TUNE at 100-200 (next step: decode each codeword once instead of once per dimension), KILL under 100 (serve on the M5). Its synthetic projections reproduce the roadmap's 2,129,920 bytes per projection per expert. --- benchmarks/bench_e8p_m1.py | 211 +++++++++++++++++++++++++++++++++++ src/ramp/kernels/e8p_m1.py | 175 +++++++++++++++++++++++++++++ src/ramp/nn/switch_linear.py | 35 +++++- tests/test_e8p_m1.py | 145 ++++++++++++++++++++++++ 4 files changed, 564 insertions(+), 2 deletions(-) create mode 100644 benchmarks/bench_e8p_m1.py create mode 100644 src/ramp/kernels/e8p_m1.py create mode 100644 tests/test_e8p_m1.py diff --git a/benchmarks/bench_e8p_m1.py b/benchmarks/bench_e8p_m1.py new file mode 100644 index 0000000..5abe0d7 --- /dev/null +++ b/benchmarks/bench_e8p_m1.py @@ -0,0 +1,211 @@ +"""Kill test: single-token E8P (16-bit codes) on a GPU without the NAX kernels. + +Times one routed FFN at decode shape (gate, up and down for top-6 experts) two +ways in the same process: + +- ``production``: today's dispatch (the scalar kernel on M1-M4; NAX on M5); +- ``e8p_m1``: ``RAMP_E8P_M1=1``, the M=1 E8P kernels, swept over row packing + and threadgroup codebook staging. + +Effective bandwidth is the expert bytes one FFN reads (6 experts x 3 +projections of 2-bit codes plus their scales) over the median time. Verdict +for the best ``e8p_m1`` configuration: **PASS** at >= 200 GB/s, **TUNE** at +100-200 (decode the codeword once instead of per dimension, then re-measure), +**KILL** under 100 (serve on the M5 instead). Outputs of the two paths are +compared on every configuration. + +Synthetic codes at DeepSeek-V4-Flash shapes need no download; ``--artifact-dir`` +uses one real layer of the shipped artifact instead:: + + RAMP_E8P_M1=0 uv run python benchmarks/bench_e8p_m1.py --output-json artifacts/benchmarks/e8p-m1-m4max.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import statistics +import sys +import time +from collections.abc import Sequence +from pathlib import Path + +import mlx.core as mx +import numpy as np + +from keep.vq.e8 import e8p_packed_abs_grid +from ramp.kernels import e8p_m1 +from ramp.nn.switch_linear import QuantizedVQSwitchLinear + +PROJECTIONS = ("gate_proj", "up_proj", "down_proj") +PASS_GB_S = 200.0 +KILL_GB_S = 100.0 + + +def synthetic_projection(*, experts: int, out: int, inp: int, group_size: int, seed: int) -> QuantizedVQSwitchLinear: + rng = np.random.default_rng(seed) + codes = rng.integers(0, 1 << 16, size=(experts, out, inp // 8), dtype=np.uint16) + scales = rng.uniform(0.01, 0.05, size=(experts, out, inp // group_size)).astype(np.float16) + return QuantizedVQSwitchLinear( + input_dims=inp, + output_dims=out, + num_experts=experts, + codes=mx.array(codes), + scales=mx.array(scales), + codebook=mx.array(e8p_packed_abs_grid()), + group_size=group_size, + code_bits=16, + ) + + +def projection_bytes_per_expert(layer: QuantizedVQSwitchLinear) -> int: + per_expert = layer.codes.nbytes // layer.num_experts + layer.scales.nbytes // layer.num_experts + return int(per_expert) + + +def _ffn(layers, x, routes, intermediate): + gate = layers["gate_proj"](x, routes) + up = layers["up_proj"](x, routes) + act = (gate * mx.sigmoid(gate) * up).astype(x.dtype) # [1, k, I] + return layers["down_proj"](act.reshape(1, routes.shape[-1], intermediate), routes) + + +def _median_ms(fn, routes_list, iterations, warmup) -> float: + for i in range(warmup): + mx.eval(fn(routes_list[i % len(routes_list)])) + samples = [] + for i in range(iterations): + started = time.perf_counter() + mx.eval(fn(routes_list[i % len(routes_list)])) + samples.append((time.perf_counter() - started) * 1e3) + return statistics.median(samples) + + +def run(layers, *, top_k: int, iterations: int, warmup: int, seed: int, sweep: Sequence[tuple[int, bool]]) -> dict: + rng = np.random.default_rng(seed) + experts = layers["gate_proj"].num_experts + hidden = layers["gate_proj"].input_dims + intermediate = layers["gate_proj"].output_dims + x = mx.array(rng.normal(size=(1, hidden)).astype(np.float32)).astype(mx.float16) + # Fresh routes per call so consecutive calls do not reuse cached experts. + routes_list = [ + mx.array(np.sort(rng.choice(experts, size=top_k, replace=False))[None].astype(np.uint32)) + for _ in range(64) + ] + ffn_bytes = top_k * sum(projection_bytes_per_expert(layers[p]) for p in PROJECTIONS) + + os.environ[e8p_m1.E8P_M1_ENV] = "0" + production_ms = _median_ms(lambda r: _ffn(layers, x, r, intermediate), routes_list, iterations, warmup) + reference = [np.array(_ffn(layers, x, r, intermediate).astype(mx.float32)) for r in routes_list[:4]] + + configs = [] + os.environ[e8p_m1.E8P_M1_ENV] = "1" + original_token, original_route = e8p_m1.e8p_m1_token_layout, e8p_m1.e8p_m1_per_route + import ramp.nn.switch_linear as switch_linear + + try: + for rows_per_tg, tg_codebook in sweep: + switch_linear.e8p_m1_token_layout = ( + lambda *a, _r=rows_per_tg, _c=tg_codebook, **k: original_token( + *a, rows_per_threadgroup=_r, use_threadgroup_codebook=_c, **k + ) + ) + switch_linear.e8p_m1_per_route = ( + lambda *a, _r=rows_per_tg, _c=tg_codebook, **k: original_route( + *a, rows_per_threadgroup=_r, use_threadgroup_codebook=_c, **k + ) + ) + got = [np.array(_ffn(layers, x, r, intermediate).astype(mx.float32)) for r in routes_list[:4]] + diff = max( + float(np.max(np.abs(g - w)) / max(float(np.max(np.abs(w))), 1e-12)) + for g, w in zip(got, reference) + ) + ms = _median_ms(lambda r: _ffn(layers, x, r, intermediate), routes_list, iterations, warmup) + configs.append( + { + "rows_per_threadgroup": rows_per_tg, + "threadgroup_codebook": tg_codebook, + "median_ms_per_ffn": ms, + "effective_gb_per_s": ffn_bytes / (ms / 1e3) / 1e9, + "speedup_over_production": production_ms / ms, + "max_relative_diff_vs_production": diff, + } + ) + finally: + switch_linear.e8p_m1_token_layout = original_token + switch_linear.e8p_m1_per_route = original_route + os.environ.pop(e8p_m1.E8P_M1_ENV, None) + + best = max(configs, key=lambda c: c["effective_gb_per_s"]) + agrees = all(c["max_relative_diff_vs_production"] <= 1e-2 for c in configs) + gb_s = best["effective_gb_per_s"] + verdict = "KILL" if not agrees or gb_s < KILL_GB_S else ("PASS" if gb_s >= PASS_GB_S else "TUNE") + return { + "ffn_bytes": ffn_bytes, + "production_median_ms_per_ffn": production_ms, + "production_effective_gb_per_s": ffn_bytes / (production_ms / 1e3) / 1e9, + "configs": configs, + "best": best, + "outputs_agree": agrees, + "verdict": verdict if mx.default_device() == mx.gpu else "NOT_ON_GPU", + } + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--artifact-dir", type=Path, help="real E8P artifact; default synthetic") + parser.add_argument("--layer", type=int, default=20) + parser.add_argument("--experts", type=int, default=64, help="synthetic experts per projection") + parser.add_argument("--hidden", type=int, default=4096) + parser.add_argument("--intermediate", type=int, default=2048) + parser.add_argument("--group-size", type=int, default=512) + parser.add_argument("--top-k", type=int, default=6) + parser.add_argument("--iterations", type=int, default=200) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--seed", type=int, default=20261105) + parser.add_argument("--sweep", default="8:1,16:1,32:1,64:1,32:0", help="rows_per_threadgroup:tg_codebook,...") + parser.add_argument("--output-json", type=Path) + args = parser.parse_args(argv) + + if args.artifact_dir is not None: + from keep.io.load import load_quantized_vq_switch_linear + + layers = { + p: load_quantized_vq_switch_linear( + args.artifact_dir.expanduser() / f"layer-{args.layer:05d}-{p}.safetensors", + f"model.layers.{args.layer}.ffn.switch_mlp.{p}", + ) + for p in PROJECTIONS + } + if any(layer.code_bits != 16 for layer in layers.values()): + parser.error("the artifact layer is not 16-bit E8P") + else: + layers = { + "gate_proj": synthetic_projection(experts=args.experts, out=args.intermediate, inp=args.hidden, group_size=args.group_size, seed=1), + "up_proj": synthetic_projection(experts=args.experts, out=args.intermediate, inp=args.hidden, group_size=args.group_size, seed=2), + "down_proj": synthetic_projection(experts=args.experts, out=args.hidden, inp=args.intermediate, group_size=args.group_size, seed=3), + } + sweep = [] + for item in args.sweep.split(","): + rows, codebook = item.split(":") + sweep.append((int(rows), codebook.strip() == "1")) + result = { + "record_type": "ramp_e8p_m1_kill_test_v1", + "device": str(mx.default_device()), + "mlx_version": mx.__version__, + "platform": platform.platform(), + "source": "synthetic" if args.artifact_dir is None else str(args.artifact_dir), + **run(layers, top_k=args.top_k, iterations=args.iterations, warmup=args.warmup, seed=args.seed, sweep=sweep), + } + text = json.dumps(result, indent=2, sort_keys=True) + if args.output_json is not None: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(text + "\n") + print(text) + return 1 if result["verdict"] == "KILL" else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/ramp/kernels/e8p_m1.py b/src/ramp/kernels/e8p_m1.py new file mode 100644 index 0000000..1f943a7 --- /dev/null +++ b/src/ramp/kernels/e8p_m1.py @@ -0,0 +1,175 @@ +"""Opt-in single-token (M=1) kernels for 16-bit E8P codes on GPUs without NAX. + +The shipped DeepSeek-V4-Flash artifact uses 16-bit E8P codes. Every M=1 fast +path in ``QuantizedVQSwitchLinear`` is gated on 8-bit codes, and the E8P +tensor-unit kernels are M5-only, so on an M1-M4 GPU a decoded token's routed +projections fall through to the scalar kernel, which spends a whole 256-thread +threadgroup per output element. ``gather_vqmm_m1.metal`` already carries an +E8P branch; nothing dispatched to it. + +With ``RAMP_E8P_M1=1`` (opt-in until ``benchmarks/bench_e8p_m1.py`` passes on +the target Mac), M=1 calls with 16-bit codes use: + +- :func:`e8p_m1_token_layout` for gate/up (one activation row, top-k + experts): the existing M=1 kernel with ``CODE_BITS=16``. +- :func:`e8p_m1_per_route` for down (one activation row *per route*): the + same kernel source with the activation read offset by route. It is built as + a separate kernel from a text substitution, so the production 8-bit kernel + stays byte-identical. + +The E8P branch decodes each codeword once per dimension. If the kill test +lands under its bandwidth bar, decoding a codeword once into eight values is +the first optimisation to try. +""" + +from __future__ import annotations + +import os +from functools import cache +from pathlib import Path + +__all__ = [ + "E8P_M1_ENV", + "e8p_m1_enabled", + "e8p_m1_per_route", + "e8p_m1_token_layout", + "per_route_x_source", + "should_use_e8p_m1", +] + +E8P_M1_ENV = "RAMP_E8P_M1" +_KERNEL_DIR = Path(__file__).resolve().parent +_ON = {"1", "on", "true", "yes"} +_OFF = {"", "0", "off", "false", "no"} +_X_INDEX_REWRITES = ( + ("uint x_base = cw * 8u;", "uint x_base = route * (codewords * 8u) + cw * 8u;"), + ("uint x_idx = cw * 8u + dim;", "uint x_idx = route * (codewords * 8u) + cw * 8u + dim;"), +) + + +def e8p_m1_enabled() -> bool: + value = os.environ.get(E8P_M1_ENV, "").strip().lower() + if value in _ON: + return True + if value in _OFF: + return False + raise ValueError(f"{E8P_M1_ENV} must be on/off (1/0), got {value!r}") + + +def should_use_e8p_m1( + *, code_bits: int, tokens: int, sorted_indices: bool, route_strategy: str +) -> bool: + """The opt-in decision, kept pure so the dispatch table is testable.""" + + return ( + code_bits == 16 + and tokens == 1 + and not sorted_indices + and route_strategy in ("auto", "direct") + and e8p_m1_enabled() + ) + + +def per_route_x_source() -> str: + """``gather_vqmm_m1.metal`` with the activation row chosen by route.""" + + source = (_KERNEL_DIR / "gather_vqmm_m1.metal").read_text() + for old, new in _X_INDEX_REWRITES: + count = source.count(old) + if count != 1: + raise RuntimeError( + f"gather_vqmm_m1.metal changed: expected one {old!r}, found {count}" + ) + source = source.replace(old, new) + return source + + +@cache +def _per_route_kernel(): + import mlx.core as mx + + return mx.fast.metal_kernel( + name="ramp_e8p_m1_per_route_x", + input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], + source=per_route_x_source(), + header=(_KERNEL_DIR / "common_e8.metal").read_text(), + ) + + +def _launch_shape(output_dims: int, rows_per_threadgroup: int) -> tuple[int, int]: + lanes_per_row = 256 // rows_per_threadgroup + if rows_per_threadgroup * lanes_per_row != 256 or lanes_per_row < 1: + raise ValueError("rows_per_threadgroup must divide 256") + return lanes_per_row, (output_dims + rows_per_threadgroup - 1) // rows_per_threadgroup + + +def e8p_m1_token_layout( + x, + codes, + scales, + codebook, + rhs_indices, + *, + output_dims: int, + rows_per_threadgroup: int = 32, + use_threadgroup_codebook: bool = True, +): + """One activation row against ``top_k`` experts: ``[1, top_k, out]``.""" + + from ramp.kernels.gather_vqmm import gather_vqmm_m1_kernel_unchecked + + if x.shape[0] != 1 or rhs_indices.ndim != 2 or rhs_indices.shape[0] != 1: + raise ValueError("token-layout E8P M=1 needs x [1, in] and rhs_indices [1, top_k]") + return gather_vqmm_m1_kernel_unchecked( + x, + codes, + scales, + codebook, + rhs_indices, + output_dims=output_dims, + code_bits=16, + rows_per_threadgroup=rows_per_threadgroup, + use_threadgroup_codebook=use_threadgroup_codebook, + use_decoded_codebook=False, + ) + + +def e8p_m1_per_route( + x, + codes, + scales, + codebook, + rhs_indices, + *, + output_dims: int, + rows_per_threadgroup: int = 32, + use_threadgroup_codebook: bool = True, +): + """Route ``r``'s activation row against expert ``rhs_indices[r]``: ``[routes, out]``.""" + + import mlx.core as mx + + from ramp.kernels.gather_vqmm import _qmv_output_dtype + + routes = int(rhs_indices.size) + if x.ndim != 2 or x.shape[0] != routes: + raise ValueError(f"per-route E8P M=1 needs x [{routes}, in], got {tuple(x.shape)}") + lanes_per_row, row_tiles = _launch_shape(output_dims, rows_per_threadgroup) + output_dtype = _qmv_output_dtype(x.dtype) + (out,) = _per_route_kernel()( + inputs=[mx.contiguous(x), codes, scales, codebook, rhs_indices.reshape(1, routes)], + template=[ + ("ROWS_PER_TG", rows_per_threadgroup), + ("LANES_PER_ROW", lanes_per_row), + ("OUT_T", output_dtype), + ("CODE_BITS", 16), + ("USE_THREADGROUP_CODEBOOK", int(use_threadgroup_codebook)), + ("USE_DECODED_CODEBOOK", 0), + ], + grid=(256, row_tiles, routes), + threadgroup=(256, 1, 1), + output_shapes=[(1, routes, output_dims)], + output_dtypes=[output_dtype], + ) + return out.reshape(routes, output_dims) diff --git a/src/ramp/nn/switch_linear.py b/src/ramp/nn/switch_linear.py index aadaf41..172407f 100644 --- a/src/ramp/nn/switch_linear.py +++ b/src/ramp/nn/switch_linear.py @@ -14,6 +14,7 @@ m1_use_decoded_codebook, m1_use_threadgroup_codebook, ) +from ramp.kernels.e8p_m1 import e8p_m1_per_route, e8p_m1_token_layout, should_use_e8p_m1 from ramp.ops.vq_switch import RouteStrategy from ramp.ops.vq_switch import gather_vqmm, vq_switch_qmv from keep.quant.rotation import apply_rotation_mx, apply_rotation_np, validate_rotation_matrix_mx @@ -332,7 +333,23 @@ def __call__(self, x: mx.array, indices: mx.array, *, sorted_indices: bool = Fal top_k = indices.shape[-1] flat_x = x.reshape((-1, self.input_dims)) flat_indices = indices.reshape((-1, top_k)) - if ( + if should_use_e8p_m1( + code_bits=self.code_bits, + tokens=flat_x.shape[0], + sorted_indices=sorted_indices, + route_strategy=self.route_strategy, + ): + # Opt-in (RAMP_E8P_M1=1): 16-bit codes at M=1 on GPUs without + # the NAX tensor-unit kernels, instead of the scalar fallback. + y = e8p_m1_token_layout( + flat_x, + self.codes, + scales, + self.codebook, + flat_indices, + output_dims=self.output_dims, + ) + elif ( flat_x.shape[0] == 1 and self.code_bits == 8 and not sorted_indices @@ -374,7 +391,21 @@ def __call__(self, x: mx.array, indices: mx.array, *, sorted_indices: bool = Fal route_shape = indices.shape flat_x = x.reshape((-1, self.input_dims)) flat_indices = indices.reshape((-1,)) - if ( + if indices.ndim >= 2 and should_use_e8p_m1( + code_bits=self.code_bits, + tokens=route_shape[0], + sorted_indices=sorted_indices, + route_strategy=self.route_strategy, + ): + y = e8p_m1_per_route( + flat_x, + self.codes, + scales, + self.codebook, + flat_indices, + output_dims=self.output_dims, + ).reshape((*route_shape, self.output_dims)) + elif ( route_shape[0] == 1 and flat_x.shape[0] <= 8 and self.code_bits == 8 diff --git a/tests/test_e8p_m1.py b/tests/test_e8p_m1.py new file mode 100644 index 0000000..b5a6622 --- /dev/null +++ b/tests/test_e8p_m1.py @@ -0,0 +1,145 @@ +"""Opt-in M=1 kernels for 16-bit E8P codes on GPUs without the NAX kernels.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") + +import ramp.nn.switch_linear as switch_linear # noqa: E402 +from ramp.kernels import e8p_m1 # noqa: E402 +from ramp.nn.switch_linear import QuantizedVQSwitchLinear # noqa: E402 + + +def _metal_available() -> bool: + metal = getattr(mx, "metal", None) + return bool(metal is not None and metal.is_available()) + + +@pytest.mark.parametrize( + ("kwargs", "expected"), + [ + ({}, True), + ({"code_bits": 8}, False), + ({"tokens": 2}, False), + ({"sorted_indices": True}, False), + ({"route_strategy": "sorted_tiled"}, False), + ({"route_strategy": "direct"}, True), + ], +) +def test_decision_table_with_the_flag_on(monkeypatch, kwargs, expected): + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "1") + args = {"code_bits": 16, "tokens": 1, "sorted_indices": False, "route_strategy": "auto", **kwargs} + assert e8p_m1.should_use_e8p_m1(**args) is expected + + +def test_flag_is_off_by_default_and_validated(monkeypatch): + monkeypatch.delenv(e8p_m1.E8P_M1_ENV, raising=False) + assert e8p_m1.should_use_e8p_m1(code_bits=16, tokens=1, sorted_indices=False, route_strategy="auto") is False + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "maybe") + with pytest.raises(ValueError, match=e8p_m1.E8P_M1_ENV): + e8p_m1.e8p_m1_enabled() + + +def test_per_route_source_only_offsets_the_activation_reads(): + original = (Path(e8p_m1.__file__).parent / "gather_vqmm_m1.metal").read_text() + rewritten = e8p_m1.per_route_x_source() + assert rewritten.count("route * (codewords * 8u)") == 2 + # Nothing else changed: undoing the two rewrites restores the file exactly. + restored = rewritten + for old, new in e8p_m1._X_INDEX_REWRITES: + restored = restored.replace(new, old) + assert restored == original + + +def _layer(code_bits=16, experts=4, out=32, inp=64, seed=0): + rng = np.random.default_rng(seed) + weight = mx.array(rng.normal(size=(experts, out, inp)).astype(np.float32)) + return QuantizedVQSwitchLinear.from_weights(weight, group_size=32, code_bits=code_bits) + + +def test_switch_linear_routes_m1_16bit_calls_to_the_new_kernels(monkeypatch): + calls = [] + + def fake_token(x, codes, scales, codebook, rhs, *, output_dims): + calls.append(("token", tuple(x.shape), tuple(rhs.shape))) + return mx.zeros((1, rhs.shape[1], output_dims)) + + def fake_route(x, codes, scales, codebook, rhs, *, output_dims): + calls.append(("route", tuple(x.shape), tuple(rhs.shape))) + return mx.zeros((rhs.size, output_dims)) + + monkeypatch.setattr(switch_linear, "e8p_m1_token_layout", fake_token) + monkeypatch.setattr(switch_linear, "e8p_m1_per_route", fake_route) + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "1") + layer = _layer() + indices = mx.array([[0, 2, 3]], dtype=mx.uint32) + y = layer(mx.zeros((1, 64)), indices) # gate/up: one row, three experts + assert y.shape == (1, 3, 32) + y = layer(mx.zeros((1, 3, 64)), indices) # down: one row per route + assert y.shape == (1, 3, 32) + assert calls == [("token", (1, 64), (1, 3)), ("route", (3, 64), (3,))] + + +def test_switch_linear_leaves_8bit_and_flag_off_calls_alone(monkeypatch): + def boom(*args, **kwargs): # pragma: no cover - must not be called + raise AssertionError("opt-in kernel used without the flag or for 8-bit codes") + + monkeypatch.setattr(switch_linear, "e8p_m1_token_layout", boom) + monkeypatch.setattr(switch_linear, "e8p_m1_per_route", boom) + monkeypatch.setattr(switch_linear, "gather_vqmm", lambda *a, **k: mx.zeros((3, 32))) + monkeypatch.setattr(switch_linear, "gather_vqmm_m1_kernel", lambda *a, **k: mx.zeros((1, 3, 32))) + indices = mx.array([[0, 2, 3]], dtype=mx.uint32) + monkeypatch.delenv(e8p_m1.E8P_M1_ENV, raising=False) + _layer(code_bits=16)(mx.zeros((1, 64)), indices) + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "1") + _layer(code_bits=8)(mx.zeros((1, 64)), indices) + + +def _dense(layer) -> np.ndarray: + from keep.vq.e8 import decode_weight_matrix + + codes = np.array(layer.codes) + scales = np.array(layer.effective_scales().astype(mx.float32)) + return np.stack( + [decode_weight_matrix(codes[e], scales[e], code_bits=16) for e in range(codes.shape[0])] + ) + + +@pytest.mark.skipif(not _metal_available(), reason="needs a Metal GPU") +@pytest.mark.parametrize("rows_per_threadgroup", [8, 32]) +def test_token_and_per_route_kernels_match_the_dense_decode(rows_per_threadgroup): + layer = _layer(experts=6, out=96, inp=128, seed=3) + dense = _dense(layer) + rng = np.random.default_rng(4) + rhs = mx.array([[5, 0, 3]], dtype=mx.uint32) + + x = rng.normal(size=(1, 128)).astype(np.float32) + got = e8p_m1.e8p_m1_token_layout( + mx.array(x), layer.codes, layer.effective_scales(), layer.codebook, rhs, + output_dims=96, rows_per_threadgroup=rows_per_threadgroup, + ) + want = np.stack([x[0] @ dense[e].T for e in (5, 0, 3)])[None] + np.testing.assert_allclose(np.array(got.astype(mx.float32)), want, rtol=1e-4, atol=1e-3) + + xr = rng.normal(size=(3, 128)).astype(np.float32) + got = e8p_m1.e8p_m1_per_route( + mx.array(xr), layer.codes, layer.effective_scales(), layer.codebook, rhs.reshape(-1), + output_dims=96, rows_per_threadgroup=rows_per_threadgroup, + ) + want = np.stack([xr[r] @ dense[e].T for r, e in enumerate((5, 0, 3))]) + np.testing.assert_allclose(np.array(got.astype(mx.float32)), want, rtol=1e-4, atol=1e-3) + + +def test_benchmark_synthetic_layers_match_the_shipped_byte_accounting(): + from benchmarks import bench_e8p_m1 as bench + + gate = bench.synthetic_projection(experts=2, out=2048, inp=4096, group_size=512, seed=0) + down = bench.synthetic_projection(experts=2, out=4096, inp=2048, group_size=512, seed=0) + # The roadmap's V4-Flash decode table: 2,129,920 B per projection per expert. + assert bench.projection_bytes_per_expert(gate) == 2_129_920 + assert bench.projection_bytes_per_expert(down) == 2_129_920 + assert gate.code_bits == 16 and gate.codes.dtype == mx.uint16 From 7deda771be1e417f32c5fc9c590ebad4090382ad Mon Sep 17 00:00:00 2001 From: Jack Mazac <114446521+jackmaxwil@users.noreply.github.com> Date: Thu, 5 Nov 2026 05:25:06 +0000 Subject: [PATCH 7/8] docs: Phase 0 kill-test runbook and one-command Mac runner scripts/phase0_kill_tests.sh runs every Phase 0 kill test that needs no model download (mxfp8 residents, fused sinkhorn, single-token E8P, and the Metal parity tests). It refuses to start with custom wired limits set, writes JSON evidence under artifacts/benchmarks/phase0--/ and prints one verdict per test. If KEEP_RESIDENT_DIR and KEEP_VQ_DIR are set, it also runs the real-weight variants. docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.md lists each opt-in switch, its kill criterion and the decision a pass unlocks. It also gives the calibration steps the REAP and residency tests need, which must write into a new output directory because the runner skips existing sessions. --- .../2026-11-05-phase0-kill-tests.md | 70 +++++++++++++++++++ scripts/phase0_kill_tests.sh | 70 +++++++++++++++++++ 2 files changed, 140 insertions(+) create mode 100644 docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.md create mode 100755 scripts/phase0_kill_tests.sh diff --git a/docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.md b/docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.md new file mode 100644 index 0000000..44f303c --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.md @@ -0,0 +1,70 @@ +# Phase 0 kill tests: toward DeepSeek-V4.1-Flash on one 128 GB Mac + +Date: 2026-11-05. Target machine: the M4 Max (128 GB, 409.4 GB/s measured +MLX read). Phase 0 is the work that needs no V4.1 download and can be proven +on the V4-Flash artifact. Every runtime change below is **opt-in**; defaults +are unchanged until the kill test that guards it passes on the target Mac. + +Run everything that needs no model download in one go: + +```bash +scripts/phase0_kill_tests.sh +``` + +It writes JSON evidence to `artifacts/benchmarks/phase0--/`. That +path is gitignored, so force-add it (`git add -f`). + +## What each test decides + +| # | Change | Opt-in switch | Kill test | KILL when | If it passes | +|---|---|---|---|---|---| +| 1 | FP8 residents kept as source bytes (weights-only mxfp8, incl. `wo_a`) | `resident_format="mxfp8"`; `--resident-format mxfp8` on the headline harness | `benchmarks/bench_dsv4_mxfp8_residents.py` | any weight differs from the bf16 decode; mxfp8 < 1.5x bf16 at M=1; > 1.1x affine-8's time | bind the composite with mxfp8 and rerun the headline series | +| 3 | Single-token E8P kernels for 16-bit codes on GPUs without NAX | `RAMP_E8P_M1=1` | `benchmarks/bench_e8p_m1.py` | outputs disagree, or < 100 GB/s effective on one routed FFN | >= 200 GB/s: make it the M4 default. 100-200: decode each codeword once, then re-measure | +| 5 | Headline rows also gate memory coming back in (schema v2) | always on for new rows; v1 rows keep their rule | the next headline series | — | quote only v2 rows | +| 6 | One-dispatch hyper-connection sinkhorn | `RAMP_HC_SINKHORN=metal` | `benchmarks/bench_hc_sinkhorn.py` | < 3x faster at decode shape, or > 2e-6 relative difference | set it for the headline | +| R3 | REAP saliency from the real gate, in the calibration pass | always collected in `calibration` mode | `produce_dsv4_teacher_cache.py reap-select --keep 128` on V4-Flash stats | median pruned-set Jaccard vs route frequency >= 0.9 (REAP adds nothing) | use REAP order for the prune and the hot set | +| 7 | Expert hot set in RAM, rest on SSD | none (simulation only) | `benchmarks/replay_dsv4_expert_residency.py` | best fetching policy misses > 5% of routed slots at the hot-set size you'd ship | worth building the SSD tier | + +Tests R3 and 7 need a calibration run, because the per-session REAP statistics +and recorded routes are new. Rerun the V4-Flash calibration split (5.11 h last +time) into a **new** output directory. The runner resumes by skipping sessions +that already exist, so the old directory would keep sessions without the new +arrays, and `finalize-reap` refuses those. Then: + +```bash +uv run python benchmarks/produce_dsv4_teacher_cache.py finalize-reap --out-dir ~/keep-artifacts/dsv4-teacher-calibration-reap +``` + +```bash +uv run python benchmarks/produce_dsv4_teacher_cache.py reap-select --stats ~/keep-artifacts/dsv4-teacher-calibration-reap/reap-stats.npz --keep 128 +``` + +```bash +uv run python benchmarks/replay_dsv4_expert_residency.py --calibration-dir ~/keep-artifacts/dsv4-teacher-calibration-reap --reap-stats ~/keep-artifacts/dsv4-teacher-calibration-reap/reap-stats.npz --hot 96,128,160 --expert-mb 6.39 +``` + +The replay's default `--expert-mb` is V4.1's 8.99 MB. A V4-Flash expert is +6.39 MB (3 x 2,129,920 B). + +## What was verified where + +On Linux with CPU MLX 0.31.2: + +- the mxfp8 byte copy dequantizes bit-identically at all eight V4-Flash + resident shapes, and `mx.quantize(mode="mxfp8")` does not (608 of 98,304 + values changed); +- REAP statistics match a float64 reference, and slicing the router equals + masking it; +- the residency replay reproduces hand-counted LRU and decayed-count misses; +- the sinkhorn reference reproduces the model's compiled path. + +Everything Metal (kernel parity, every speed number) is untested until the +script above runs on the Mac. + +## Known gaps + +- `benchmarks/bench_dsv4_mtp_headline.py` validates speculative rows against + `EXPECTED_VERIFY_IMPLEMENTATION = "nax_e8p_m32n64"`, so MTP headline rows + cannot validate on an M4 until the verify window has a non-NAX E8P kernel. +- The E8P M=1 branch decodes each codeword once per dimension. Decoding once + per codeword is the first optimisation if the kill test lands in TUNE. diff --git a/scripts/phase0_kill_tests.sh b/scripts/phase0_kill_tests.sh new file mode 100755 index 0000000..d75f554 --- /dev/null +++ b/scripts/phase0_kill_tests.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash +# Phase 0 kill tests that need no model download. Run on the target Mac: +# +# scripts/phase0_kill_tests.sh # writes artifacts/benchmarks/phase0--/ +# +# Each benchmark exits 1 on KILL. The script runs them all, then prints one +# verdict line per test. Optional real-weight variants, if the files exist: +# KEEP_RESIDENT_DIR=~/keep-artifacts/dsv4-residents-release-precision-20260819 +# KEEP_VQ_DIR=~/keep-artifacts/dsv4-vq-e8p-g512 +set -uo pipefail + +for name in GLM_MLX_WIRED_LIMIT_GB GLM_SINGLE_HOST_MLX_WIRED_LIMIT_GB; do + if [[ -n "${!name:-}" ]]; then + echo "unset $name first: kill tests run under the default wired-memory limits" >&2 + exit 2 + fi +done + +repo="$(cd "$(dirname "$0")/.." && pwd)" +cd "$repo" +out="artifacts/benchmarks/phase0-$(hostname -s)-$(date +%Y%m%d-%H%M%S)" +mkdir -p "$out" +status=0 + +run() { + local label="$1"; shift + echo "== $label" + if uv run python "$@"; then echo "$label: ok"; else echo "$label: KILL or error (exit $?)"; status=1; fi +} + +run "mxfp8 residents (synthetic V4-Flash shapes)" \ + benchmarks/bench_dsv4_mxfp8_residents.py --synthetic --output-json "$out/mxfp8-residents-synthetic.json" +if [[ -n "${KEEP_RESIDENT_DIR:-}" && -d "${KEEP_RESIDENT_DIR}" ]]; then + run "mxfp8 residents (real layer 20)" \ + benchmarks/bench_dsv4_mxfp8_residents.py --checkpoint "$KEEP_RESIDENT_DIR" --layer 20 \ + --output-json "$out/mxfp8-residents-layer20.json" +fi +run "fused hyper-connection sinkhorn" \ + benchmarks/bench_hc_sinkhorn.py --output-json "$out/hc-sinkhorn.json" +run "E8P single-token kernels (synthetic)" \ + benchmarks/bench_e8p_m1.py --output-json "$out/e8p-m1-synthetic.json" +if [[ -n "${KEEP_VQ_DIR:-}" && -d "${KEEP_VQ_DIR}" ]]; then + run "E8P single-token kernels (real layer 20)" \ + benchmarks/bench_e8p_m1.py --artifact-dir "$KEEP_VQ_DIR" --layer 20 \ + --output-json "$out/e8p-m1-layer20.json" +fi + +echo "== Metal parity tests" +if uv run --group dev python -m pytest -q \ + tests/test_hc_sinkhorn.py tests/test_e8p_m1.py tests/test_fp8_block_dequant.py \ + tests/test_bench_dsv4_mxfp8_residents.py "tests/test_deepseek_v4_flash_adapter.py" -k "mxfp8 or sinkhorn or e8p or fp8" \ + | tee "$out/metal-parity-tests.txt"; then + echo "Metal parity tests: ok" +else + echo "Metal parity tests: FAILED"; status=1 +fi + +echo +echo "Verdicts:" +for f in "$out"/*.json; do + uv run python - "$f" <<'PY' +import json, sys +d = json.load(open(sys.argv[1])) +v = d.get("verdict") +v = v.get("verdict") if isinstance(v, dict) else v +print(f" {sys.argv[1].rsplit('/', 1)[-1]}: {v}") +PY +done +echo "Evidence: $out (force-add it: git add -f $out)" +exit $status From b392eb3165419e29e314d3463274baa49279077e Mon Sep 17 00:00:00 2001 From: Jack Mazac <114446521+jackmaxwil@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:29:48 +0000 Subject: [PATCH 8/8] docs: V4.1 plan and Phase 0 closeout status Records the ranked plan for running DeepSeek-V4.1-Flash on one 128 GB Mac, what Phase 0 built and verified on this branch, the next actions in order (Mac kill tests, REAP calibration rerun, CI fix), and the open decisions. --- .../2026-11-05-v41-plan-and-phase0-status.md | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 docs/deepseek-v4-flash/2026-11-05-v41-plan-and-phase0-status.md diff --git a/docs/deepseek-v4-flash/2026-11-05-v41-plan-and-phase0-status.md b/docs/deepseek-v4-flash/2026-11-05-v41-plan-and-phase0-status.md new file mode 100644 index 0000000..e433c8a --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-05-v41-plan-and-phase0-status.md @@ -0,0 +1,95 @@ +# DeepSeek-V4.1-Flash on one 128 GB Mac: plan and Phase 0 status + +Date: 2026-11-05. Branch `jack/hopeful-cori-iszsfq`, draft PR jackmaxwil/keep#3. +Target machine: M4 Max, 128 GB, 409.4 GB/s measured MLX read. + +## The plan in brief + +This synthesises three earlier analyses (oMLX's oQ quantization, oMLX's V4.1 +Engram/offload work, and REAP pruning) and checks them against this repo. + +- **Memory and speed come from different weights.** Each routed expert kept in + every layer costs 0.359 GB of RAM (V4.1, 2.031 bpw) but adds no per-token + reads. Resident weights and kernels set decode speed. So: prune to fit, fix + residents and kernels to go fast, and give MTP whatever memory is left. +- **Speed is a kernel problem first.** The measured V4-Flash decode floor is + 270–390 ms/token (roadmap). On an M4 Max, 16-bit E8P experts ran on the + scalar kernel, because NAX is M5-only and the M=1 fast paths were 8-bit-only. +- **Fit forces a ~50% prune**, about 176–192 of 384 experts per layer. That is + where REAP's published results lose most outside code (Kimi-K2 at 50%: −17.6% + multiple choice, −54% multi-turn tool use). KEEP's quality gate is in-domain, + so an out-of-domain teacher split must decide the expert count. +- **Plausible end state (estimate):** ~31 tok/s autoregressive and 37–53 with + MTP, at 81–87 GB in flight. Engram's 203 GB stays on SSD. + +### The top 20, by execution order + +Required for V4.1 to exist at all: +- **R1:** port V4.1 with reference parity. +- **R2:** read Engram from SSD. +- **R3:** REAP done right (real-gate saliency plus router slicing). + +Ranked levers: +1. FP8 residents as mxfp8 byte copies +2. Out-of-domain teacher split to choose the calibration mix and N +3. Non-NAX 16-bit E8P kernels +4. Standard benchmarks compared with oMLX and GGUF builds +5. Quiet-window harness that gates on inbound paging +6. Decode-floor fusion (sinkhorn, descriptors, compile) +7. REAP-ordered fast/exact residency (SSD tier) +8. MTP stack +9. Recovery distillation +10. Ship V4-Flash first +11. Persistent prefix-KV cache +12. Prune-vs-bits frontier +13. Resident trims (head and embedding to 8 bits) +14. Two-Mac RDMA calibration +15. 4-bit attention "turbo" +16. 4-bit Engram tables +17. Fit polish + +## Phase 0: done in this branch (no V4.1 download needed) + +| Plan item | Commit | Status | +|---|---|---| +| 1 mxfp8 residents (incl. `wo_a`) | `3c6d9f1` | Built, opt-in. Bit-exact on CPU at all 8 V4-Flash shapes. **Speed untested** | +| 5 inbound-paging gate (row schema v2) | `450c925` | Built. Unit-tested | +| R3 REAP saliency + router slicing | `bd7fd54` | Built. Matches a float64 reference; slicing equals masking | +| 7 residency replay (SSD-tier kill test) | `33a5681` | Built. Hand-checked miss counts | +| 6 fused sinkhorn | `5baed30` | Built, opt-in. **Metal untested** | +| 3 single-token E8P on M1–M4 | `00580c4` | Built, opt-in. **Metal untested** | +| Kill-test runner and runbook | `7deda77` | `scripts/phase0_kill_tests.sh`, `2026-11-05-phase0-kill-tests.md` | + +Regression check (Linux, CPU MLX): on the 33 test files that import changed +modules, the failure set is identical to `main`'s (308 failures needing Metal +or Mac-local files), plus 22 newly passing tests. + +## Next actions, in order + +1. On the M4 Max, run `scripts/phase0_kill_tests.sh`. Each test is PASS, KILL + or TUNE per the runbook. +2. Rerun V4-Flash calibration into a **new** directory, then run + `finalize-reap`, `reap-select --keep 128` and the residency replay. +3. If mxfp8 passes, bind the composite with `--resident-format mxfp8` and rerun + the headline series. Quote v2 rows only. +4. Fix CI. `validate-recipes` fails on `main` because strict recipes refuse + `--no-hash`, so the pytest step never runs. The proposed workflow patch is + in the PR description. +5. Then Phase 1 (V4.1 port, Engram, two-Mac calibration), which waits on your + decisions below. + +## Open decisions (yours) + +- **Serving machine.** M4 Max (this plan) or M5 Max (decision D2, 2026-10-06). +- **Scope.** Coding-agent specialist or generalist. This sets the calibration + mix and the prune depth. +- **Disk and download.** The V4.1 source is 510 GB; the M4 had 73 GiB free on + Oct 6. +- **Exact mode.** In v1 or later. It adds 69 GB to the download. + +## Known gaps + +- MTP headline rows require the M5-only `nax_e8p_m32n64`, so they won't + validate on an M4. +- The E8P M=1 branch decodes each codeword once per dimension. Decoding once + per codeword is the next optimisation if its verdict is TUNE.