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/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/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/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/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/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/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/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. 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 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(" list[Any]: """Prefill the whole backbone layer-outer, chunk-inner. @@ -1182,6 +1184,15 @@ 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. + + ``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 @@ -1284,7 +1295,7 @@ def down_hook(activated, expert_ids, order, _capture=capture): ) mask = masks[chunk_index] - if accumulator 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]) @@ -1316,6 +1327,18 @@ 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) + for observer in route_observers: + observer.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]) + for observer in route_observers: + observer.eval_pending() + continue if collect_hidden: rows = moe_input.reshape(-1, moe_input.shape[-1]).astype(mx.float32) accumulator.add_chunk( @@ -1344,6 +1367,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): ) mx.eval(hidden[chunk_index]) accumulator.eval_pending() + for observer in route_observers: + observer.eval_pending() switch_mlp.down_hook = None layer.ffn.switch_mlp = None @@ -1352,6 +1377,8 @@ def down_hook(activated, expert_ids, order, _capture=capture): expert_provider.release(layer_index) if accumulator is not None: accumulator.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) @@ -1933,6 +1960,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 +2705,20 @@ def _produce_sessions( started = time.perf_counter() 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, @@ -2696,10 +2748,15 @@ def on_layer(layer_index: int, n_chunks: int, _s=session, _i=index) -> None: collect_spaces=collect, on_layer=on_layer, dspark_taps=taps, + route_observers=observers, ) if mode == "calibration": - arrays = calibration_session_arrays(accumulator) + 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 ) @@ -2797,6 +2854,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/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/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/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/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/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 9b6b72a..add7193 100644 --- a/src/ramp/models/deepseek_v4_flash_adapter.py +++ b/src/ramp/models/deepseek_v4_flash_adapter.py @@ -139,12 +139,17 @@ scaled_dot_product_attention, ) from mlx_lm.models.cache import CacheList, RotatingKVCache, _BaseCache -from mlx_lm.models.mla import MultiLinear +from mlx_lm.models.mla import MultiLinear, QuantizedMultiLinear from mlx_lm.models.switch_layers import SwitchGLU from keep.convert.fp4_expert import dequantize_fp4_expert -from keep.convert.fp8_block import dequantize_fp8_block +from keep.convert.fp8_block import ( + MXFP8_GROUP_SIZE, + dequantize_fp8_block, + 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 @@ -803,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__() @@ -820,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 @@ -2107,7 +2150,7 @@ def append_context( f"DSpark context cache count {len(cache)} != {len(self.blocks)} stages" ) entry = self.blocks[0] - expected = entry.main_proj.weight.shape[1] + expected = _linear_input_dims(entry.main_proj) if main_hidden.shape[-1] != expected: raise ValueError( "DSpark context width must be hidden_size * " @@ -2790,6 +2833,13 @@ class DeepseekV4FlashNonVQBindReport: skipped_out_of_scope_tensors: tuple[str, ...] = () skipped_unmatched_tensors: tuple[str, ...] = () missing_model_parameters: tuple[str, ...] = () + #: ``"bf16"`` (decode FP8 to ``compute_dtype``) or ``"mxfp8"`` (byte copy). + resident_format: str = "bf16" + #: Source FP8 tensors bound as weights-only mxfp8 instead of being decoded. + mxfp8_repacked_tensors: tuple[str, ...] = () + #: The ``.scales`` parameters those bindings created. They are bookkeeping + #: for the byte copy, not logical tensors, so ``bound_count`` excludes them. + mxfp8_scale_parameters: tuple[str, ...] = () @property def bound_count(self) -> int: @@ -2822,6 +2872,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 +3038,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 +3153,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 +3162,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 +3252,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 +3298,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 +3386,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 +3401,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/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_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} 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_dsv4_teacher_runner.py b/tests/test_dsv4_teacher_runner.py index 6b725b7..664d8f2 100644 --- a/tests/test_dsv4_teacher_runner.py +++ b/tests/test_dsv4_teacher_runner.py @@ -1781,6 +1781,132 @@ 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, + route_observers=() if reap is None else (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 + + # 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] + 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_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 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"]) 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(" 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 + ) 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())