diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/chain.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/chain.py new file mode 100644 index 0000000..569887b --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/chain.py @@ -0,0 +1,29 @@ +import os, sys, time +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection, _ffn, projection_bytes_per_expert, PROJECTIONS +from ramp.kernels import e8p_m1 +layers = { + "gate_proj": synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1), + "up_proj": synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=2), + "down_proj": synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3), +} +rng = np.random.default_rng(0) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +ffn_bytes = 6 * sum(projection_bytes_per_expert(layers[p]) for p in PROJECTIONS) +x0 = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) + +def chain(n): + x = x0 + for i in range(n): # dependent, like consecutive layers in one decode step + y = _ffn(layers, x, routes[i % 64], 2048) + x = (x + y.sum(axis=0, keepdims=True) * 1e-3).astype(mx.float16) + return x + +for flag in ("0", "1"): + os.environ[e8p_m1.E8P_M1_ENV] = flag + for n in (1, 43): + for _ in range(3): mx.eval(chain(n)) + best = min((lambda a: (mx.eval(chain(n)), time.perf_counter() - a)[1])(time.perf_counter()) for _ in range(30)) + ms = best / n * 1e3 + print(f"E8P_M1={flag} ffns/eval={n:2d} {ms:.4f} ms/ffn {ffn_bytes / (ms / 1e3) / 1e9:.1f} GB/s") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/e8p_m1.decode-once.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/e8p_m1.decode-once.metal new file mode 100644 index 0000000..9e1aac9 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/e8p_m1.decode-once.metal @@ -0,0 +1,100 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes, decoding each +// codeword once into eight values (the generic M=1 kernel decodes it once per +// dimension, with an 8-step parity loop each time). +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble shuffle(d) = {0,4,1,5,2,6,3,7}[d], so nibbles 0-3 +// pair with the even activations and nibbles 4-7 with the odd ones. Bit p of +// the effective signs negates nibble p, and every value gets +-0.25 by parity. +// +// Each lane loads CHUNK (4 or 8) codewords at once; 2-byte loads leave the +// matvec latency-bound. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); + +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +uint row_slot = lane / uint(LANES_PER_ROW); +uint row_lane = lane - row_slot * uint(LANES_PER_ROW); +uint row = row_tile * uint(ROWS_PER_TG) + row_slot; + +threadgroup uint local_codebook[256]; +if (uint(USE_THREADGROUP_CODEBOOK) != 0u) { + for (uint idx = lane; idx < 256u; idx += threads_per_threadgroup.x) { + local_codebook[idx] = codebook[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +float acc = 0.0f; +if (row_slot < uint(ROWS_PER_TG) && row < out_dim) { + uint code_base = (expert * out_dim + row) * codewords; + uint scale_base = (expert * out_dim + row) * groups; + // One wide load of CHUNK codewords per lane per step; the host checks + // that codewords per group is a multiple of CHUNK. + uint chunks = codewords / uint(CHUNK); + for (uint chunk = row_lane; chunk < chunks; chunk += uint(LANES_PER_ROW)) { + uint cw0 = chunk * uint(CHUNK); + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + code_base))[chunk]; + } else { + words.xy = ((const device uint2*)(codes + code_base))[chunk]; + } + float chunk_acc = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint abs_code = uint(USE_THREADGROUP_CODEBOOK) != 0u + ? local_codebook[codeword >> 8u] + : codebook[codeword >> 8u]; + float4 v_lo = (float4((uint4(abs_code) >> uint4(0u, 4u, 8u, 12u)) & 0xFu) - 8.0f) * 0.5f; + float4 v_hi = (float4((uint4(abs_code) >> uint4(16u, 20u, 24u, 28u)) & 0xFu) - 8.0f) * 0.5f; + v_lo *= 1.0f - 2.0f * float4((uint4(effective) >> uint4(0u, 1u, 2u, 3u)) & 1u); + v_hi *= 1.0f - 2.0f * float4((uint4(effective) >> uint4(4u, 5u, 6u, 7u)) & 1u); + float offset = parity != 0u ? -0.25f : 0.25f; + uint xb = x_row + (cw0 + i) * 8u; + float4 x_even = float4(float(x[xb]), float(x[xb + 2u]), float(x[xb + 4u]), float(x[xb + 6u])); + float4 x_odd = float4(float(x[xb + 1u]), float(x[xb + 3u]), float(x[xb + 5u]), float(x[xb + 7u])); + chunk_acc += dot(v_lo + offset, x_even) + dot(v_hi + offset, x_odd); + } + acc += chunk_acc * float(scales[scale_base + cw0 / codewords_per_group]); + } +} + +if (uint(LANES_PER_ROW) <= 32u) { + float total = acc; + for (uint stride = uint(LANES_PER_ROW) >> 1; stride > 0; stride >>= 1) { + total += simd_shuffle_down(total, stride); + } + if (row_lane == 0 && row_slot < uint(ROWS_PER_TG) && row < out_dim) { + out[(route * out_dim) + row] = OUT_T(total); + } +} else { + threadgroup float partials[256]; + partials[lane] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = uint(LANES_PER_ROW) >> 1; stride > 0; stride >>= 1) { + if (row_lane < stride) { + partials[lane] += partials[lane + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (row_lane == 0 && row_slot < uint(ROWS_PER_TG) && row < out_dim) { + out[(route * out_dim) + row] = OUT_T(partials[lane]); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/kbench.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/kbench.py new file mode 100644 index 0000000..3a0d29e --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-decode-once/kbench.py @@ -0,0 +1,38 @@ +import time, statistics, sys +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import gather_vqmm_m1_kernel_unchecked + +L = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +rng = np.random.default_rng(0) +x = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 +sc = L.scales + +def old(r, rows, tg): + return gather_vqmm_m1_kernel_unchecked(x, L.codes, sc, L.codebook, r, output_dims=2048, code_bits=16, + rows_per_threadgroup=rows, use_threadgroup_codebook=tg, use_decoded_codebook=False) +def new(r, rows, tg): + return e8p_m1.e8p_m1_token_layout(x, L.codes, sc, L.codebook, r, output_dims=2048, + rows_per_threadgroup=rows, use_threadgroup_codebook=tg) +def prod(r, rows, tg): + return L(x, r) + +def once(fn, rows, tg): + a = time.perf_counter(); mx.eval([fn(r, rows, tg) for r in routes * 4]); return (time.perf_counter() - a) / (len(routes) * 4) + +cases = [("prod", prod, 0)] + [(n, f, r) for r in (4, 8, 16, 32) for n, f in (("old", old), ("new", new))] +for n, f, r in cases: + for _ in range(3): once(f, r, True) +best = {c[0] + str(c[2]): 1e9 for c in cases} +for _ in range(20): + for n, f, r in cases: + best[n + str(r)] = min(best[n + str(r)], once(f, r, True)) +ref = np.array(prod(routes[0], 0, 0).astype(mx.float32)) +for n, f, r in cases: + ms = best[n + str(r)] * 1e3 + d = np.max(np.abs(np.array(f(routes[0], r, True).astype(mx.float32)) - ref)) / np.max(np.abs(ref)) + print("%s rows=%2d %.4f ms %.1f GB/s diff %.1e" % (n, r, ms, nbytes / (ms / 1e3) / 1e9, d)) diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-run1.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-run1.json new file mode 100644 index 0000000..b48c857 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-run1.json @@ -0,0 +1,63 @@ +{ + "best": { + "effective_gb_per_s": 126.24687513292753, + "max_relative_diff_vs_production": 0.000638977624475956, + "median_ms_per_ffn": 0.30367927887032975, + "rows_per_threadgroup": 16, + "speedup_over_production": 2.4872253552301857, + "threadgroup_codebook": true + }, + "configs": [ + { + "effective_gb_per_s": 125.61180414014768, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.3052146274184938, + "rows_per_threadgroup": 8, + "speedup_over_production": 2.4747136421762974, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 126.24687513292753, + "max_relative_diff_vs_production": 0.000638977624475956, + "median_ms_per_ffn": 0.30367927887032975, + "rows_per_threadgroup": 16, + "speedup_over_production": 2.4872253552301857, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 124.18555893060979, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.3087199536736968, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.446614782349445, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 114.7250308046974, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.334177813952962, + "rows_per_threadgroup": 64, + "speedup_over_production": 2.260230244879809, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 120.24603723202925, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.31883429078017034, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.369001152341795, + "threadgroup_codebook": false + } + ], + "device": "Device(gpu, 0)", + "ffn_bytes": 38338560, + "ffns_per_sync": 43, + "mlx_version": "0.31.2", + "outputs_agree": true, + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "production_effective_gb_per_s": 50.758116817783794, + "production_median_ms_per_ffn": 0.7553188022643025, + "record_type": "ramp_e8p_m1_kill_test_v2", + "source": "synthetic", + "verdict": "TUNE" +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-run2.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-run2.json new file mode 100644 index 0000000..33bdef8 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-run2.json @@ -0,0 +1,63 @@ +{ + "best": { + "effective_gb_per_s": 123.71473313984096, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.30989486075731987, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.4497149387742128, + "threadgroup_codebook": true + }, + "configs": [ + { + "effective_gb_per_s": 122.3088833198184, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.3134568721370035, + "rows_per_threadgroup": 8, + "speedup_over_production": 2.421877257534667, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 121.26227485943866, + "max_relative_diff_vs_production": 0.000638977624475956, + "median_ms_per_ffn": 0.3161623022860176, + "rows_per_threadgroup": 16, + "speedup_over_production": 2.4011530291798957, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 123.71473313984096, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.30989486075731987, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.4497149387742128, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 114.24601252337658, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.33557897692188876, + "rows_per_threadgroup": 64, + "speedup_over_production": 2.2622217780444154, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 120.92889489763162, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.3170339068462856, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.394551666092416, + "threadgroup_codebook": false + } + ], + "device": "Device(gpu, 0)", + "ffn_bytes": 38338560, + "ffns_per_sync": 43, + "mlx_version": "0.31.2", + "outputs_agree": true, + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "production_effective_gb_per_s": 50.50168539272789, + "production_median_ms_per_ffn": 0.759154069846561, + "record_type": "ramp_e8p_m1_kill_test_v2", + "source": "synthetic", + "verdict": "TUNE" +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-tuned-run1.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-tuned-run1.json new file mode 100644 index 0000000..cb42a74 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-tuned-run1.json @@ -0,0 +1,63 @@ +{ + "best": { + "effective_gb_per_s": 205.49661679791186, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.1865654072431891, + "rows_per_simdgroup": 4, + "simdgroups": 1, + "speedup_over_production": 3.969571859025884 + }, + "configs": [ + { + "effective_gb_per_s": 200.36558927025263, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.1913430351969721, + "rows_per_simdgroup": 8, + "simdgroups": 1, + "speedup_over_production": 3.870455957270122 + }, + { + "effective_gb_per_s": 205.49661679791186, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.1865654072431891, + "rows_per_simdgroup": 4, + "simdgroups": 1, + "speedup_over_production": 3.969571859025884 + }, + { + "effective_gb_per_s": 201.96111349692686, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.18983139544129807, + "rows_per_simdgroup": 4, + "simdgroups": 4, + "speedup_over_production": 3.901276649937917 + }, + { + "effective_gb_per_s": 201.68983785241895, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.19008672131539517, + "rows_per_simdgroup": 2, + "simdgroups": 4, + "speedup_over_production": 3.8960364266132834 + }, + { + "effective_gb_per_s": 192.8922985262998, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.19875630231433397, + "rows_per_simdgroup": 16, + "simdgroups": 1, + "speedup_over_production": 3.7260946286323495 + } + ], + "device": "Device(gpu, 0)", + "ffn_bytes": 38338560, + "ffns_per_sync": 43, + "mlx_version": "0.31.2", + "outputs_agree": true, + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "production_effective_gb_per_s": 51.76795485655721, + "production_median_ms_per_ffn": 0.7405847904602673, + "record_type": "ramp_e8p_m1_kill_test_v2", + "source": "synthetic", + "verdict": "PASS" +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-tuned-run2.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-tuned-run2.json new file mode 100644 index 0000000..5a09427 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic-v2-tuned-run2.json @@ -0,0 +1,63 @@ +{ + "best": { + "effective_gb_per_s": 201.9003166348455, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.18988855807164812, + "rows_per_simdgroup": 4, + "simdgroups": 1, + "speedup_over_production": 3.9789912891068244 + }, + "configs": [ + { + "effective_gb_per_s": 196.16447165307085, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.1954408954736928, + "rows_per_simdgroup": 8, + "simdgroups": 1, + "speedup_over_production": 3.865950965057083 + }, + { + "effective_gb_per_s": 201.9003166348455, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.18988855807164812, + "rows_per_simdgroup": 4, + "simdgroups": 1, + "speedup_over_production": 3.9789912891068244 + }, + { + "effective_gb_per_s": 198.43019765127468, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.19320930208101175, + "rows_per_simdgroup": 4, + "simdgroups": 4, + "speedup_over_production": 3.9106032180134807 + }, + { + "effective_gb_per_s": 198.12068741620567, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.1935111395987617, + "rows_per_simdgroup": 2, + "simdgroups": 4, + "speedup_over_production": 3.904503482511547 + }, + { + "effective_gb_per_s": 189.32440694713122, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.20250194160495127, + "rows_per_simdgroup": 16, + "simdgroups": 1, + "speedup_over_production": 3.7311490076580545 + } + ], + "device": "Device(gpu, 0)", + "ffn_bytes": 38338560, + "ffns_per_sync": 43, + "mlx_version": "0.31.2", + "outputs_agree": true, + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "production_effective_gb_per_s": 50.74158297043335, + "production_median_ms_per_ffn": 0.7555649184681432, + "record_type": "ramp_e8p_m1_kill_test_v2", + "source": "synthetic", + "verdict": "PASS" +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic.json new file mode 100644 index 0000000..6715e39 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-m1-synthetic.json @@ -0,0 +1,62 @@ +{ + "best": { + "effective_gb_per_s": 52.75492459836312, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.7267295004567131, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.3768878351187355, + "threadgroup_codebook": true + }, + "configs": [ + { + "effective_gb_per_s": 26.706688923324624, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 1.4355414896272123, + "rows_per_threadgroup": 8, + "speedup_over_production": 1.2032773148939389, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 37.284508678850266, + "max_relative_diff_vs_production": 0.000638977624475956, + "median_ms_per_ffn": 1.0282704897690564, + "rows_per_threadgroup": 16, + "speedup_over_production": 1.6798639329282228, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 52.75492459836312, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.7267295004567131, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.3768878351187355, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 44.31135339640817, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.8652085089124739, + "rows_per_threadgroup": 64, + "speedup_over_production": 1.9964603806644017, + "threadgroup_codebook": true + }, + { + "effective_gb_per_s": 45.07216109711928, + "max_relative_diff_vs_production": 0.000643500650767237, + "median_ms_per_ffn": 0.8506039885105565, + "rows_per_threadgroup": 32, + "speedup_over_production": 2.0307387837225495, + "threadgroup_codebook": false + } + ], + "device": "Device(gpu, 0)", + "ffn_bytes": 38338560, + "mlx_version": "0.31.2", + "outputs_agree": true, + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "production_effective_gb_per_s": 22.194957548650077, + "production_median_ms_per_ffn": 1.7273545090574771, + "record_type": "ramp_e8p_m1_kill_test_v1", + "source": "synthetic", + "verdict": "KILL" +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/cpu.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/cpu.py new file mode 100644 index 0000000..b81c57b --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/cpu.py @@ -0,0 +1,23 @@ +import sys, time +import numpy as np, mlx.core as mx +S = sys.argv[1] +src = open(f"{S}/mr.py").read(); src = src[: src.index("rng = np.random.default_rng(0)")] +ns = {}; sys.argv = [sys.argv[0], S]; exec(src, ns) +from benchmarks.bench_e8p_m1 import synthetic_projection +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +x = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +one = [mx.array(np.array([[i]], dtype=np.uint32)) for i in range(64)] +f = lambda r: ns["new"](x, gate, r, 0, 4, 1, 8, 1) +for _ in range(3): mx.eval([f(r) for r in routes * 4]) +for label, rs in (("6 routes", routes), ("1 route ", one)): + best_build = best_total = 1e9 + for _ in range(20): + a = time.perf_counter(); outs = [f(r) for r in rs * 4]; b = time.perf_counter(); mx.eval(outs); c = time.perf_counter() + best_build = min(best_build, (b - a) / 256); best_total = min(best_total, (c - a) / 256) + print(f"{label}: build {best_build*1e6:.1f} us/call, build+eval {best_total*1e6:.1f} us/call") +# matmul build overhead for comparison +w = mx.random.normal((1560, 4096)).astype(mx.bfloat16); xb = x.astype(mx.bfloat16); mx.eval(w, xb) +a = time.perf_counter(); outs = [xb @ w.T for _ in range(256)]; b = time.perf_counter(); mx.eval(outs); c = time.perf_counter() +print(f"bf16 matmul: build {(b-a)/256*1e6:.1f} us/call, build+eval {(c-a)/256*1e6:.1f} us/call") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr.metal new file mode 100644 index 0000000..f5ccbd2 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr.metal @@ -0,0 +1,87 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +threadgroup uint local_codebook[256]; +if (uint(TG_CODEBOOK) != 0u) { + for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + local_codebook[idx] = codebook[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float xv[CHUNK * 8]; + for (uint i = 0; i < uint(CHUNK) * 8u; ++i) { + xv[i] = float(x[x_row + cw0 * 8u + i]); + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint abs_code = uint(TG_CODEBOOK) != 0u ? local_codebook[codeword >> 8u] : codebook[codeword >> 8u]; + float offset = parity != 0u ? -0.25f : 0.25f; + const uint shuffle[8] = {0u, 4u, 1u, 5u, 2u, 6u, 3u, 7u}; + for (uint d = 0; d < 8u; ++d) { + uint p = shuffle[d]; + float v = (float((abs_code >> (4u * p)) & 0xFu) - 8.0f) * 0.5f; + v = ((effective >> p) & 1u) != 0u ? -v : v; + part += (v + offset) * xv[i * 8u + d]; + } + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr2.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr2.metal new file mode 100644 index 0000000..690ab5c --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr2.metal @@ -0,0 +1,96 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +const device half4* abs4 = (const device half4*)abs_tab; +const device ushort4* sign4 = (const device ushort4*)sign_tab; +threadgroup half4 tg_abs[512]; +threadgroup ushort4 tg_sign[512]; +if (uint(TG_CODEBOOK) != 0u) { + for (uint idx = tg_lane; idx < 512u; idx += threads_per_threadgroup.x) { + tg_abs[idx] = abs4[idx]; + tg_sign[idx] = sign4[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float4 xe[CHUNK * 2]; + float xs[CHUNK]; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint xb = x_row + (cw0 + i) * 8u; + xe[2 * i] = float4(float(x[xb]), float(x[xb + 1u]), float(x[xb + 2u]), float(x[xb + 3u])); + xe[2 * i + 1] = float4(float(x[xb + 4u]), float(x[xb + 5u]), float(x[xb + 6u]), float(x[xb + 7u])); + xs[i] = dot(xe[2 * i] + xe[2 * i + 1], float4(1.0f)); + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint a = (codeword >> 8u) * 2u; + uint sg = (codeword & 0xFFu) * 2u; + half4 a0, a1; + ushort4 s0, s1; + if (uint(TG_CODEBOOK) != 0u) { + a0 = tg_abs[a]; a1 = tg_abs[a + 1u]; s0 = tg_sign[sg]; s1 = tg_sign[sg + 1u]; + } else { + a0 = abs4[a]; a1 = abs4[a + 1u]; s0 = sign4[sg]; s1 = sign4[sg + 1u]; + } + float4 w0 = float4(as_type(as_type(a0) ^ s0)); + float4 w1 = float4(as_type(as_type(a1) ^ s1)); + float offset = (popcount(codeword & 0xFFu) & 1u) != 0u ? -0.25f : 0.25f; + part += dot(w0, xe[2 * i]) + dot(w1, xe[2 * i + 1]) + offset * xs[i]; + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr3.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr3.metal new file mode 100644 index 0000000..3b50a31 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr3.metal @@ -0,0 +1,89 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +threadgroup uint local_codebook[256]; +if (uint(TG_CODEBOOK) != 0u) { + for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + local_codebook[idx] = codebook[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float xv[CHUNK * 8]; + for (uint i = 0; i < uint(CHUNK) * 8u; ++i) { + xv[i] = float(x[x_row + cw0 * 8u + i]); + } + uint group = cw0 / codewords_per_group; + // Issue every row's code and scale load before any decode. + uint4 words[ROWS_PER_SG]; + float row_scale[ROWS_PER_SG]; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = min(row0 + r, out_dim - 1u); + uint base = expert * out_dim + row; + if (uint(CHUNK) == 8u) { + words[r] = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words[r] = uint4(((const device uint2*)(codes + base * codewords))[cw0 / 4u], 0u, 0u); + } + row_scale[r] = float(scales[base * groups + group]); + } + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[r][i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint abs_code = uint(TG_CODEBOOK) != 0u ? local_codebook[codeword >> 8u] : codebook[codeword >> 8u]; + float offset = parity != 0u ? -0.25f : 0.25f; + const uint shuffle[8] = {0u, 4u, 1u, 5u, 2u, 6u, 3u, 7u}; + for (uint d = 0; d < 8u; ++d) { + uint p = shuffle[d]; + float v = (float((abs_code >> (4u * p)) & 0xFu) - 8.0f) * 0.5f; + v = ((effective >> p) & 1u) != 0u ? -v : v; + part += (v + offset) * xv[i * 8u + d]; + } + } + acc[r] += part * row_scale[r]; + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr4.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr4.metal new file mode 100644 index 0000000..f8f1291 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr4.metal @@ -0,0 +1,87 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +threadgroup uint local_codebook[256]; +if (uint(TG_CODEBOOK) != 0u) { + for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + local_codebook[idx] = codebook[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + X_T xv[CHUNK * 8]; + for (uint i = 0; i < uint(CHUNK) * 8u; ++i) { + xv[i] = x[x_row + cw0 * 8u + i]; + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint abs_code = uint(TG_CODEBOOK) != 0u ? local_codebook[codeword >> 8u] : codebook[codeword >> 8u]; + float offset = parity != 0u ? -0.25f : 0.25f; + const uint shuffle[8] = {0u, 4u, 1u, 5u, 2u, 6u, 3u, 7u}; + for (uint d = 0; d < 8u; ++d) { + uint p = shuffle[d]; + float v = (float((abs_code >> (4u * p)) & 0xFu) - 8.0f) * 0.5f; + v = ((effective >> p) & 1u) != 0u ? -v : v; + part += (v + offset) * float(xv[i * 8u + d]); + } + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr5.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr5.metal new file mode 100644 index 0000000..46b9020 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr5.metal @@ -0,0 +1,96 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +threadgroup uint local_codebook[256]; +if (uint(TG_CODEBOOK) != 0u) { + for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + local_codebook[idx] = codebook[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float2 xv[CHUNK * 4]; + float xs[CHUNK]; + for (uint i = 0; i < uint(CHUNK); ++i) { + float sum = 0.0f; + for (uint k = 0; k < 4u; ++k) { + uint xi = x_row + (cw0 + i) * 8u + 2u * k; + xv[i * 4u + k] = float2(float(x[xi]), float(x[xi + 1u])); + sum += xv[i * 4u + k].x + xv[i * 4u + k].y; + } + xs[i] = sum; + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint abs_code = uint(TG_CODEBOOK) != 0u ? local_codebook[codeword >> 8u] : codebook[codeword >> 8u]; + // Nibbles k and k+4 are dims 2k and 2k+1: one half2 per pair. + // (1024 + n) * 0.5 - 516 = (n - 8) / 2, exact in half. + uint spread = (effective & 0x0Fu) | ((effective & 0xF0u) << 12u); + for (uint k = 0; k < 4u; ++k) { + half2 h = as_type(((abs_code >> (4u * k)) & 0x000F000Fu) | 0x64006400u); + half2 v = fma(h, half2(0.5h), half2(-516.0h)); + v = as_type(as_type(v) ^ ((spread << (15u - k)) & 0x80008000u)); + part += dot(float2(v), xv[i * 4u + k]); + } + part += (parity != 0u ? -0.25f : 0.25f) * xs[i]; + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr6.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr6.metal new file mode 100644 index 0000000..b42aafc --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr6.metal @@ -0,0 +1,98 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +threadgroup uint4 decoded[256]; +for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + uint c = codebook[idx]; + uint4 entry; + for (uint k = 0; k < 4u; ++k) { + half2 h = as_type(((c >> (4u * k)) & 0x000F000Fu) | 0x64006400u); + entry[k] = as_type(fma(h, half2(0.5h), half2(-516.0h))); + } + decoded[idx] = entry; +} +threadgroup_barrier(mem_flags::mem_threadgroup); + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float2 xv[CHUNK * 4]; + float xs[CHUNK]; + for (uint i = 0; i < uint(CHUNK); ++i) { + float sum = 0.0f; + for (uint k = 0; k < 4u; ++k) { + uint xi = x_row + (cw0 + i) * 8u + 2u * k; + xv[i * 4u + k] = float2(float(x[xi]), float(x[xi + 1u])); + sum += xv[i * 4u + k].x + xv[i * 4u + k].y; + } + xs[i] = sum; + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint4 entry = decoded[codeword >> 8u]; + // Nibbles k and k+4 are dims 2k and 2k+1: one half2 per pair. + // (1024 + n) * 0.5 - 516 = (n - 8) / 2, exact in half. + uint spread = (effective & 0x0Fu) | ((effective & 0xF0u) << 12u); + for (uint k = 0; k < 4u; ++k) { + half2 v = as_type(entry[k] ^ ((spread << (15u - k)) & 0x80008000u)); + part += dot(float2(v), xv[i * 4u + k]); + } + part += (parity != 0u ? -0.25f : 0.25f) * xs[i]; + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr7.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr7.metal new file mode 100644 index 0000000..6bc0763 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_mr7.metal @@ -0,0 +1,100 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +threadgroup uint local_codebook[256]; +if (uint(TG_CODEBOOK) != 0u) { + for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + local_codebook[idx] = codebook[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float2 xv[CHUNK * 4]; + half2 xh[CHUNK * 4]; + float xs[CHUNK]; + for (uint i = 0; i < uint(CHUNK); ++i) { + float sum = 0.0f; + for (uint k = 0; k < 4u; ++k) { + uint xi = x_row + (cw0 + i) * 8u + 2u * k; + xv[i * 4u + k] = float2(float(x[xi]), float(x[xi + 1u])); + xh[i * 4u + k] = half2(xv[i * 4u + k]); + sum += xv[i * 4u + k].x + xv[i * 4u + k].y; + } + xs[i] = sum; + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint abs_code = uint(TG_CODEBOOK) != 0u ? local_codebook[codeword >> 8u] : codebook[codeword >> 8u]; + // Nibbles k and k+4 are dims 2k and 2k+1: one half2 per pair. + // (1024 + n) * 0.5 - 516 = (n - 8) / 2, exact in half. + uint spread = (effective & 0x0Fu) | ((effective & 0xF0u) << 12u); + half2 hacc = half2(0.0h); + for (uint k = 0; k < 4u; ++k) { + half2 h = as_type(((abs_code >> (4u * k)) & 0x000F000Fu) | 0x64006400u); + half2 v = fma(h, half2(0.5h), half2(-516.0h)); + v = as_type(as_type(v) ^ ((spread << (15u - k)) & 0x80008000u)); + hacc = fma(v, xh[i * 4u + k], hacc); + } + part += float(hacc.x) + float(hacc.y); + part += (parity != 0u ? -0.25f : 0.25f) * xs[i]; + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_stub.metal b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_stub.metal new file mode 100644 index 0000000..643f7e5 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/e8p_stub.metal @@ -0,0 +1,76 @@ +// Single-token (M=1) gather matvec for 16-bit E8P codes. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK consecutive +// codewords (one 8- or 16-byte load per row), so a simdgroup reads one +// contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// CHUNK * 8 activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d]; bit p of the effective signs +// negates nibble p, and every value gets +-0.25 by parity. +// +// PER_ROUTE=0: one activation row for every route (gate/up). +// PER_ROUTE=1: route r reads activation row r (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint top_k = uint(rhs_indices_shape[1]); +uint expert = uint(rhs_indices[route % top_k]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = uint(PER_ROUTE) != 0u ? route * codewords * 8u : 0u; + +threadgroup uint local_codebook[256]; +if (uint(TG_CODEBOOK) != 0u) { + for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + local_codebook[idx] = codebook[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float xv[CHUNK * 8]; + for (uint i = 0; i < uint(CHUNK) * 8u; ++i) { + xv[i] = float(x[x_row + cw0 * 8u + i]); + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + part += float(codeword) * xv[i * 8u]; + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h.py new file mode 100644 index 0000000..7693e61 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h.py @@ -0,0 +1,29 @@ +import sys, importlib.util, time +import numpy as np, mlx.core as mx +S = sys.argv[1]; sys.argv = [sys.argv[0], S] +def load(name): + src = open(f"{S}/{name}.py").read() + src = src[: src.index("rng = np.random.default_rng(0)")] # defs only + ns = {}; exec(src, ns); return ns +A, B = load("mr"), load("mr2") +from benchmarks.bench_e8p_m1 import synthetic_projection +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)); xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +cfgA = [(4, 1, 8, 1), (8, 1, 8, 0), (2, 1, 8, 0), (8, 2, 8, 0)] +cfgB = [(2, 2, 4, 1), (8, 2, 8, 1), (4, 1, 4, 0), (4, 1, 8, 1)] +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nb = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases = {"old": lambda r: A["old"](x, L, r, pr)} + for c in cfgA: cases[f"arith {c}"] = lambda r, c=c: A["new"](x, L, r, pr, *c) + for c in cfgB: cases[f"table {c}"] = lambda r, c=c: B["new"](x, L, r, pr, *c) + for f in cases.values(): once(f); once(f) + best = {k: 1e9 for k in cases} + for _ in range(25): + for k, f in cases.items(): best[k] = min(best[k], once(f)) + for k in sorted(best, key=best.get): + print(f"{name} {k:22s} {best[k]*1e3:.4f} ms {nb / best[k] / 1e9:6.1f} GB/s x{best['old'] / best[k]:.2f}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h3.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h3.py new file mode 100644 index 0000000..5dc6fe4 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h3.py @@ -0,0 +1,29 @@ +import sys, importlib.util, time +import numpy as np, mlx.core as mx +S = sys.argv[1]; sys.argv = [sys.argv[0], S] +def load(name): + src = open(f"{S}/{name}.py").read() + src = src[: src.index("rng = np.random.default_rng(0)")] # defs only + ns = {}; exec(src, ns); return ns +A, B = load("mr"), load("mr3") +from benchmarks.bench_e8p_m1 import synthetic_projection +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)); xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +cfgA = [(4, 1, 8, 1), (8, 1, 8, 0)] +cfgB = [(4, 1, 8, 1), (8, 1, 8, 0), (8, 2, 8, 0), (16, 1, 8, 0), (16, 2, 8, 1)] +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nb = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases = {"old": lambda r: A["old"](x, L, r, pr)} + for c in cfgA: cases[f"arith {c}"] = lambda r, c=c: A["new"](x, L, r, pr, *c) + for c in cfgB: cases[f"pre {c}"] = lambda r, c=c: B["new"](x, L, r, pr, *c) + for f in cases.values(): once(f); once(f) + best = {k: 1e9 for k in cases} + for _ in range(25): + for k, f in cases.items(): best[k] = min(best[k], once(f)) + for k in sorted(best, key=best.get): + print(f"{name} {k:22s} {best[k]*1e3:.4f} ms {nb / best[k] / 1e9:6.1f} GB/s x{best['old'] / best[k]:.2f}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h4.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h4.py new file mode 100644 index 0000000..ce9645f --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h4.py @@ -0,0 +1,29 @@ +import sys, importlib.util, time +import numpy as np, mlx.core as mx +S = sys.argv[1]; sys.argv = [sys.argv[0], S] +def load(name): + src = open(f"{S}/{name}.py").read() + src = src[: src.index("rng = np.random.default_rng(0)")] # defs only + ns = {}; exec(src, ns); return ns +A, B = load("mr"), load("mr4") +from benchmarks.bench_e8p_m1 import synthetic_projection +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)); xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +cfgA = [(4, 1, 8, 1), (8, 1, 8, 0)] +cfgB = [(4, 1, 8, 1), (8, 1, 8, 0), (8, 2, 8, 1), (16, 1, 8, 1), (4, 2, 8, 1), (2, 1, 8, 1)] +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nb = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases = {"old": lambda r: A["old"](x, L, r, pr)} + for c in cfgA: cases[f"arith {c}"] = lambda r, c=c: A["new"](x, L, r, pr, *c) + for c in cfgB: cases[f"xhalf {c}"] = lambda r, c=c: B["new"](x, L, r, pr, *c) + for f in cases.values(): once(f); once(f) + best = {k: 1e9 for k in cases} + for _ in range(25): + for k, f in cases.items(): best[k] = min(best[k], once(f)) + for k in sorted(best, key=best.get): + print(f"{name} {k:22s} {best[k]*1e3:.4f} ms {nb / best[k] / 1e9:6.1f} GB/s x{best['old'] / best[k]:.2f}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h5.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h5.py new file mode 100644 index 0000000..3322616 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h5.py @@ -0,0 +1,29 @@ +import sys, importlib.util, time +import numpy as np, mlx.core as mx +S = sys.argv[1]; sys.argv = [sys.argv[0], S] +def load(name): + src = open(f"{S}/{name}.py").read() + src = src[: src.index("rng = np.random.default_rng(0)")] # defs only + ns = {}; exec(src, ns); return ns +A, B = load("mr"), load("mr5") +from benchmarks.bench_e8p_m1 import synthetic_projection +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)); xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +cfgA = [(4, 1, 8, 1), (8, 1, 8, 0)] +cfgB = [(4, 1, 8, 1), (8, 1, 8, 0), (4, 4, 8, 1), (2, 1, 8, 1), (8, 2, 8, 1), (4, 1, 4, 1)] +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nb = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases = {"old": lambda r: A["old"](x, L, r, pr)} + for c in cfgA: cases[f"arith {c}"] = lambda r, c=c: A["new"](x, L, r, pr, *c) + for c in cfgB: cases[f"half2 {c}"] = lambda r, c=c: B["new"](x, L, r, pr, *c) + for f in cases.values(): once(f); once(f) + best = {k: 1e9 for k in cases} + for _ in range(25): + for k, f in cases.items(): best[k] = min(best[k], once(f)) + for k in sorted(best, key=best.get): + print(f"{name} {k:22s} {best[k]*1e3:.4f} ms {nb / best[k] / 1e9:6.1f} GB/s x{best['old'] / best[k]:.2f}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h67.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h67.py new file mode 100644 index 0000000..2e197d0 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2h67.py @@ -0,0 +1,33 @@ +import sys, importlib.util, time +import numpy as np, mlx.core as mx +S = sys.argv[1]; sys.argv = [sys.argv[0], S] +def load(name): + src = open(f"{S}/{name}.py").read() + src = src[: src.index("rng = np.random.default_rng(0)")] # defs only + ns = {}; exec(src, ns); return ns +A, B, C, D = load("mr"), load("mr5"), load("mr6"), load("mr7") +from benchmarks.bench_e8p_m1 import synthetic_projection +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)); xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +cfgA = [(4, 1, 8, 1), (8, 1, 8, 0), (2, 1, 8, 0), (8, 2, 8, 0)] +cfgB = [(2, 2, 4, 1), (8, 2, 8, 1), (4, 1, 4, 0), (4, 1, 8, 1)] +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nb = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases = {"old": lambda r: A["old"](x, L, r, pr)} + for c in [(4, 4, 8, 1), (8, 1, 8, 0), (4, 1, 8, 1), (8, 4, 8, 1), (2, 4, 8, 1)]: + cases[f"h2 {c}"] = lambda r, c=c: B["new"](x, L, r, pr, *c) + cases[f"tab {c}"] = lambda r, c=c: C["new"](x, L, r, pr, *c) + cases[f"hacc {c}"] = lambda r, c=c: D["new"](x, L, r, pr, *c) + for f in cases.values(): once(f); once(f) + best = {k: 1e9 for k in cases} + for _ in range(25): + for k, f in cases.items(): best[k] = min(best[k], once(f)) + ref = np.array(cases["old"](routes[0]).astype(mx.float32)).reshape(-1) + for k in sorted(best, key=best.get)[:9] + ["old"]: + d = np.max(np.abs(np.array(cases[k](routes[0]).astype(mx.float32)).reshape(-1) - ref)) / np.max(np.abs(ref)) + print(f"{name} {k:22s} {best[k]*1e3:.4f} ms {nb / best[k] / 1e9:6.1f} GB/s x{best['old'] / best[k]:.2f} diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2hs.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2hs.py new file mode 100644 index 0000000..1a30a15 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/h2hs.py @@ -0,0 +1,29 @@ +import sys, importlib.util, time +import numpy as np, mlx.core as mx +S = sys.argv[1]; sys.argv = [sys.argv[0], S] +def load(name): + src = open(f"{S}/{name}.py").read() + src = src[: src.index("rng = np.random.default_rng(0)")] # defs only + ns = {}; exec(src, ns); return ns +A, B = load("mr"), load("stub") +from benchmarks.bench_e8p_m1 import synthetic_projection +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)); xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +cfgA = [(4, 1, 8, 1), (8, 1, 8, 0)] +cfgB = [(4, 1, 8, 1), (8, 1, 8, 0), (1, 1, 8, 1), (4, 4, 8, 1)] +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nb = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases = {"old": lambda r: A["old"](x, L, r, pr)} + for c in cfgA: cases[f"arith {c}"] = lambda r, c=c: A["new"](x, L, r, pr, *c) + for c in cfgB: cases[f"stub {c}"] = lambda r, c=c: B["new"](x, L, r, pr, *c) + for f in cases.values(): once(f); once(f) + best = {k: 1e9 for k in cases} + for _ in range(25): + for k, f in cases.items(): best[k] = min(best[k], once(f)) + for k in sorted(best, key=best.get): + print(f"{name} {k:22s} {best[k]*1e3:.4f} ms {nb / best[k] / 1e9:6.1f} GB/s x{best['old'] / best[k]:.2f}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr.py new file mode 100644 index 0000000..23ac27a --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr.py @@ -0,0 +1,46 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_mr.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, L.codebook, r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr2.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr2.py new file mode 100644 index 0000000..8cf83fa --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr2.py @@ -0,0 +1,59 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +SHUF = np.array([0, 4, 1, 5, 2, 6, 3, 7]) +_T = {} +def tables(cb): + if id(cb) not in _T: + c = np.array(cb).astype(np.uint32) + nib = (c[:, None] >> (4 * SHUF[None, :]).astype(np.uint32)) & 0xF + abs_tab = ((nib.astype(np.float32) - 8) * 0.5).astype(np.float16) + s = np.arange(256, dtype=np.uint32) + par = np.array([bin(v).count("1") & 1 for v in s], dtype=np.uint32) + eff = s ^ par + sign_tab = (((eff[:, None] >> SHUF[None, :].astype(np.uint32)) & 1) << 15).astype(np.uint16) + _T[id(cb)] = (mx.array(abs_tab), mx.array(sign_tab)) + return _T[id(cb)] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "abs_tab", "sign_tab", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_mr2.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, *tables(L.codebook), r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr3.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr3.py new file mode 100644 index 0000000..c6aa7ec --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr3.py @@ -0,0 +1,46 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_mr3.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, L.codebook, r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr4.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr4.py new file mode 100644 index 0000000..ad8d317 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr4.py @@ -0,0 +1,46 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_mr4.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, L.codebook, r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("X_T", x.dtype), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr5.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr5.py new file mode 100644 index 0000000..eafd7b2 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr5.py @@ -0,0 +1,46 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_mr5.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, L.codebook, r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr6.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr6.py new file mode 100644 index 0000000..7a88f09 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr6.py @@ -0,0 +1,46 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_mr6.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, L.codebook, r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr7.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr7.py new file mode 100644 index 0000000..b93620b --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/mr7.py @@ -0,0 +1,46 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_mr7.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, L.codebook, r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/stub.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/stub.py new file mode 100644 index 0000000..329f59c --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-tuning/stub.py @@ -0,0 +1,46 @@ +import sys, time, itertools +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.kernels.gather_vqmm import _qmv_output_dtype +S = sys.argv[1] +K = mx.fast.metal_kernel(name="e8p_mr", input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], source=open(S + "/e8p_stub.metal").read()) +def new(x, L, r, per_route, R, SG, C, tg): + routes = r.size; out_dim = L.codes.shape[1] + (o,) = K(inputs=[x, L.codes, L.scales, L.codebook, r.reshape(1, routes)], + template=[("ROWS_PER_SG", R), ("SIMDGROUPS", SG), ("CHUNK", C), ("TG_CODEBOOK", tg), + ("PER_ROUTE", per_route), ("OUT_T", _qmv_output_dtype(x.dtype))], + grid=(32 * SG, (out_dim + R * SG - 1) // (R * SG), routes), threadgroup=(32 * SG, 1, 1), + output_shapes=[(1, routes, out_dim)], output_dtypes=[_qmv_output_dtype(x.dtype)]) + return o +def old(x, L, r, per_route): + f = e8p_m1.e8p_m1_per_route if per_route else e8p_m1.e8p_m1_token_layout + return f(x, L.codes, L.scales, L.codebook, r, output_dims=L.codes.shape[1], rows_per_threadgroup=16) +rng = np.random.default_rng(0) +gate = synthetic_projection(experts=64, out=2048, inp=4096, group_size=512, seed=1) +down = synthetic_projection(experts=64, out=4096, inp=2048, group_size=512, seed=3) +routes = [mx.array(np.sort(rng.choice(64, 6, replace=False))[None].astype(np.uint32)) for _ in range(64)] +xg = mx.array(rng.normal(size=(1, 4096)).astype(np.float16)) +xd = mx.array(rng.normal(size=(6, 2048)).astype(np.float16)) +def once(fn): + a = time.perf_counter(); mx.eval([fn(r) for r in routes * 4]); return (time.perf_counter() - a) / 256 +cases = {} +for name, L, x, pr in (("gate", gate, xg, 0), ("down", down, xd, 1)): + nbytes = 6 * (L.codes.nbytes + L.scales.nbytes) // 64 + cases[(name, "old")] = (lambda r, L=L, x=x, pr=pr: old(x, L, r, pr), nbytes, L, x, pr) + for R, SG, C, tg in itertools.product((1, 2, 4, 8), (1, 2, 4), (4, 8), (1, 0)): + cases[(name, f"R{R} SG{SG} C{C} tg{tg}")] = (lambda r, L=L, x=x, pr=pr, a=(R, SG, C, tg): new(x, L, r, pr, *a), nbytes, L, x, pr) +for k, (fn, *_ ) in cases.items(): + for _ in range(2): once(fn) +best = {k: 1e9 for k in cases} +for _ in range(8): + for k, (fn, *_ ) in cases.items(): best[k] = min(best[k], once(fn)) +for name in ("gate", "down"): + ref = np.array(cases[(name, "old")][0](routes[0]).astype(mx.float32)).reshape(-1) + rows = sorted(((k[1], best[k]) for k in cases if k[0] == name), key=lambda t: t[1]) + for cfg, t in rows[:8] + [r for r in rows if r[0] == "old"]: + got = np.array(cases[(name, cfg)][0](routes[0]).astype(mx.float32)).reshape(-1) + d = np.max(np.abs(got - ref)) / np.max(np.abs(ref)) + print(f"{name} {cfg:18s} {t*1e3:.4f} ms {cases[(name, cfg)][1] / t / 1e9:6.1f} GB/s diff {d:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-verify/verify_bench.py b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-verify/verify_bench.py new file mode 100644 index 0000000..8edddf2 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-verify/verify_bench.py @@ -0,0 +1,42 @@ +import os, sys, time, statistics, warnings +warnings.filterwarnings("ignore") +import numpy as np, mlx.core as mx +sys.path.insert(0, ".") +from benchmarks.bench_e8p_m1 import synthetic_projection +from ramp.kernels import e8p_m1 +from ramp.models.glm4_moe_adapter import QuantizedVQSwitchGLU +E = 256 +glu = QuantizedVQSwitchGLU( + gate_proj=synthetic_projection(experts=E, out=2048, inp=4096, group_size=512, seed=1), + up_proj=synthetic_projection(experts=E, out=2048, inp=4096, group_size=512, seed=2), + down_proj=synthetic_projection(experts=E, out=4096, inp=2048, group_size=512, seed=3), +) +per_expert = sum((p.codes.nbytes + p.scales.nbytes) // E for p in (glu.gate_proj, glu.up_proj, glu.down_proj)) +rng = np.random.default_rng(0) +def routes(m): + return [mx.array(np.stack([np.sort(rng.choice(E, 6, replace=False)) for _ in range(m)])[None].astype(np.uint32)) for _ in range(32)] +def chain(x, rs): + for r in rs: + y = glu(x, r) # [1, m, 6, hidden] + x = (x + y.sum(axis=2) * 1e-3).astype(x.dtype) + return x +def ms_per_ffn(m, rs, n=43, samples=12): + x = mx.array(rng.normal(size=(1, m, 4096)).astype(np.float16)) + def one(i): + seq = [rs[(i * n + j) % len(rs)] for j in range(n)] + t = time.perf_counter(); mx.eval(chain(x, seq)); return (time.perf_counter() - t) * 1e3 / n + one(0); one(1) + return statistics.median(one(i) for i in range(samples)) +for m in (1, 2, 4, 8, 11, 16, 21, 32): + rs = routes(m) + distinct = statistics.mean(len(set(np.array(r).reshape(-1).tolist())) for r in rs) + out = {} + for flag in ("0", "1"): + os.environ[e8p_m1.E8P_M1_ENV] = flag + out[flag] = ms_per_ffn(m, rs) + x = mx.array(rng.normal(size=(1, m, 4096)).astype(np.float16)) + os.environ[e8p_m1.E8P_M1_ENV] = "0"; a = np.array(glu(x, rs[0]).astype(mx.float32)) + os.environ[e8p_m1.E8P_M1_ENV] = "1"; b = np.array(glu(x, rs[0]).astype(mx.float32)) + diff = np.max(np.abs(a - b)) / np.max(np.abs(a)) + gbs = distinct * per_expert / (out["1"] / 1e3) / 1e9 + print(f"M={m:2d} distinct experts {distinct:5.1f} old {out['0']:.3f} ms new {out['1']:.3f} ms x{out['0']/out['1']:.2f} {gbs:.0f} GB/s distinct diff {diff:.1e}") diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-verify/verify_bench_m4max.txt b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-verify/verify_bench_m4max.txt new file mode 100644 index 0000000..4b6a3d2 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/e8p-verify/verify_bench_m4max.txt @@ -0,0 +1,8 @@ +M= 1 distinct experts 6.0 old 0.770 ms new 0.195 ms x3.96 197 GB/s distinct diff 3.0e-04 +M= 2 distinct experts 11.8 old 1.445 ms new 0.314 ms x4.60 240 GB/s distinct diff 5.7e-04 +M= 4 distinct experts 23.1 old 2.809 ms new 0.537 ms x5.23 275 GB/s distinct diff 2.8e-04 +M= 8 distinct experts 43.8 old 5.532 ms new 0.998 ms x5.54 280 GB/s distinct diff 4.7e-04 +M=11 distinct experts 59.0 old 4.444 ms new 1.360 ms x3.27 277 GB/s distinct diff 4.5e-04 +M=16 distinct experts 79.7 old 5.366 ms new 1.932 ms x2.78 263 GB/s distinct diff 5.0e-04 +M=21 distinct experts 99.4 old 7.073 ms new 2.539 ms x2.79 250 GB/s distinct diff 5.1e-04 +M=32 distinct experts 136.2 old 39.129 ms new 4.493 ms x8.71 194 GB/s distinct diff 5.2e-04 diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/hc-sinkhorn.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/hc-sinkhorn.json new file mode 100644 index 0000000..e686b81 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/hc-sinkhorn.json @@ -0,0 +1,28 @@ +{ + "device": "Device(gpu, 0)", + "mlx_version": "0.31.2", + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "reasons": [], + "record_type": "ramp_hc_sinkhorn_kill_test_v1", + "rows": [ + { + "compiled_ms_per_call": 1.407559666025918, + "compiled_ms_per_decode_step": 121.05013127822895, + "fused_ms_per_call": 0.3815831679967232, + "fused_ms_per_decode_step": 32.816152447718196, + "max_relative_diff": 5.69557300877932e-07, + "rows": 1, + "speedup": 3.688736254839221 + }, + { + "compiled_ms_per_call": 0.9860470839776099, + "compiled_ms_per_decode_step": 84.80004922207445, + "fused_ms_per_call": 1.9785255839815363, + "fused_ms_per_decode_step": 170.15320022241212, + "max_relative_diff": 1.4791691000937135e-06, + "rows": 1024, + "speedup": 0.4983746947529043 + } + ], + "verdict": "PASS" +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/metal-parity-tests-after-fix.txt b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/metal-parity-tests-after-fix.txt new file mode 100644 index 0000000..d6f6f3b --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/metal-parity-tests-after-fix.txt @@ -0,0 +1,11 @@ +...................................................................... [100%] +=============================== warnings summary =============================== +:491 + :491: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute + +:491 + :491: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute + +-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html +70 passed, 80 deselected, 2 warnings in 5.96s +:0: DeprecationWarning: builtin type swigvarlink has no __module__ attribute diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/metal-parity-tests.txt b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/metal-parity-tests.txt new file mode 100644 index 0000000..3fb8e66 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/metal-parity-tests.txt @@ -0,0 +1,64 @@ +........F............................................................. [100%] +=================================== FAILURES =================================== +__________________ test_hyper_connection_output_with_the_flag __________________ + +monkeypatch = <_pytest.monkeypatch.MonkeyPatch object at 0x119e616e0> + + @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]) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +tests/test_hc_sinkhorn.py:88: +_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ +:48: in __init__ + ??? +_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ + +self = DeepseekV4FlashVQModelArgs(model_type='deepseek_v4', vocab_size=129280, hidden_size=64, intermediate_size=18432, moe_i... dspark_target_layer_ids=[], dspark_markov_rank=256, n_mtp_layers=0, tie_word_embeddings=False, topk_method='noaux_tc') + + def __post_init__(self) -> None: + if not self.compress_ratios: + n = self.num_hidden_layers + self.compress_ratios = ( + [0] + + [4 if i % 2 else 128 for i in range(max(n - 2, 0))] + + ([0] if n >= 2 else []) + ) + # The drafter ratios live past ``num_hidden_layers``; keep them so the + # MTP blocks can be built from the same list, then truncate the + # backbone view exactly as upstream does. + self.mtp_compress_ratios = list(self.compress_ratios[self.num_hidden_layers :]) + self.compress_ratios = list(self.compress_ratios[: self.num_hidden_layers]) + if len(self.compress_ratios) != self.num_hidden_layers: + raise ValueError( + "`compress_ratios` must have one entry per hidden layer, " + f"got {len(self.compress_ratios)} for {self.num_hidden_layers} layers." + ) + bad = [r for r in self.compress_ratios if r not in (0, 4, 128)] + if bad: + raise ValueError(f"Unsupported DeepSeek-V4 compress ratios: {bad}") + if self.num_hash_layers > self.num_hidden_layers: +> raise ValueError( + "`num_hash_layers` cannot exceed `num_hidden_layers`: " + f"{self.num_hash_layers} > {self.num_hidden_layers}" + ) +E ValueError: `num_hash_layers` cannot exceed `num_hidden_layers`: 3 > 1 + +src/ramp/models/deepseek_v4_flash_adapter.py:300: ValueError +=============================== warnings summary =============================== +:491 + :491: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute + +:491 + :491: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute + +-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html +=========================== short test summary info ============================ +FAILED tests/test_hc_sinkhorn.py::test_hyper_connection_output_with_the_flag +1 failed, 69 passed, 80 deselected, 2 warnings in 11.74s diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic-v2-run1.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic-v2-run1.json new file mode 100644 index 0000000..507b596 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic-v2-run1.json @@ -0,0 +1,429 @@ +{ + "created_utc": "2026-11-06T03:49:50+00:00", + "device": "Device(gpu, 0)", + "device_info": { + "architecture": "applegpu_g16s", + "device_name": "Apple M4 Max", + "max_buffer_length": 86586540032, + "max_recommended_working_set_size": 115448725504, + "memory_size": 137438953472, + "resource_limit": 499000 + }, + "exactness": { + "attn.wkv": { + "bf16_bytes": 4194304, + "bit_mismatches": 0, + "elements": 2097152, + "mxfp8_bytes": 2162688 + }, + "attn.wo_a": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "attn.wo_b": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "attn.wq_a": { + "bf16_bytes": 8388608, + "bit_mismatches": 0, + "elements": 4194304, + "mxfp8_bytes": 4325376 + }, + "attn.wq_b": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "ffn.shared_experts.w1": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + }, + "ffn.shared_experts.w2": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + }, + "ffn.shared_experts.w3": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + } + }, + "layer": null, + "mlx_version": "0.31.2", + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "record_type": "dsv4_mxfp8_resident_kill_test_v2", + "settings": { + "calls_per_sync": 43, + "copies": 8, + "iterations": 20, + "rounds": 5, + "seed": 20261105, + "warmup": 3 + }, + "shapes": { + "attn.wkv": [ + 512, + 4096 + ], + "attn.wo_a": [ + 8192, + 4096 + ], + "attn.wo_b": [ + 4096, + 8192 + ], + "attn.wq_a": [ + 1024, + 4096 + ], + "attn.wq_b": [ + 32768, + 1024 + ], + "ffn.shared_experts.w1": [ + 2048, + 4096 + ], + "ffn.shared_experts.w2": [ + 4096, + 2048 + ], + "ffn.shared_experts.w3": [ + 2048, + 4096 + ] + }, + "source": "synthetic", + "timings": { + "attn.wkv": { + "affine8_g64": { + "effective_gb_per_s": 117.88515055513322, + "median_ms_per_call": 0.018901651221609046, + "round_ms": [ + 0.019171511413269608, + 0.018474813767297323, + 0.01889244188976929, + 0.019241279047957166, + 0.018901651221609046 + ], + "weight_bytes": 2228224 + }, + "bf16": { + "effective_gb_per_s": 165.36217826043818, + "median_ms_per_call": 0.02536434899517443, + "round_ms": [ + 0.02575241845523462, + 0.02498158150850687, + 0.024488860765623663, + 0.02536434899517443, + 0.025727721147759015 + ], + "weight_bytes": 4194304 + }, + "mxfp8": { + "effective_gb_per_s": 103.28533410732432, + "median_ms_per_call": 0.02093896503982589, + "round_ms": [ + 0.021193325749685077, + 0.020817825760215868, + 0.020693313820439194, + 0.021076557797726332, + 0.02093896503982589 + ], + "weight_bytes": 2162688 + } + }, + "attn.wo_a": { + "affine8_g64": { + "effective_gb_per_s": 377.4318395226111, + "median_ms_per_call": 0.09445833728572917, + "round_ms": [ + 0.09079940710487494, + 0.10062548836833948, + 0.09445833728572917, + 0.09613808110453795, + 0.0907858487808817 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 400.91202351724274, + "median_ms_per_call": 0.1673904998190051, + "round_ms": [ + 0.17180233695652597, + 0.1666923490120098, + 0.16632412768748783, + 0.17667877918304226, + 0.1673904998190051 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 359.651200009676, + "median_ms_per_call": 0.09621268606658073, + "round_ms": [ + 0.09504505832473811, + 0.09621268606658073, + 0.10114632578036034, + 0.1009491395454334, + 0.09025629071063947 + ], + "weight_bytes": 34603008 + } + }, + "attn.wo_b": { + "affine8_g64": { + "effective_gb_per_s": 374.3957633791575, + "median_ms_per_call": 0.09522432539893617, + "round_ms": [ + 0.08996367436931144, + 0.09794525604500154, + 0.09233043035004981, + 0.09522432539893617, + 0.09756298839970123 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 384.7842163658154, + "median_ms_per_call": 0.17440648848288368, + "round_ms": [ + 0.17165697652164327, + 0.17440648848288368, + 0.17524320905659954, + 0.17554553490955122, + 0.172934593053448 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 370.70419374851963, + "median_ms_per_call": 0.09334398850495385, + "round_ms": [ + 0.08617103476142293, + 0.09418702359462894, + 0.09601744200471182, + 0.09334398850495385, + 0.08957752339314495 + ], + "weight_bytes": 34603008 + } + }, + "attn.wq_a": { + "affine8_g64": { + "effective_gb_per_s": 194.17178634635061, + "median_ms_per_call": 0.022951058358452174, + "round_ms": [ + 0.0277771282843639, + 0.02407123298156833, + 0.022951058358452174, + 0.021333337171263126, + 0.021545058062170136 + ], + "weight_bytes": 4456448 + }, + "bf16": { + "effective_gb_per_s": 272.3368393730946, + "median_ms_per_call": 0.03080232560277245, + "round_ms": [ + 0.044809104766436784, + 0.03175581379003043, + 0.03080232560277245, + 0.030590604711865442, + 0.03065068652409367 + ], + "weight_bytes": 8388608 + }, + "mxfp8": { + "effective_gb_per_s": 178.17737895723826, + "median_ms_per_call": 0.02427567419227819, + "round_ms": [ + 0.037029069823992636, + 0.030775197787037075, + 0.02427567419227819, + 0.023396325644192306, + 0.02350194209125326 + ], + "weight_bytes": 4325376 + } + }, + "attn.wq_b": { + "affine8_g64": { + "effective_gb_per_s": 430.18001625193773, + "median_ms_per_call": 0.08287596506835504, + "round_ms": [ + 0.08287596506835504, + 0.0827180233182992, + 0.08377180229006118, + 0.08250193042290765, + 0.09378973295814683 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 451.4740676890608, + "median_ms_per_call": 0.14864389519315474, + "round_ms": [ + 0.14864389519315474, + 0.15096608136488168, + 0.14745058129043426, + 0.14806202332608315, + 0.18916956959887904 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 417.38857949296505, + "median_ms_per_call": 0.08290358121929214, + "round_ms": [ + 0.08290358121929214, + 0.08399224451602286, + 0.08090261605481595, + 0.08119670924633135, + 0.08568072083763519 + ], + "weight_bytes": 34603008 + } + }, + "ffn.shared_experts.w1": { + "affine8_g64": { + "effective_gb_per_s": 270.20676991221814, + "median_ms_per_call": 0.03298546517874266, + "round_ms": [ + 0.03286337224106023, + 0.03354796540089558, + 0.03319427907077033, + 0.03205620931450624, + 0.03298546517874266 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 353.94466951119, + "median_ms_per_call": 0.0474006743007881, + "round_ms": [ + 0.04676745346270848, + 0.04999709334549342, + 0.048844488256416005, + 0.0474006743007881, + 0.047367255996674476 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 265.70166976009136, + "median_ms_per_call": 0.032558139389229204, + "round_ms": [ + 0.03290165122327676, + 0.0334568836098147, + 0.03255086038586532, + 0.03223401172151573, + 0.032558139389229204 + ], + "weight_bytes": 8650752 + } + }, + "ffn.shared_experts.w2": { + "affine8_g64": { + "effective_gb_per_s": 283.11890788381345, + "median_ms_per_call": 0.031481104764849124, + "round_ms": [ + 0.031481104764849124, + 0.031687499993199184, + 0.031428779461854246, + 0.03148061642964739, + 0.03174175562466993 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 351.02757450305774, + "median_ms_per_call": 0.047794581447771295, + "round_ms": [ + 0.04740746530736688, + 0.04811966311197381, + 0.047794581447771295, + 0.04782510459758774, + 0.04771656998389855 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 279.2398129187639, + "median_ms_per_call": 0.03097965118074573, + "round_ms": [ + 0.030670534998048532, + 0.03064632534894139, + 0.03097965118074573, + 0.03115019759672239, + 0.03109349995082634 + ], + "weight_bytes": 8650752 + } + }, + "ffn.shared_experts.w3": { + "affine8_g64": { + "effective_gb_per_s": 275.7845823914847, + "median_ms_per_call": 0.032318325856765515, + "round_ms": [ + 0.03201065099050919, + 0.031834302421356014, + 0.03246802343921952, + 0.032318325856765515, + 0.03248691825510197 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 354.386644461155, + "median_ms_per_call": 0.04734155832963108, + "round_ms": [ + 0.047691360313640255, + 0.04733527867058509, + 0.04734155832963108, + 0.04790067438545182, + 0.0469578605356388 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 267.56901299741423, + "median_ms_per_call": 0.032330918678104184, + "round_ms": [ + 0.0321501974276332, + 0.0324093954648476, + 0.03353826739543746, + 0.032330918678104184, + 0.03195251168975563 + ], + "weight_bytes": 8650752 + } + } + }, + "verdict": { + "bit_mismatches": 0, + "exact": true, + "mxfp8_speedup_over_bf16": 1.6664370215227893, + "mxfp8_time_over_affine8": 1.005708639691768, + "reasons": [], + "speed_valid": true, + "summed_median_ms": { + "affine8_g64": 0.4111962331334389, + "bf16": 0.6891443721711809, + "mxfp8": 0.41354360427100995 + }, + "verdict": "PASS" + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic-v2-run2.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic-v2-run2.json new file mode 100644 index 0000000..29242fc --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic-v2-run2.json @@ -0,0 +1,429 @@ +{ + "created_utc": "2026-11-06T03:50:00+00:00", + "device": "Device(gpu, 0)", + "device_info": { + "architecture": "applegpu_g16s", + "device_name": "Apple M4 Max", + "max_buffer_length": 86586540032, + "max_recommended_working_set_size": 115448725504, + "memory_size": 137438953472, + "resource_limit": 499000 + }, + "exactness": { + "attn.wkv": { + "bf16_bytes": 4194304, + "bit_mismatches": 0, + "elements": 2097152, + "mxfp8_bytes": 2162688 + }, + "attn.wo_a": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "attn.wo_b": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "attn.wq_a": { + "bf16_bytes": 8388608, + "bit_mismatches": 0, + "elements": 4194304, + "mxfp8_bytes": 4325376 + }, + "attn.wq_b": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "ffn.shared_experts.w1": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + }, + "ffn.shared_experts.w2": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + }, + "ffn.shared_experts.w3": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + } + }, + "layer": null, + "mlx_version": "0.31.2", + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "record_type": "dsv4_mxfp8_resident_kill_test_v2", + "settings": { + "calls_per_sync": 43, + "copies": 8, + "iterations": 20, + "rounds": 5, + "seed": 20261105, + "warmup": 3 + }, + "shapes": { + "attn.wkv": [ + 512, + 4096 + ], + "attn.wo_a": [ + 8192, + 4096 + ], + "attn.wo_b": [ + 4096, + 8192 + ], + "attn.wq_a": [ + 1024, + 4096 + ], + "attn.wq_b": [ + 32768, + 1024 + ], + "ffn.shared_experts.w1": [ + 2048, + 4096 + ], + "ffn.shared_experts.w2": [ + 4096, + 2048 + ], + "ffn.shared_experts.w3": [ + 2048, + 4096 + ] + }, + "source": "synthetic", + "timings": { + "attn.wkv": { + "affine8_g64": { + "effective_gb_per_s": 123.19335756867, + "median_ms_per_call": 0.018087208953274543, + "round_ms": [ + 0.017982069335249797, + 0.01816520924939839, + 0.018087208953274543, + 0.018070732124236433, + 0.018106604625774157 + ], + "weight_bytes": 2228224 + }, + "bf16": { + "effective_gb_per_s": 165.3812097704874, + "median_ms_per_call": 0.02536143015171293, + "round_ms": [ + 0.025331395337098214, + 0.025791662940105726, + 0.025838662580543653, + 0.025216569896614135, + 0.02536143015171293 + ], + "weight_bytes": 4194304 + }, + "mxfp8": { + "effective_gb_per_s": 101.77589158529817, + "median_ms_per_call": 0.021249511709631703, + "round_ms": [ + 0.021006302166253675, + 0.020820244085476843, + 0.021312011508704272, + 0.021485941752596476, + 0.021249511709631703 + ], + "weight_bytes": 2162688 + } + }, + "attn.wo_a": { + "affine8_g64": { + "effective_gb_per_s": 389.6368286215456, + "median_ms_per_call": 0.09149952309725937, + "round_ms": [ + 0.09149952309725937, + 0.09543944230075752, + 0.08811966311673872, + 0.10774080197565084, + 0.08725677887627552 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 404.80547814406657, + "median_ms_per_call": 0.16578052329646725, + "round_ms": [ + 0.16559400023123638, + 0.1665716977291849, + 0.16527616271348455, + 0.18161095355280005, + 0.16578052329646725 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 365.816813102497, + "median_ms_per_call": 0.0945910815485255, + "round_ms": [ + 0.0945910815485255, + 0.09285950042350694, + 0.09141472062680783, + 0.1081787790902782, + 0.09522288374407867 + ], + "weight_bytes": 34603008 + } + }, + "attn.wo_b": { + "affine8_g64": { + "effective_gb_per_s": 387.107456589037, + "median_ms_per_call": 0.09209738379658394, + "round_ms": [ + 0.09140310447587263, + 0.09209738379658394, + 0.08843895345709699, + 0.12042149996146734, + 0.12089680206182218 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 384.48089172054546, + "median_ms_per_call": 0.1745440812407841, + "round_ms": [ + 0.1764244188115997, + 0.16961869753918848, + 0.16992925605357645, + 0.1961807209354024, + 0.1745440812407841 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 348.2641136251952, + "median_ms_per_call": 0.09935852316164867, + "round_ms": [ + 0.09935852316164867, + 0.09093652303373917, + 0.09315649993706841, + 0.12642346514757116, + 0.13854118642061514 + ], + "weight_bytes": 34603008 + } + }, + "attn.wq_a": { + "affine8_g64": { + "effective_gb_per_s": 194.24146461236384, + "median_ms_per_call": 0.02294282535859925, + "round_ms": [ + 0.023552802388642935, + 0.023326546841755856, + 0.02294282535859925, + 0.021597860194201214, + 0.021567348888966928 + ], + "weight_bytes": 4456448 + }, + "bf16": { + "effective_gb_per_s": 269.53868627697466, + "median_ms_per_call": 0.031122092772166918, + "round_ms": [ + 0.03372044214534811, + 0.031814918931855195, + 0.031122092772166918, + 0.03086579090449959, + 0.03101744216617717 + ], + "weight_bytes": 8388608 + }, + "mxfp8": { + "effective_gb_per_s": 177.73748468501907, + "median_ms_per_call": 0.024335755666089788, + "round_ms": [ + 0.028541662898131235, + 0.026405999791635156, + 0.024335755666089788, + 0.024085278963930038, + 0.02344090695936926 + ], + "weight_bytes": 4325376 + } + }, + "attn.wq_b": { + "affine8_g64": { + "effective_gb_per_s": 422.57369042139044, + "median_ms_per_call": 0.08436773232249326, + "round_ms": [ + 0.0854161856316983, + 0.08461288336209606, + 0.08425774420936441, + 0.08436773232249326, + 0.08346753494088474 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 410.1321441161742, + "median_ms_per_call": 0.16362741853511176, + "round_ms": [ + 0.16440843031633384, + 0.16242732547039468, + 0.16278100029435444, + 0.16418846525782504, + 0.16362741853511176 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 387.49845290049126, + "median_ms_per_call": 0.08929844168664584, + "round_ms": [ + 0.0915271395866131, + 0.08929844168664584, + 0.08705813948032555, + 0.09061433753908374, + 0.08242344179158206 + ], + "weight_bytes": 34603008 + } + }, + "ffn.shared_experts.w1": { + "affine8_g64": { + "effective_gb_per_s": 284.94760510813404, + "median_ms_per_call": 0.03127906969640144, + "round_ms": [ + 0.03127906969640144, + 0.030744174457436733, + 0.031008244002086308, + 0.03279989543316756, + 0.03204747681170269 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 354.70603402628257, + "median_ms_per_call": 0.04729893035526106, + "round_ms": [ + 0.04660512752516932, + 0.04780716310136107, + 0.04729893035526106, + 0.04804845348401212, + 0.046610465032308424 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 273.0814011550195, + "median_ms_per_call": 0.03167829066135942, + "round_ms": [ + 0.031283430194720456, + 0.03167829066135942, + 0.03128827936665783, + 0.031721418477096704, + 0.03222287206120009 + ], + "weight_bytes": 8650752 + } + }, + "ffn.shared_experts.w2": { + "affine8_g64": { + "effective_gb_per_s": 281.3307754453691, + "median_ms_per_call": 0.031681197998655396, + "round_ms": [ + 0.031019860829854775, + 0.03179651177434121, + 0.03254069773436979, + 0.030814441595074918, + 0.031681197998655396 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 351.82300887874, + "median_ms_per_call": 0.04768652298628504, + "round_ms": [ + 0.04734787216875702, + 0.047846418753981074, + 0.047673453335910165, + 0.04843604645345273, + 0.04768652298628504 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 279.4058272822925, + "median_ms_per_call": 0.03096124402323175, + "round_ms": [ + 0.030500488761855765, + 0.030575569876101473, + 0.03096124402323175, + 0.031158907425611518, + 0.03127859320578187 + ], + "weight_bytes": 8650752 + } + }, + "ffn.shared_experts.w3": { + "affine8_g64": { + "effective_gb_per_s": 279.9692533063395, + "median_ms_per_call": 0.031835267247177314, + "round_ms": [ + 0.031857558073974106, + 0.03157461638735651, + 0.03156395355607708, + 0.031835267247177314, + 0.032773244108051755 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 357.8732852182121, + "median_ms_per_call": 0.046880325223969, + "round_ms": [ + 0.04709204645329264, + 0.0463788838228644, + 0.04638033698388744, + 0.046880325223969, + 0.04720494190371747 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 269.8172499246007, + "median_ms_per_call": 0.03206152313248103, + "round_ms": [ + 0.03181781408615237, + 0.03169282565575613, + 0.032378871976614515, + 0.03206152313248103, + 0.032399220968769904 + ], + "weight_bytes": 8650752 + } + } + }, + "verdict": { + "bit_mismatches": 0, + "exact": true, + "mxfp8_speedup_over_bf16": 1.6581920421850846, + "mxfp8_time_over_affine8": 1.0488970824576456, + "reasons": [], + "speed_valid": true, + "summed_median_ms": { + "affine8_g64": 0.4037902084704445, + "bf16": 0.7023013245617581, + "mxfp8": 0.4235343715896137 + }, + "verdict": "PASS" + } +} diff --git a/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic.json b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic.json new file mode 100644 index 0000000..9928525 --- /dev/null +++ b/artifacts/benchmarks/phase0-MAXBOOK-20261105-192445/mxfp8-residents-synthetic.json @@ -0,0 +1,431 @@ +{ + "created_utc": "2026-11-06T03:25:00+00:00", + "device": "Device(gpu, 0)", + "device_info": { + "architecture": "applegpu_g16s", + "device_name": "Apple M4 Max", + "max_buffer_length": 86586540032, + "max_recommended_working_set_size": 115448725504, + "memory_size": 137438953472, + "resource_limit": 499000 + }, + "exactness": { + "attn.wkv": { + "bf16_bytes": 4194304, + "bit_mismatches": 0, + "elements": 2097152, + "mxfp8_bytes": 2162688 + }, + "attn.wo_a": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "attn.wo_b": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "attn.wq_a": { + "bf16_bytes": 8388608, + "bit_mismatches": 0, + "elements": 4194304, + "mxfp8_bytes": 4325376 + }, + "attn.wq_b": { + "bf16_bytes": 67108864, + "bit_mismatches": 0, + "elements": 33554432, + "mxfp8_bytes": 34603008 + }, + "ffn.shared_experts.w1": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + }, + "ffn.shared_experts.w2": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + }, + "ffn.shared_experts.w3": { + "bf16_bytes": 16777216, + "bit_mismatches": 0, + "elements": 8388608, + "mxfp8_bytes": 8650752 + } + }, + "layer": null, + "mlx_version": "0.31.2", + "platform": "macOS-26.7.1-arm64-arm-64bit-Mach-O", + "record_type": "dsv4_mxfp8_resident_kill_test_v1", + "settings": { + "copies": 8, + "iterations": 200, + "rounds": 5, + "seed": 20261105, + "warmup": 20 + }, + "shapes": { + "attn.wkv": [ + 512, + 4096 + ], + "attn.wo_a": [ + 8192, + 4096 + ], + "attn.wo_b": [ + 4096, + 8192 + ], + "attn.wq_a": [ + 1024, + 4096 + ], + "attn.wq_b": [ + 32768, + 1024 + ], + "ffn.shared_experts.w1": [ + 2048, + 4096 + ], + "ffn.shared_experts.w2": [ + 4096, + 2048 + ], + "ffn.shared_experts.w3": [ + 2048, + 4096 + ] + }, + "source": "synthetic", + "timings": { + "attn.wkv": { + "affine8_g64": { + "effective_gb_per_s": 12.346655118439584, + "median_ms_per_call": 0.18047187506454065, + "round_ms": [ + 0.1531339599750936, + 0.14204958497430198, + 0.18047187506454065, + 0.445007084927056, + 0.3591016650898382 + ], + "weight_bytes": 2228224 + }, + "bf16": { + "effective_gb_per_s": 12.865132878226667, + "median_ms_per_call": 0.3260210399457719, + "round_ms": [ + 0.17775187487131916, + 0.16968333497061394, + 0.3260210399457719, + 0.33897583503858186, + 0.4410322899639141 + ], + "weight_bytes": 4194304 + }, + "mxfp8": { + "effective_gb_per_s": 7.474577455789129, + "median_ms_per_call": 0.2893391650286503, + "round_ms": [ + 0.18327979007153772, + 0.16327021003235132, + 0.2893391650286503, + 0.3063174999260809, + 0.3682972899696324 + ], + "weight_bytes": 2162688 + } + }, + "attn.wo_a": { + "affine8_g64": { + "effective_gb_per_s": 102.68116105737253, + "median_ms_per_call": 0.34720666997600347, + "round_ms": [ + 0.24603687488706782, + 0.3152260449132882, + 0.34720666997600347, + 0.4145656249602325, + 0.49060249992180616 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 181.9718413458899, + "median_ms_per_call": 0.36878707993309945, + "round_ms": [ + 0.259654375113314, + 0.3418068749306258, + 0.39245832987944596, + 0.36878707993309945, + 0.37461583502590656 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 75.11911070382045, + "median_ms_per_call": 0.4606418749608565, + "round_ms": [ + 0.22329687504679896, + 0.4196122899884358, + 0.4606418749608565, + 0.49901457998203114, + 0.6001195800490677 + ], + "weight_bytes": 34603008 + } + }, + "attn.wo_b": { + "affine8_g64": { + "effective_gb_per_s": 94.23450790150733, + "median_ms_per_call": 0.3783283299708273, + "round_ms": [ + 0.356820000015432, + 0.5951014599122573, + 0.3783283299708273, + 0.5414641699462663, + 0.3499670849123504 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 124.70975174653974, + "median_ms_per_call": 0.5381204200966749, + "round_ms": [ + 0.3941564599517733, + 0.46919395492295735, + 0.5381204200966749, + 0.5693583299580496, + 0.6413285399321467 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 72.3405002191153, + "median_ms_per_call": 0.4783352049707901, + "round_ms": [ + 0.4622772899165284, + 0.4783352049707901, + 0.6258887500735, + 0.7353650000004563, + 0.471238544996595 + ], + "weight_bytes": 34603008 + } + }, + "attn.wq_a": { + "affine8_g64": { + "effective_gb_per_s": 25.627142711392384, + "median_ms_per_call": 0.17389562504831702, + "round_ms": [ + 0.17389562504831702, + 0.16125479000038467, + 0.21068166504846886, + 0.16948375006904826, + 0.1889118750113994 + ], + "weight_bytes": 4456448 + }, + "bf16": { + "effective_gb_per_s": 40.897470608003545, + "median_ms_per_call": 0.20511312497546896, + "round_ms": [ + 0.23084375003236346, + 0.20371875012642704, + 0.20511312497546896, + 0.20693083002697676, + 0.19877270999131724 + ], + "weight_bytes": 8388608 + }, + "mxfp8": { + "effective_gb_per_s": 20.422224049928882, + "median_ms_per_call": 0.21179750008741394, + "round_ms": [ + 0.2476037500309758, + 0.1843235400156118, + 0.2059089599060826, + 0.2374814600625541, + 0.21179750008741394 + ], + "weight_bytes": 4325376 + } + }, + "attn.wq_b": { + "affine8_g64": { + "effective_gb_per_s": 117.4911368600051, + "median_ms_per_call": 0.3034406249935273, + "round_ms": [ + 0.3034406249935273, + 0.29054125014226884, + 0.31788437510840595, + 0.2908995850884821, + 0.3122981249180157 + ], + "weight_bytes": 35651584 + }, + "bf16": { + "effective_gb_per_s": 213.11194810661834, + "median_ms_per_call": 0.3148995849187486, + "round_ms": [ + 0.3262793750036508, + 0.3105731248797383, + 0.32184896001126617, + 0.29918750005890615, + 0.3148995849187486 + ], + "weight_bytes": 67108864 + }, + "mxfp8": { + "effective_gb_per_s": 86.6203516645367, + "median_ms_per_call": 0.3994789600255899, + "round_ms": [ + 0.4192366650386248, + 0.3994789600255899, + 0.39770604504155926, + 0.39270124994800426, + 0.4122983349952847 + ], + "weight_bytes": 34603008 + } + }, + "ffn.shared_experts.w1": { + "affine8_g64": { + "effective_gb_per_s": 18.53428043671034, + "median_ms_per_call": 0.4808870800479781, + "round_ms": [ + 0.5062785399786662, + 0.4808870800479781, + 0.5387225000595208, + 0.4101327050011605, + 0.3056120849214494 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 34.19429269589055, + "median_ms_per_call": 0.49064375009038486, + "round_ms": [ + 0.29815021000104025, + 0.5594412499340251, + 0.5425206250220072, + 0.4029964600340463, + 0.49064375009038486 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 20.047966791213195, + "median_ms_per_call": 0.43150270998012275, + "round_ms": [ + 0.23859770997660235, + 0.4383347899420187, + 0.43150270998012275, + 0.4810493749391753, + 0.4037372949824203 + ], + "weight_bytes": 8650752 + } + }, + "ffn.shared_experts.w2": { + "affine8_g64": { + "effective_gb_per_s": 34.47281436689568, + "median_ms_per_call": 0.25854854509816505, + "round_ms": [ + 0.2853339599096216, + 0.25854854509816505, + 0.34017124999081716, + 0.22430207987781614, + 0.18124374997569248 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 43.488465694280784, + "median_ms_per_call": 0.3857854199304711, + "round_ms": [ + 0.6273304150090553, + 0.3663843750837259, + 0.3857854199304711, + 0.27615958504611626, + 0.45533145996159874 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 25.411608398029266, + "median_ms_per_call": 0.34042520506773144, + "round_ms": [ + 0.3799189599521924, + 0.36401020508492365, + 0.34042520506773144, + 0.22804603999247774, + 0.32679124997230247 + ], + "weight_bytes": 8650752 + } + }, + "ffn.shared_experts.w3": { + "affine8_g64": { + "effective_gb_per_s": 17.94197381660171, + "median_ms_per_call": 0.49676228998578154, + "round_ms": [ + 0.49676228998578154, + 0.45178312502685003, + 0.3772306250175461, + 0.6179714600148145, + 0.6743810400075745 + ], + "weight_bytes": 8912896 + }, + "bf16": { + "effective_gb_per_s": 42.092094426890746, + "median_ms_per_call": 0.39858353993622586, + "round_ms": [ + 0.353656039951602, + 0.5094829149311408, + 0.5573883350007236, + 0.39488395501393825, + 0.39858353993622586 + ], + "weight_bytes": 16777216 + }, + "mxfp8": { + "effective_gb_per_s": 10.516395663417452, + "median_ms_per_call": 0.8225966649479233, + "round_ms": [ + 0.48426271008793265, + 0.4548260450246744, + 0.827316249924479, + 1.3013379200128838, + 0.8225966649479233 + ], + "weight_bytes": 8650752 + } + } + }, + "verdict": { + "bit_mismatches": 0, + "exact": true, + "mxfp8_speedup_over_bf16": 0.8817270082742493, + "mxfp8_time_over_affine8": 1.3109614365218596, + "reasons": [ + "mxfp8 is 0.882x bf16, below the 1.5x bar", + "mxfp8 takes 1.311x affine-8's time, above 1.1x" + ], + "speed_valid": true, + "summed_median_ms": { + "affine8_g64": 2.6195410401851404, + "bf16": 3.0279539598268457, + "mxfp8": 3.434117285069078 + }, + "verdict": "KILL" + } +} 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..1c547c1 --- /dev/null +++ b/benchmarks/bench_dsv4_mxfp8_residents.py @@ -0,0 +1,447 @@ +"""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. Each timed sample chains ``--calls-per-sync`` +dependent calls (default 43, one per V4-Flash layer) before one ``mx.eval``, +as one decode step does; v1 synced after every sub-millisecond call and so +mostly measured CPU-GPU round trips. + +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_v2" +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], mx.array], + x: mx.array, + *, + iterations: int, + warmup: int, + calls_per_sync: int, +) -> float: + """Median ms per call over ``iterations`` samples of chained calls. + + Each call's activation depends on the previous call's output, like + consecutive layers in one decode step, so the GPU cannot overlap them. + """ + + def sample(index: int) -> float: + started = time.perf_counter() + y = x + for j in range(calls_per_sync): + out = fn(index * calls_per_sync + j, y) + y = y + out.reshape(-1)[:1].astype(y.dtype) * 1e-6 + mx.eval(y) + return (time.perf_counter() - started) * 1_000.0 / calls_per_sync + + for i in range(warmup): + sample(i) + return statistics.median(sample(i) for i in range(iterations)) + + +def time_variants( + codes: np.ndarray, + scales: np.ndarray, + *, + copies: int, + iterations: int, + warmup: int, + rounds: int, + seed: int, + groups: int = 1, + calls_per_sync: int = 43, +) -> 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, x: mx.array) -> mx.array: + return x @ bf16_copies[i % copies].swapaxes(-1, -2) + + def mxfp8(i: int, x: mx.array) -> 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, x: mx.array) -> 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], + x, + iterations=iterations, + warmup=warmup, + calls_per_sync=calls_per_sync, + ) + ) + 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, + calls_per_sync: int = 43, +) -> 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), + calls_per_sync=calls_per_sync, + ) + 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, + "calls_per_sync": calls_per_sync, + }, + "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=20, help="timed samples per round") + parser.add_argument("--calls-per-sync", type=int, default=43, help="chained calls per mx.eval") + parser.add_argument("--warmup", type=int, default=3) + 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, + calls_per_sync=args.calls_per_sync, + ) + + 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..47faae7 --- /dev/null +++ b/benchmarks/bench_e8p_m1.py @@ -0,0 +1,245 @@ +"""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 kernel, swept over rows per + simdgroup and simdgroups per threadgroup. + +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 per FFN. +Each timed sample chains ``--ffns-per-sync`` dependent FFNs (default 43, one +per V4-Flash layer) before one ``mx.eval``, as one decode step does; syncing +after every FFN measured mostly CPU-GPU round trips (v1 rows). Verdict +for the best ``e8p_m1`` configuration: **PASS** at >= 200 GB/s, **TUNE** at +100-200 (tune; decoding each codeword once was tried and is slower), +**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 _chain(layers, x, routes_seq, intermediate): + """Dependent FFNs, each feeding the next like consecutive decoder layers.""" + + for routes in routes_seq: + y = _ffn(layers, x, routes, intermediate) # [1, k, hidden] + x = (x + y.sum(axis=1) * 1e-3).astype(x.dtype) + return x + + +def _median_ms(layers, x, routes_list, intermediate, iterations, warmup, ffns_per_sync) -> float: + """Median ms per FFN, with ``ffns_per_sync`` chained FFNs per ``mx.eval``.""" + + def sample(i): + seq = [routes_list[(i * ffns_per_sync + j) % len(routes_list)] for j in range(ffns_per_sync)] + started = time.perf_counter() + mx.eval(_chain(layers, x, seq, intermediate)) + return (time.perf_counter() - started) * 1e3 / ffns_per_sync + + for i in range(warmup): + sample(i) + return statistics.median(sample(i) for i in range(iterations)) + + +def run( + layers, + *, + top_k: int, + iterations: int, + warmup: int, + seed: int, + sweep: Sequence[tuple[int, bool]], + ffns_per_sync: int = 43, +) -> 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(layers, x, routes_list, intermediate, iterations, warmup, ffns_per_sync) + 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_sg, simdgroups in sweep: + switch_linear.e8p_m1_token_layout = ( + lambda *a, _r=rows_per_sg, _s=simdgroups, **k: original_token( + *a, rows_per_simdgroup=_r, simdgroups=_s, **k + ) + ) + switch_linear.e8p_m1_per_route = ( + lambda *a, _r=rows_per_sg, _s=simdgroups, **k: original_route( + *a, rows_per_simdgroup=_r, simdgroups=_s, **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(layers, x, routes_list, intermediate, iterations, warmup, ffns_per_sync) + configs.append( + { + "rows_per_simdgroup": rows_per_sg, + "simdgroups": simdgroups, + "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, + "ffns_per_sync": ffns_per_sync, + "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=50, help="timed samples, each --ffns-per-sync FFNs") + parser.add_argument("--ffns-per-sync", type=int, default=43, help="chained FFNs per mx.eval (V4-Flash has 43 layers)") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--seed", type=int, default=20261105) + parser.add_argument("--sweep", default="4:1,8:1,4:4,2:4,16:1", help="rows_per_simdgroup:simdgroups,...") + 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, simdgroups = item.split(":") + sweep.append((int(rows), int(simdgroups))) + result = { + "record_type": "ramp_e8p_m1_kill_test_v2", + "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, + ffns_per_sync=args.ffns_per_sync, + ), + } + 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..4f3f2d8 --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.md @@ -0,0 +1,97 @@ +# Phase 0 kill tests: toward DeepSeek-V4.1-Flash on one 128 GB Mac + +Date: 2026-11-05. Target machine: the M4 Max (40-core GPU, 128 GB; 546 GB/s +theoretical, 474-477 GB/s measured MLX read on 2026-11-05). 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 per routed FFN, 43 chained FFNs per sync | >= 200 GB/s: make it the M4 default. 100-200: tune (decoding each codeword once was tried and is slower) | +| 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. + On a synthetic V4-Flash FFN (256 experts, top-6, 43 chained FFNs per sync) + the same kernel now also serves MTP verify windows (<= 32 tokens): 2.6-5.5x + the previous M4 path at 2-21 tokens and 9x at 32, outputs within 5.7e-4. + It still is not traced as a verify implementation, so the MTP headline + harness, which expects `nax_e8p_m32n64`, needs updating before M4 rows can + validate. +- Decoding each E8P codeword once was tried on the M4 Max on 2026-11-05 and + made the kernel slower (113 vs 154 GB/s, kernel alone, 256 calls per sync). + The compiler already lifts the per-codeword work out of the 8-dimension loop. +- `bench_e8p_m1.py` v1 synced after every FFN and called KILL at 52.8 GB/s, + mostly CPU-GPU round trips. v2 (`ramp_e8p_m1_kill_test_v2`) chains 43 + dependent FFNs per `mx.eval`, as one decode step does. On the M4 Max it + gives **TUNE**: 122-126 GB/s best (8-16 rows per threadgroup), 2.4-2.5x + production's 50.5 GB/s, outputs within 6.4e-4. Quote only v2 rows. +- The tuned kernel (`src/ramp/kernels/e8p_m1.metal`, 2026-11-05) then + **PASSES** v2, narrowly: 201.9-205.5 GB/s best (4 rows per simdgroup), 3.97x + production, outputs unchanged at 6.4e-4. Each simdgroup owns several rows, + lanes load 8 codewords at once, and the abs grid is decoded once per + threadgroup into half2 pairs, with signs applied by XOR. A loads-only stub + of the same kernel reads at ~410 GB/s, so decode ALU is still half the + time. Since 2026-11-05 it is the default on any Metal GPU without NAX + (`RAMP_E8P_M1` unset or `auto`); `RAMP_E8P_M1=0` turns it off, `=1` forces + it on. It clears the bar by only 1-3%, measured while other jobs were + loading the machine, so rerun the kill test on an idle machine. +- `bench_dsv4_mxfp8_residents.py` v1 also synced after every call. v2 + (`dsv4_mxfp8_resident_kill_test_v2`) chains 43 dependent calls per + `mx.eval`. On the M4 Max the verdict flips from KILL to **PASS** in two runs: + mxfp8 is 1.66x bf16 and takes 1.01-1.05x affine-8's time, still bit-exact. + The large tensors reach 350-450 GB/s. 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..0df2a31 --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-05-v41-plan-and-phase0-status.md @@ -0,0 +1,125 @@ +# 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, 40-core GPU, 128 GB. Memory bandwidth is 546 GB/s +theoretical; MLX reads 474-477 GB/s (1 GiB reductions) and a bf16 matvec +458 GB/s, measured 2026-11-05. The 409.4 GB/s used for the estimates below is +an older, lower measurement, so those estimates are conservative. + +## 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. + On a synthetic V4-Flash FFN (256 experts, top-6, 43 chained FFNs per sync) + the same kernel now also serves MTP verify windows (<= 32 tokens): 2.6-5.5x + the previous M4 path at 2-21 tokens and 9x at 32, outputs within 5.7e-4. + It still is not traced as a verify implementation, so the MTP headline + harness, which expects `nax_e8p_m32n64`, needs updating before M4 rows can + validate. +- Decoding each E8P codeword once was tried on the M4 Max on 2026-11-05 and + made the kernel slower (113 vs 154 GB/s, kernel alone, 256 calls per sync). + The compiler already lifts the per-codeword work out of the 8-dimension loop. +- `bench_e8p_m1.py` v1 synced after every FFN and called KILL at 52.8 GB/s, + mostly CPU-GPU round trips. v2 (`ramp_e8p_m1_kill_test_v2`) chains 43 + dependent FFNs per `mx.eval`, as one decode step does. On the M4 Max it + gives **TUNE**: 122-126 GB/s best (8-16 rows per threadgroup), 2.4-2.5x + production's 50.5 GB/s, outputs within 6.4e-4. Quote only v2 rows. +- The tuned kernel (`src/ramp/kernels/e8p_m1.metal`, 2026-11-05) then + **PASSES** v2, narrowly: 201.9-205.5 GB/s best (4 rows per simdgroup), 3.97x + production, outputs unchanged at 6.4e-4. Each simdgroup owns several rows, + lanes load 8 codewords at once, and the abs grid is decoded once per + threadgroup into half2 pairs, with signs applied by XOR. A loads-only stub + of the same kernel reads at ~410 GB/s, so decode ALU is still half the + time. Since 2026-11-05 it is the default on any Metal GPU without NAX + (`RAMP_E8P_M1` unset or `auto`); `RAMP_E8P_M1=0` turns it off, `=1` forces + it on. It clears the bar by only 1-3%, measured while other jobs were + loading the machine, so rerun the kill test on an idle machine. +- `bench_dsv4_mxfp8_residents.py` v1 also synced after every call. v2 + (`dsv4_mxfp8_resident_kill_test_v2`) chains 43 dependent calls per + `mx.eval`. On the M4 Max the verdict flips from KILL to **PASS** in two runs: + mxfp8 is 1.66x bf16 and takes 1.01-1.05x affine-8's time, still bit-exact. + The large tensors reach 350-450 GB/s. diff --git a/docs/deepseek-v4-flash/2026-11-05-v41-runpod-reap-runbook.md b/docs/deepseek-v4-flash/2026-11-05-v41-runpod-reap-runbook.md new file mode 100644 index 0000000..59d7c06 --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-05-v41-runpod-reap-runbook.md @@ -0,0 +1,321 @@ +# DeepSeek-V4.1-Flash REAP on RunPod: runbook + +Date: 2026-11-05. Budget: $497. Kit: `scripts/runpod/v41_reap/`. + +> **2026-11-06:** the GPU stage now runs on one GPU (or several, data-parallel), +> streaming the weights one layer at a time with no `convert.py` step. See +> `2026-11-06-v41-reap-single-gpu-design.md`. The 8×H200 / TP8 steps below are kept for the record. + +## Goal + +The V4-Flash artifacts are gone. This run makes new V4.1 artifacts on rented +GPUs so that the Mac downloads a smaller, pruned checkpoint instead of the +510 GB source. + +| Output | Where (on the volume) | Size (est.) | Used for | +|---|---|---|---| +| Pruned HF checkpoint (FP8 dense, FP4 experts, Engram, drafter untouched) | `/workspace/v41-reapN/` | ~366 GB at N=192, ~342 GB at 160 | V4.1 port, E8P conversion | +| REAP stats, survivor tables, reap-select reports | `/workspace/out/` | < 50 MB | prune choice, hot set | +| Route traces (first 2,048 tokens of each session) | in each session npz | small | residency replay (SSD tier) | +| E8P importance stats (imatrix `hidden` + `down`) | `out/calibration/sessions/*.npz` | ~0.9 GB per session, ~25 sessions | E8P/VQ fitting (`finalize`) | +| Teacher top-2048 logits for the eval split | `out/eval/sessions/*.npz` | ~5 GB | KL gates, MTP and distillation | +| Candidate scores (KL, NLL change, top-1 agreement per domain) | `out/candidate-scores.json` | tiny | choosing N | + +## What is verified, and what is not + +**Verified on CPU on this Mac**, in 7 tests in `tests/test_runpod_v41_reap.py` +plus a manual end-to-end: + +- The patched `MoE.forward` gives bit-identical output to a copy of the + reference Gate/Expert/MoE. +- REAP and imatrix sums match a direct float64 computation. +- KEEP's own validators accept the session files, and `finalize-reap`, + `reap-select` and `finalize` run unchanged on them. +- The head capture matches full logits. +- A masked-router student scored against its own teacher gives KL 0. +- The pruner copies bytes exactly, renumbers experts, slices `weight`, `bias` + and `bias_vl`, and leaves `mtp.*` alone. +- The pack builder keeps eval and calibration groups disjoint and drops + non-English text. +- Real rows from Open-SWE-Traces, Nemotron-Terminal and Toucan render through + V4.1's own `encoding.py` and tokenizer. + +**Not verified until the pod runs:** + +- the reference code under torchrun with 8 GPUs and NCCL; +- throughput (`progress.jsonl` reports tok/s per session); +- GPU memory headroom at 16K tokens; +- that each dataset streams with the field names assumed. They were checked + against one row of each via the HF datasets-server. + +## Calibration mix + +The research is in `docs/research/2026-11-05-v41-reap-calibration-data.md`. In +summary: + +- **Size:** about 3.2M calibration tokens and 0.45M eval tokens. +- **Language:** English only. A record is dropped if it contains any CJK, + Hangul or Kana, or more than 2% of its letters are outside Latin. +- **Agentic and tool-loop share:** 60%. The sources are: + - Open-SWE-Traces (DeepSeek-V4-Flash-0731 OpenHands trajectories; Python 400K, TS 300K, Rust 300K); + - SWE-rebench-V2 PRs as issue-to-patch pairs (400K); + - Nemotron-Terminal (400K); + - Toucan multi-turn MCP traces (450K). +- **The rest:** + - cyber: CIRCL advisories plus security StackExchange, 300K; + - networking: RFC 9000-9800 plus networkengineering, serverfault and unix StackExchange, 250K; + - business: finepdfs business-filtered plus Enron, 250K; + - English: WildChat human prompts plus finewiki, 150K. +- **Held-out split:** one group bucket in ten goes to eval. Groups are + repository, task, RFC, URL host, CVE and tool set. + +Differences from the research mix: + +- **No Primus-FineWeb.** It is gated. +- **No WildChat responses**, only the human prompts. The responses are OpenAI + output. +- **No stackv2 plain-code files.** The shards have no language index, so I + doubled SWE-rebench PRs instead. +- **No 32K-64K long slice.** `--max-seq-len` is 16384. Longer records are + truncated, not windowed. +- **Assistant turns are not regenerated by V4.1.** That would add a generation + stage, so it is a follow-up. + +Jack's five candidates: + +- Open-SWE-Traces is the top source. +- CoderForge is suited to eval only and is not used yet. +- Vyber07/cyber-security, Fable-5.1 and Axion are rejected. Fable-5.1 is + Claude-distilled with placeholder tool arguments. + +## Cost and time (estimates) + +The prices come from runpod.io/pricing on 2026-11-05. Secure H200 is $5.29/GPU/h. + +| Stage | Pod | Time | Cost | +|---|---|---|---| +| Network volume, 1.5 TB, ~5 days | — | — | ~$15 | +| `fetch`: 510 GB download plus sha256 of every file | CPU | 1-3 h | < $2 | +| `pack` | CPU | 1-2 h | < $2 | +| `gpu`: install, convert (~1 h), teacher pass (~3.7M tokens), 4 student passes (4 × 0.45M) | 8×H200 | 2.5-3.5 h | $105-150 | +| `prune N`, then the Mac pulls ~366 GB | CPU | 1-4 h | < $2 | +| **Expected total** | | | **~$125-170** | +| **Worst case: the watchdog fires at `MAX_GPU_HOURS=6`** | | | **~$270** | + +The throughput behind the GPU estimate is a guess, about 3K tok/s with +reference prefill plus a Python expert loop. Check it on the first sessions. If +`tok_per_s` is under 1,000 after the first 3 sessions, stop and rethink. At +that rate the teacher pass alone takes more than an hour. + +## Single-pod mode (what actually ran, 2026-11-06) + +On 2026-11-06, 8×H200 was out of capacity in every data center that offers +network volumes, so everything runs on one 8×H200 pod with a 2 TB volume disk: + +```bash +MAX_GPU_HOURS=8 bash /workspace/keep/scripts/runpod/v41_reap/pod.sh all +``` + +That command runs these steps: + +- **fetch and pack:** run in parallel. +- **convert:** writes the TP8 checkpoint. +- **teacher pass:** runs over the pack. +- **scoring:** checks N = 256/192/128/96. +- **cleanup:** deletes the TP8 checkpoint. +- **prune:** writes REAP 50% (N=192, ~366 GB) and REAP 75% (N=96, ~293 GB). + +The pod stays up so the Mac can pull the results. The watchdog stops it at +8 h, and a failed stage stops it at once. + +The Mac cannot hold both checkpoints at once (423 GB free), so pull one at a +time. + +## Getting an 8×H200 pod (lessons from 2026-11-06) + +No 8×H200 pod could be created on 2026-11-06, even though the catalog listed +stock. These are the lessons from that day. + +**Tools:** + +- **Use the official Runpod plugin.** Install it with + `claude plugin marketplace add runpod/runpod-plugins-official` and + `claude plugin install runpod@runpod`, then sign in through `/mcp` → + runpod. It signs in with OAuth, so no API key ends up on disk. Its MCP tools + are REST v2 and read live state. +- **Check stock for the GPU count you need:** + `get-gpu-type id="NVIDIA H200" include=AVAILABILITY product=POD count=8`. + This returns stock per data center, the host CUDA versions that have + capacity, the minimum RAM and vCPU, and `maxCount`. `runpodctl gpu list` is + not enough: its `stockStatus` is the best status across all data centers and + is counted for a single GPU. +- **Get the price and CUDA matrix for many types at once:** + `get-capacity gpuCount=8`. +- **Create the pod with REST v2 `create-pod`.** The body must include `disk`, + `mounts` or `templateId`. Without one of them you get *"You must either + provide a template id or pod configuration parameters"*. + +**What the catalog showed, and what happened:** + +- **Catalog:** 8×H200 secure was "Low", only in EUR-IS-5 and only on CUDA 13.2 + hosts, at $5.29/GPU/h ($42.32/h). Those hosts had 1510 GB RAM and 160 vCPU. + B200, H200 NVL and every H100 variant showed no 8-GPU stock. +- **Every H200 create failed with the same 400** (*"There are no longer any + instances available"*): + - with the host specs pinned (`allowedCudaVersions`, `minRamPerGpu`, + `minVcpuCountPerGpu`); + - with 1400 GB and 1100 GB volume disks; + - with community cloud; + - with no data-center filter; + - with a bare single H200. +- **The control worked.** A 1×RTX 4090 with the same image and ports was + created at once. So the account, the image (`runpod/pytorch:1.5.0-cu1281-torch291-ubuntu2404`) + and the request shape were all fine; H200 itself was unplaceable. + +**What to take from it:** + +- **Catalog stock is not a reservation.** A "Low" reading can be a host that is + already held. Runpod's create-pod docs say a create can fail for capacity on + a GPU the catalog just reported as available. +- **The 400 means one of two things.** It covers both a broken request and no + capacity, so separate them with a cheap control pod. If the control succeeds, + the request is fine and the problem is capacity. Disk size, ports and CUDA + pinning were not the cause. +- **Do not run tight retry loops against the API.** A loop of creates every + 90 s was stopped as spam. Check by hand from time to time, then create once. +- **The image does not limit placement.** A CUDA 12.8 image runs on CUDA 13.2 + hosts. + +**Fix the local Claude CLI first if plugin install fails.** Old npm installs of +`@anthropic-ai/claude-code` crash on Node 26. Remove them with +`npm uninstall -g --prefix @anthropic-ai/claude-code`, so that +`claude` resolves to a working install. + +## Failure: the H100 sat idle for hours on prep work (2026-11-06) + +The 1×H100 SXM pod `m13sexafpkqia6` ($3.99/h, EUR-NO-2) ran `pod.sh all`. +Its GPU stayed at 0% for more than 2 hours: + +- **19:14–20:10:** the 510 GB weight fetch and sha256 check, at about 160–210 MB/s. +- **19:14 to 21:12 and later:** `build_pack.py`. Most of that time it was + streaming all 35 GB (gzip) of `common-pile/stackexchange_filtered` at about + 7 MB/s, unauthenticated, to keep posts from five sites only. + +That is at least $9 of GPU time spent on work that needs no GPU. It also pushes +the sweeps toward the 8 h watchdog. **This should have been caught before the +pod was rented.** The single-GPU design copied "fetch and pack" from the +8×H200 single-pod mode as GPU-pod time, and nobody questioned that or timed +the pack. + +**Rules for the next run:** + +- **Do no prep on a GPU pod.** Build and validate the pack on the Mac or on a + cheap CPU pod first. It needs only the tokenizer and public data. Download + the weights to a network volume from a CPU pod. Then create the GPU pod on + that volume. Its first command is `pod.sh smoke`. +- **Time slow inputs on cheap hardware.** Before any GPU spend, measure every + dataset stream and download in the pack. +- **Don't stream a whole corpus to keep a slice of it.** For StackExchange, + read only the per-site files, or pre-filter once on a CPU and keep the + result. Set `HF_TOKEN` (Jack enters it himself) so streams aren't throttled. +- **Count idle GPU time in the cost estimate.** If prep has to share the pod, + say up front how many idle GPU-hours it costs. + +## Steps (network-volume mode) + +Only you can do the steps marked **(you)**. I do not enter credentials or +launch pods. + +1. **(you)** In RunPod: + - Set an account spend limit at or below $497. + - Add your SSH public key. + - Create a **1.5 TB network volume** in a data center that offers 8×H200. +2. **(you)** Start a cheap **CPU pod** with the volume at `/workspace` and an + exposed SSH port. +3. Copy the kit onto the volume. It lives only in local commits, so do not push: + + ```bash + git archive HEAD | ssh -p PORT root@HOST 'mkdir -p /workspace/keep && tar x -C /workspace/keep' + ``` + +4. On the CPU pod: + + ```bash + bash /workspace/keep/scripts/runpod/v41_reap/pod.sh fetch + ``` + + ```bash + bash /workspace/keep/scripts/runpod/v41_reap/pod.sh pack + ``` + + - **Check the fetch:** it must end with `fetch verified`, and the printed + model.py sha256 must be `e5decd76…8b8d2d`. + - **Check the pack:** read `/workspace/pack/manifest.json` and confirm every + domain reached its budget. A source that falls short shows up in its + `filtered`, `off_domain` or `non_english` counts. + + Then stop the CPU pod. +5. **(you)** Start an **8×H200 pod** with the same volume. Pick a PyTorch 2.10+ + / CUDA 12.8+ image and **at least 1 TB of system RAM**, because `convert.py` + holds the whole checkpoint in memory. +6. On the GPU pod, run inside `tmux`, since the pod stops itself when the stage + ends or at the limit: + + ```bash + MAX_GPU_HOURS=6 bash /workspace/keep/scripts/runpod/v41_reap/pod.sh gpu + ``` + + - **Watch the run:** follow `/workspace/out/progress.jsonl` for the first + sessions. + - **Results:** `/workspace/out/candidate-scores.json` holds the per-domain + KL, NLL change and top-1 agreement for N = 256/224/192/160. + - **Watchdog:** it calls `runpodctl stop pod`. If that fails it logs the + failure, so also keep the spend limit from step 1 in place. +7. **Choose N** from the scores. Take the smallest N whose worst domain stays + close to N=256, and which also fits the Mac's disk: 423 GB free, minus + about 30 GB of stats and logits. +8. **(you)** Start a CPU pod on the volume again, then run: + + ```bash + bash /workspace/keep/scripts/runpod/v41_reap/pod.sh prune 192 + ``` + +9. Pull the results to the Mac. Run each transfer as its own command. + + The pruned checkpoint: + + ```bash + rsync -avP --partial -e "ssh -p PORT" root@HOST:/workspace/v41-reap192/ ~/keep-artifacts/v41-reap192/ + ``` + + The stats, logits and pack: + + ```bash + rsync -avP --partial -e "ssh -p PORT" root@HOST:/workspace/out/ ~/keep-artifacts/v41-runpod-out/ + ``` + + ```bash + rsync -avP --partial -e "ssh -p PORT" root@HOST:/workspace/pack/ ~/keep-artifacts/v41-pack/ + ``` + + Then verify the checkpoint: + + ```bash + cd ~/keep-artifacts/v41-reap192 && shasum -a 256 -c SHA256SUMS --quiet + ``` + +10. **(you)** Delete the network volume once the Mac copies verify. + +## Notes + +- **Survivors:** `reap-select` picks them from calibration stats only. The + per-layer Jaccard against frequency is in `reap-select-N.json`; a value of + 0.9 or more means REAP added little. +- **Masked student pass:** it sets the router bias of pruned experts to -inf, + which equals slicing them (`tests/test_reap.py`). One converted checkpoint + therefore scores every N. +- **Eval sessions carry REAP stats and routes but no imatrix**, so the + per-split saliency can be compared later. +- **Sizes:** routed experts are about 289 GB of the 510 GB source, roughly + 0.75 GB per expert slot across the 40 layers. Engram (~196 GB) is untouched + and is most of what is left after pruning. diff --git a/docs/deepseek-v4-flash/2026-11-06-v41-reap-single-gpu-design.md b/docs/deepseek-v4-flash/2026-11-06-v41-reap-single-gpu-design.md new file mode 100644 index 0000000..963c3cc --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-06-v41-reap-single-gpu-design.md @@ -0,0 +1,175 @@ +# V4.1-Flash REAP on one GPU: design + +Date: 2026-11-06. This replaces the TP8 plan in +`2026-11-05-v41-runpod-reap-runbook.md`, which needed an 8×H200 pod that +Runpod could not place. Every figure below comes from the reference config and +code at revision `1bae0e10`. I fetched `inference/{config,model,kernel,convert,engram}.py`; +`model.py` has sha256 `e5decd76…8b8d2d`, which matches the pin in `collect.py`. + +## Why 8 GPUs was never about compute + +| Quantity | Value | +|---|---| +| Active parameters per token | about 15.6B: 374M per layer × 40 layers, plus the head | +| Forward cost | about 31 GFLOP per token | +| Total work | 3.2M calibration tokens, plus 0.45M eval tokens × (1 teacher + 4 students) ≈ 5.5M token passes ≈ 1.7e17 FLOP | +| Work on one H100 | 15 minutes at 200 TFLOPS effective; 1 hour at 50 TFLOPS | +| Weights | 510 GB on disk. Routed experts are 6.8 GB per layer in FP4, Engram tables are about 196 GB. | + +The 8×H200 pod existed only to hold all 510 GB in GPU memory at once. The work +itself fits on one GPU in under an hour. So the fix is to **stream weights +through one GPU, one layer at a time, over large token batches**. Holding every +weight resident is what forced the 8-GPU pod. + +## Design + +1. **Load layer weights straight from the HF shards.** There is no `convert.py` + step and no TP8 copy. With one GPU, `convert.py` only renames tensors, + dequantizes `wo_a` to bf16 and views the experts as `float4_e2m1fn_x2`. That + is about 20 lines, applied per tensor as it loads. This saves the hour of + conversion, the 510 GB TP8 copy, and the need for more than 550 GB of RAM. +2. **Run the reference code unchanged, one layer at a time.** Build the + reference `Block` modules with `world_size=1`. For each chunk of records + (whole records, so attention sees full context): + - keep the chunk's hyper-connection residual in GPU memory, as + `[tokens, 4, 5120]` bf16 (40 KB per token); + - for each layer, load its weights, prefetching the next layer's; + - run attention one micro-batch of records at a time; + - run the MoE over the whole chunk at once, so each expert sees about 15K + tokens instead of about 100, which makes the expert GEMMs efficient. + + The REAP, imatrix and route hooks in `make_moe_forward` stay as they are. +3. **Save and restore the cross-layer attention state.** Layers pass + compressed KV, indexer keys, top-k indices and candidate blocks to later + layers through the module-level `shared_attn` object. KV comes from source + layers 2, 8, 14 and 20; the indexer from 2, 8, 14, 20, 24, 28, 32 and 36; + candidates from layer 20 into every layer after it. Done one layer at a + time, that state has to be saved per micro-batch and restored before the + next layer. The candidate mask is the large part, up to about 16 KB per + token at 16K context. It gets bit-packed, or kept in host RAM. +4. **Gather the Engram rows on the host, at run time.** The rows Engram + fetches depend only on the token ids (`NgramHashState` → `embed(hash_ids)`); + the gate needs the hidden state, but the lookup does not. A stand-in for + `ParallelEngramEmbedding` memory-maps the two 196 GB tables straight from the + shard files, gathers the FP8 rows and scales a chunk needs (24 × 256 values + per token per Engram layer) and dequantizes them exactly as the module does. + The tables never reach the GPU. I first planned to precompute the rows in + the pack stage (about 48 GB on disk); the run-time gather does the same + thing with no extra stage and no extra disk. +5. **Do the whole job in two sweeps over the weights.** + - **Sweep A, calibration** (4 chunks on an 80 GB GPU, 2 on an H200): writes + the REAP stats, imatrix and routes, then `finalize-reap` and + `reap-select` produce the survivors for 256/192/128/96. + - **Sweep B, eval:** the teacher and the 4 masked-router students run side + by side through each layer, so each layer is read once for all 5 streams. + This writes the teacher's top-2048 logits and the student + KL / NLL change / top-1 agreement, in the same npz schema as now. That + replaces the separate student `torchrun` pass. + + Every chunk is one pass over the backbone weights, about 320 GB without the + Engram tables. A chunk holds `--chunk-tokens` resident tokens (default 512K, + counting each stream), so one 80 GB GPU makes about 7 calibration passes and + 5 eval passes. That is about 4 TB of reads. From NVMe this is 20-30 minutes, + overlapped with compute by a one-layer prefetch, and less when the pod's + page cache holds the shards. +6. **Use as many GPUs as are free, data-parallel by chunk.** Chunks are + independent, so any pod with 1 to 8 GPUs of the Hopper or Blackwell + class (80 GB or more) works. The stats are sums, so the shards add together + exactly as the per-session npz files do today. +7. **Skip the MTP drafter.** It isn't needed for backbone REAP, and the pruner + already leaves `mtp.*` alone. + +## REAP-specific streamlining + +- **Calibration budget:** the per-session npz files already let me compute a + split-half survivor Jaccard offline, at no extra cost. If two halves of the + calibration set pick nearly the same survivors (Jaccard ≥ 0.98 at N=96), then + 3.2M tokens is more than needed and the next run can use fewer. +- **Candidates:** all four are cheap now. Each extra student costs memory + during Sweep B, not another pass over the weights. +- **E8P imatrix from pruned routing:** the teacher imatrix counts traffic + under the full 384-expert routing. After pruning, the surviving experts take + over the pruned experts' tokens. Sweep B can also write a pruned-routing + imatrix for 192 and 96 at no extra cost. The catch is that it comes from the + eval split, so it would leak a little into the KL measurement if it were + used to fit E8P. I'll write it, but not use it by default. + +## What changes in the kit + +| File | Change | +|---|---| +| `stream.py` (new) | `HFCheckpoint` (per-layer loads with `convert.py`'s fixups), `EngramRows` (host gather), `LayerStreamer` (the reference `Block`s one layer at a time, with `shared_attn` saved and restored per record). | +| `torch_kernel.py` (new) | Plain-torch stand-ins for `kernel.py`, for the CPU tests. | +| `collect.py` | The `torchrun` driver is replaced by `--sweep calibration|eval`, `--shard I/N` and `--chunk-tokens`. `SessionStats`, the MoE hook, the head capture and the npz schemas are unchanged. | +| `build_pack.py` | No change. | +| `pod.sh` | `convert` and the TP8 directory are gone. `gpu` runs one `collect.py` per GPU for sweep A, then select, then sweep B. | +| `prune.py`, `score.py`, `finalize-reap`, `reap-select` | No change. | + +## Validation before any spend + +1. **Run the real `model.py` on the CPU, at tiny size.** Replace + `kernel.py`'s tilelang functions with plain-torch versions (`act_quant`, + `fp4_act_quant`, `fp8_gemm`, `fp4_gemm`, `sparse_attn`, + `hc_split_sinkhorn`; the repo already has a torch Sinkhorn). Use a tiny + random `ModelArgs` that still has hyper-connections, KV and indexer sources, + a candidate source and Engram. The test then checks that the + one-layer-at-a-time driver matches `Transformer.forward` logits and REAP + stats with: + - one record per micro-batch; + - several records per micro-batch; + - records split across chunks. + + This is the test that catches mistakes in saving and restoring + `shared_attn`. + + **Status (2026-11-06): done.** `tests/test_runpod_v41_stream.py` runs the + pinned `model.py` and `engram.py` (copied to `tests/fixtures/v41_reference/`) + on a 6-layer toy with every feature above. The streamed hidden state equals + `Transformer.forward` bit for bit, with one MoE batch per chunk and with one + per record, and a record's output does not depend on its chunk-mates. Breaking + any one `shared_attn` field, the clone of block buffers, or the Engram lookup + fails the test. Both sweeps are checked too: + - calibration REAP, imatrix and route stats match the reference per session + (route counts exactly, sums to 1e-6); + - eval teacher logits match; + - a keep-all student scores KL 0; + - a keep-4-of-8 student's NLL matches the reference with that router mask. +2. **First 10 minutes on the pod:** run the same equivalence test with the real + tilelang kernels at tiny size, then one real 16K record through all 40 + layers. Check the teacher NLL is sensible, and measure GPU memory and tok/s. + Stop if anything is off. + +## Cost + +These prices are from the Runpod catalog read of 2026-11-06. 1×H100 SXM costs +$2.69/h, and 1×H200 costs $3.59/h on community cloud or $5.29/h on secure. + +| Stage | Time on 1×H100 | +|---|---| +| fetch and pack, in parallel | 1-3 h | +| sweeps A and B | about 1-2 h | +| prune 192 and 96 | about 1 h | + +The total is about 3-6 h, roughly **$10-25**, against $125-340 for the TP8 +plan. The disk needed is about 1.2 TB: + +| Item | Size | +|---|---| +| HF source | 510 GB | +| out | about 30 GB | +| REAP-192 | 366 GB | +| REAP-96 | 293 GB | + +## Risks + +- **GPU kernels:** the tilelang kernels are written for Hopper and Blackwell. + H100 and H200 are the safe target. The RTX PRO 6000 (SM120) is not verified. +- **`shared_attn` state:** saving and restoring it is the subtle part, and the + CPU equivalence test covers it. +- **Engram host gather:** this relies on `embed(hash_ids)` being a pure row + gather. I checked the code: `ParallelEngramEmbedding.forward` gathers FP8 + rows and their scales, then dequantizes, and the CPU test runs the real + module's math through the stand-in. +- **Indexer memory at 16K context:** the indexer's `bshd,btd->bsht` einsum is + about 16 GB for one 16K record. It should fit on 80 GB, but the first pod + minutes must check `max_gpu_gb` in the progress log before the full run. diff --git a/docs/deepseek-v4-flash/2026-11-08-v41-reap-recovery-log.md b/docs/deepseek-v4-flash/2026-11-08-v41-reap-recovery-log.md new file mode 100644 index 0000000..4d9b239 --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-08-v41-reap-recovery-log.md @@ -0,0 +1,1116 @@ +# DeepSeek-V4.1-Flash REAP recovery: progress log + +Started 2026-11-08. Branch `jack/hopeful-cori-iszsfq`. This log records what was run, what +it showed, and what was learned, in order. The plan is the five-step recovery list from the +same day (non-uniform budgets, KEEP importance on the pruned model, residual sidecars, router +KD, LoRA KD), applied to the published REAP checkpoints +([REAP-50](https://huggingface.co/jackmaxwil/DeepSeek-V4.1-Flash-REAP-50), +[REAP-75](https://huggingface.co/jackmaxwil/DeepSeek-V4.1-Flash-REAP-75)). + +Compute rules for this campaign: the Mac first. A Runpod pod only for a real experiment that +the Mac cannot do in reasonable time, never for setup or idle waiting. Mac memory cap 64 GB for +these runs (Jack, 2026-11-08). + +## Starting point + +REAP degradation, measured on the pod by router masking inside the full model (437,350 +held-out positions; KL over the teacher's top 2,048 logits): + +| Experts kept / 384 | KL | top-1 | NLL change | +|---:|---:|---:|---:| +| 256 | 0.511 | 0.842 | +0.025 | +| 192 (REAP-50) | 0.810 | 0.803 | -0.138 | +| 128 | 1.114 | 0.755 | -0.109 | +| 96 (REAP-75) | 1.349 | 0.720 | -0.131 | + +- Fit: KL ≈ 1.84 · p^1.17 for pruned fraction p (worst miss 0.03); 1 − top-1 ≈ 0.33 · p^0.69. + The marginal cost is ~0.005 KL per expert per layer, nearly flat from 33% to 75% pruned. +- Agentic and tool-use text loses the most and loses it early (exponent 0.85–1.1); business, + code-diffs, cyber, networking and chat lose late (exponent 2–2.5). +- KEEP's own quant loss on V4-Flash (unpruned, 2.031 bpw) was KL 0.149 / top-1 0.899. If the + two add, REAP-50 + KEEP ≈ 0.96 KL. KEEP's gate (KL ≤ 0.30, top-1 ≥ 0.85) fails on REAP alone + at every tested level, so REAP loss is the main target. +- Both exported checkpoints were run on the Mac (CPU, reference code, dequantize-then-matmul + stand-ins) over 5 eval sessions and matched the pod's masked-router numbers: REAP-75 KL 1.378 + vs 1.385, REAP-50 0.809 vs 0.841. Per-position differences of 0.1–0.3 KL are run-to-run + chaos from rounding; the same session moved 0.17 KL between a solo and a batched run. + +## Tooling + +- `scripts/runpod/v41_reap/torch_kernel.py`: exact CPU ports of the FP8/FP4 kernels (unit tests + in `tests/test_runpod_v41_torch_kernel.py`). +- `scripts/runpod/v41_reap/local_parity.py`: one exported checkpoint against the pod's numbers. +- `scripts/runpod/v41_reap/local_eval.py`: any number of router-masked variants of a local + checkpoint in one layer-streamed pass, scored per position against the pod teacher's top-2,048 + logits. Variants may keep a different count per layer. +- Mac eval set (`~/Models/v41-reap-run/mac-eval-sessions.json`): 2 eval sessions per domain + (22 sessions, 11 domains), each cut to its first 1,024 tokens, 22,506 positions. Cutting is + exact because the model is causal; a smoke run matched the pod's per-position values for the + same positions (REAP-50, eval-00070: 0.480 local vs 0.465 pod). Domains are equally weighted + here, unlike the pod's position-weighted "all", and only early context (≤1,024 tokens) is + covered. +- Cost: about 4.5 minutes fixed per pass over REAP-50's weights, plus per-token compute. + +## Step 1: non-uniform per-layer budgets + +Question: at a fixed total expert budget, does moving experts between layers beat REAP's +uniform count? + +Allocation rule tried first (`mass-*`): REAP order within each layer is kept, so every variant +is nested inside REAP-50. Across layers, experts are removed greedily by their share of the +layer's routed output mass (Σ g·‖f‖ over calibration, normalized per layer), lowest first, +with a floor per layer. Variants (`~/Models/v41-reap-run/alloc/`): + +| Variant | Mean kept | Per-layer range | Kept mass share (mean) | +|---|---:|---|---:| +| uniform-96 (REAP-75) | 96 | 96 | 0.650 | +| mass-96 | 96 | 24–180 | 0.700 | +| mass-96-f64 | 96 | 64–180 | 0.678 | +| uniform-128 | 128 | 128 | 0.725 | +| mass-128-f64 | 128 | 64–192 | 0.764 | + +A frequency-based rule (`freq-*`) was built and rejected before running: it would keep only +10% of the output mass in layers 0–1. + +Mass-based allocation cuts layers 20–31 and 36–39 hardest, where a few experts carry most of +the output. It treats every layer as equally sensitive, which is the assumption this run tests. + +Results (`~/Models/v41-reap-run/evals/step1`, 22 sessions × 1,024 tokens, one pass of six +variants, 8,211 s, peak 22.5 GB footprint): + +| Variant | KL (domain mean) | top-1 | NLL change | +|---|---:|---:|---:| +| REAP-50 (192) | 0.588 | 0.819 | -0.043 | +| uniform-128 | **0.955** | **0.754** | +0.160 | +| mass-128-f64 | 1.046 | 0.718 | +0.394 | +| uniform-96 (REAP-75) | **1.191** | **0.707** | +0.306 | +| mass-96-f64 | 1.344 | 0.673 | +0.584 | +| mass-96 | 1.457 | 0.636 | +0.815 | + +Per domain (KL): + +| Domain | REAP-50 | uni-128 | mass-128-f64 | uni-96 | mass-96-f64 | mass-96 | +|---|---:|---:|---:|---:|---:|---:| +| agentic (system prompt, see below) | 0.346 | 0.572 | 0.722 | 0.768 | 1.109 | 1.291 | +| bash | 0.265 | 0.507 | 0.599 | 0.671 | 0.888 | 1.053 | +| business | 0.711 | 1.020 | 1.304 | 1.385 | 1.511 | 1.724 | +| business-email | 0.413 | 0.711 | 0.972 | 0.992 | 1.108 | 1.356 | +| code-diffs | 0.263 | 0.461 | 0.625 | 0.647 | 0.813 | 1.049 | +| cyber | 0.566 | 1.049 | 1.414 | 1.518 | 1.953 | 2.060 | +| english-chat | 0.321 | 0.755 | 0.866 | 1.042 | 1.105 | 1.229 | +| english-prose | 0.359 | 0.604 | 0.706 | 0.773 | 0.894 | 0.945 | +| networking | 0.336 | 0.658 | 0.970 | 0.896 | 1.104 | 1.293 | +| tool-use | 2.544 | 3.595 | **2.609** | 3.645 | 3.195 | **2.740** | + +Learnings: + +- **Mass-based reallocation is worse than uniform**, by 0.09 KL at 128 and 0.15–0.27 at 96, on + every domain except tool-use. The layers it cuts hardest (20–31, 36–39) are more sensitive than + their output-mass share says, so "low mass share = cheap to cut" is false across layers. Within + a layer REAP's ranking holds; across layers, mass share is not a sensitivity measure. +- **Tool-use runs the other way.** Mass allocation cuts tool-use KL by ~1.0 at both budgets + (3.60 → 2.61 at 128, 3.65 → 2.74 at 96), even though it hurts everything else. Tool-use leans on + the experts uniform pruning removes from the early layers (mass keeps more there). Tool-use is + also the worst domain under every variant (KL 2.5 even at REAP-50), so a domain-aware budget + (or keeping tool-use experts in the early layers) is a candidate, not a global reallocation. +- The REAP-50 → uniform-128 → uniform-96 curve on the Mac (0.59 / 0.95 / 1.19) has the same + shape as the pod's (0.81 / 1.11 / 1.35). Mac numbers are lower because they cover only the first + 1,024 tokens and weight domains equally. +- **Eval-set flaw found:** all four agentic sessions (agentic-python and agentic-rust) share a + ~3,900-token system prompt, so at 1,024 tokens they score the same text and returned identical + numbers. The "agentic" row is the system prompt only, and the domain mean double-counts it. The + 108 eval sessions have only 100 distinct 1,024-token prefixes. A later pass that needs agentic + coverage must run those sessions past token ~3,900 (about 4× their cost). +- Next for step 1, if it is pursued: measure sensitivity directly (cut one block of layers at a + time and score) instead of inferring it from mass. Deprioritized behind steps 2–3, because the + best case here is a small shift around uniform, while tool-use is the outlier worth chasing. + +## Step 2: KEEP quantization of the pruned model (running) + +- `scripts/runpod/v41_reap/keep_quant.py` fits KEEP's importance-aware E8P codes to REAP-50's + 192 experts per layer from their FP4 weights (the numpy fit with the Metal search, which is + byte-identical to the shipped V4-Flash path). +- V4.1's down projection has input width 2304, which group size 512 does not divide, so down + uses group 256: 2.031 bpw for gate/up, 2.0625 for down, 2.042 averaged. +- Weighted relative MSE (pod importance, layer 0): gate/up 0.084, down 0.047; layer 1: 0.089 / + 0.061. About 500 s per layer while the step-1 eval shares the CPU, so ~5.5 h for 40 layers. +- `local_eval.py --expert-override` runs these codes in the eval: `fp4-bf16` (exact FP4 decode, + the control that isolates removing FP8 activation quantization), a code directory, or a + per-expert mix of two directories. +- Later layers fit slightly worse: layers 16–17 gate/up 0.092, down 0.082; about 350 s per layer + once the CPU eval ended. +- 2-bit fit done for all 40 layers (`quant/e8p-podimp`, about 4.5 h): mean weighted relative MSE + gate/up 0.090, down 0.074, worst layer 17 (0.092). The 1-bit fit started right after. +- The "new importance" half (imatrix measured on the pruned model, 96 calibration records of + 1,024 tokens, `--capture-imatrix`) started when step 1 finished (output + `~/Models/v41-reap-run/imatrix/reap50-pruned.npz`). +- The first capture (all 96 records in one pass) hit the 64 GB guard after 56 minutes: peak RSS + 64.4 GB, of which only 32 GB was the process's own footprint (the rest is mmapped checkpoint + pages, which RSS counts). The importance stats are sums, so it was rerun as four passes of 24 + records (`imatrix/chunk{0..3}.npz`) to be added together. +- Capture done: 76,471 tokens, four passes of 15–38 minutes at ~31.5 GB peak RSS, summed into + `imatrix/reap50-pruned.npz`. 2 of 7,680 surviving experts never routed; 29 routed fewer than + 20 times. +- Pruned-model importance vs the full model's (cosine per expert, mean over routed experts), + against a noise floor from splitting the pod's 24 sessions into two halves: + + | Space | Pod half vs half (noise) | Pruned vs pod | Worst layers | + |---|---:|---:|---| + | gate/up input (hidden) | 0.995 | 0.976 | 15, 21–27 (~0.96) | + | down input (expert intermediate) | 0.895 | 0.788 | 6–9 (~0.73) | + + Both shifts are larger than the noise, and the down input shifts much more. The comparison is + confounded: the Mac records are the first 1,024 tokens of each calibration session, while the + pod summed whole sessions, so data and pruning both differ. The full model cannot run on this + Mac to separate them, so the test is end to end: requantize with the pruned importance only if + the 2-bit codes with pod importance show a quant loss worth recovering. +- Eval started (`evals/step2-fp4-bf16`, then `evals/step2-e8p-podimp`): REAP-50 and uniform-96, + first with the exact FP4 weights in bf16 (control: isolates dropping FP8 activation + quantization), then with the 2-bit codes. + +### Result: KEEP 2-bit on the pruned model (pod importance) + +Same 22 sessions × 1,024 tokens. Each pass ~65 min (two variants), peak RSS 52–54 GB. + +| Experts | Weights | KL | top-1 | NLL change | +|---|---|---:|---:|---:| +| 192 (REAP-50) | FP4 + FP8 activations (shipped) | 0.588 | 0.819 | -0.043 | +| 192 | FP4 decoded to bf16 (control) | 0.578 | 0.821 | -0.028 | +| 192 | KEEP 2-bit, pod importance | **0.747** | **0.788** | -0.092 | +| 96 (REAP-75) | FP4 + FP8 activations | 1.191 | 0.707 | +0.306 | +| 96 | FP4 decoded to bf16 | 1.192 | 0.708 | +0.306 | +| 96 | KEEP 2-bit, pod importance | 1.216 | 0.689 | +0.313 | + +Per domain, REAP-50 (control → 2-bit): agentic 0.33 → 0.42, bash 0.26 → 0.34, business 0.71 → +0.75, business-email 0.42 → 0.48, code-diffs 0.26 → 0.35, cyber 0.55 → 0.78, english-chat 0.33 → +0.46, english-prose 0.36 → 0.41, networking 0.34 → 0.45, tool-use 2.47 → 3.36. + +Learnings: + +- The control matches the shipped path (0.578 vs 0.588), so the override path is sound and FP8 + activation quantization costs nothing measurable here. +- **KEEP 2-bit adds +0.17 KL / −0.033 top-1 on REAP-50**, close to its +0.149 on unpruned + V4-Flash. On REAP-50, REAP and KEEP losses roughly add. +- **On REAP-75 it adds only +0.02 KL**, with the same per-expert codes. The losses are far from + additive once the prune error is large: KL against the teacher is already high, and extra weight + noise moves the distribution in directions that matter less. Two consequences: the quant loss is + worth recovering only at mild prune levels, and the step-2b comparisons (all at high KL) will + favour whatever has less prune error. +- Tool-use takes the biggest quant hit at REAP-50 (+0.89 KL) and even improves slightly at REAP-75 + (3.65 → 3.48). Tool-use is the fragile domain under every change so far. + +Next (queued behind the 1-bit fit): requantize at 2 bits with the pruned-model importance +(`quant/e8p-prunedimp`, blended: pruned importance per expert where it routed ≥20 times, the pod's +otherwise, each normalized to mean 1, floored at 1e-4; 29 experts fall back), on the GPU alongside +the step-2b evals on the CPU. + +## Step 2b: is REAP the right tool, or a low-bit floor for every expert? (Jack's idea) + +Idea: pruning is KEEP's precision ladder with a hard 0-bit rung. Give every expert a low-bit +base instead, and let the runtime raise precision per expert (nested residual stages, hot +residuals in RAM, the rest on SSD). + +Memory facts: one expert slot across all 40 layers costs 0.359 GB at 2.031 bpw, so all 384 +experts are ~138 GB at 2 bpw (does not fit) and ~70 GB at 1.03 bpw (fits). Decode reads 240 +experts per token (6 × 40 layers), so speed depends on bits read, not bits stored. + +### Residency replay (done, CPU, no model) + +Recorded full-model routing from the 24 calibration sessions (49,152 tokens, first 2,048 +tokens each) replayed through a per-layer hot set of 1-bit residual stages (4.6 MB per expert +per layer), resetting at each session: + +| Hot set / layer | Residual RAM | Decayed-count miss rate | Misses / token | Gate mass on misses (LRU, 6 sessions) | +|---:|---:|---:|---:|---:| +| 48 | 8.8 GB | 26.5% | 63.5 | 27.8% | +| 96 | 17.7 GB | 13.6% | 32.7 | 14.5% | +| 128 | 23.6 GB | 9.2% | 22.0 | 8.8% | +| 192 | 35.3 GB | 4.4% | 10.5 | 3.7% | +| 256 | 47.1 GB | 2.0% | 4.9 | 1.7% | + +- Fetching every miss from the SSD (~0.7 ms per 4.6 MB at 6–7 GB/s) costs ~23 ms/token at 96 + and ~7 ms at 192, against a ~32 ms budget at 31 tok/s. Fetch-on-miss only works at 192+ hot, + i.e. 70 GB base + 35 GB residuals, which does not leave room on a 128 GB machine. +- The workable form is serve-from-base: a miss runs at the 1-bit base stage with no stall, and + the residual is promoted in the background. At hot 128 (23.6 GB) about 9% of the routed gate + mass would run at base precision. +- LRU and decayed counts behave alike; a static REAP-ordered hot set is far worse (58% misses at + 128). REAP ranks by mean contribution per routed token, not by frequency. +- Side finding: under the full model's routing, REAP-50's 192 survivors receive only ~60% of + routed slots (the static-192 miss rate is 40%). REAP-50 reroutes 40% of all expert slots. + +### Static frontier test (queued) + +Inside REAP-50, REAP-75's 96 experts per layer are "hot" and the other 96 "cold": + +| Variant | Memory (relative) | Question | +|---|---:|---| +| hot 96 @ 2.03, cold dropped (REAP-75 + KEEP) | ~195 | prune baseline | +| all 192 @ 1.03 | ~198 | REAP vs a low-bit floor at equal memory | +| hot 96 @ 2.03, cold 96 @ 1.03 | ~294 | the cost of a cold expert at 1 bit vs 0 | + +KEEP's 1-bit tier is the E8 codebook with 8-bit indices (`--code-bits 8`): 1.031 bpw gate/up, +1.0625 down. Smoke fit on layer 0: weighted relative MSE 0.32 against 0.094 at 2 bits, so the +1-bit floor is coarse at the weight level; the KL test decides. + +The 1-bit fit finished in 3.3 h (`quant/e8-1bit-podimp`, much faster than the smoke estimate once +the CPU was free): mean weighted relative MSE gate/up 0.311, down 0.253, about 3.4× the 2-bit +error. Evals running: `evals/step2b-1bit` (REAP-50 and uniform-96, all experts at 1 bit), then +`evals/step2b-mixed` (REAP-75's 96 at 2 bits, the other 96 at 1 bit). The 2-bit row of the table +is `evals/step2-e8p-podimp` uniform-96 (KL 1.216). + +## Step 2c: precision as a learned lever (static allocation first) + +Framing (Jack, 2026-11-08): precision should be a routed, per-token dimension of the MoE, with +nested codes, so that pruning (the 0-bit rung), quantization and offloading become one +optimization: minimize E[KL] + λ·E[bits read per token] subject to bits stored ≤ RAM. Static +decisions are how many rungs each expert stores; dynamic decisions are which rung each routed +expert runs at for a token. + +Evidence from this campaign that per-axis budgets are the wrong shape: + +- Losses do not add: the same 2-bit codes cost +0.17 KL on REAP-50 and +0.02 on REAP-75. +- Proxies fail: moving experts between layers by mass share lost to uniform (step 1). +- Allocation is input-dependent: the reallocation that hurt every domain helped tool-use by ~1 KL. +- Decode reads 240 experts per token, so bits read (bandwidth) and bits stored (RAM) are separate + levers. + +Closest prior work found so far (not a full literature search): HOBBIT (arXiv 2411.01433) serves +cache-miss experts from low-precision copies with layer-level prefetch (up to 9.93× decode on +consumer GPUs); Mixtral offloading, Pre-gated MoE, Fate and ProMoE predict next-layer experts; +Matryoshka-style quantization trains nested codes; HAWQ / Bayesian Bits learn per-layer +bit-widths. The combination tested here, where the router itself chooses (expert, rung) under a +bits penalty with nested codes, has not been found in that set. + +Test ladder: + +1. Static allocation without training (Mac). `scripts/runpod/v41_reap/precision_alloc.py`: + within each layer, rungs 0/1/2 bits under a fixed unit budget (1-bit expert = 1 unit), greedy + by estimated error reduction per unit, cost = expert's routed output-mass share × relative + output error at that rung (layer-mean KEEP errors; dropped = 1). Uniform budget across layers. +2. Check the estimator against measured KL (step 2b points). +3. Router over (expert, rung) trained by KD with a bits penalty (pod; merges with step 4). +4. Trainable nested codebooks (pod; merges with step 5). + +Step 1 allocations (`~/Models/v41-reap-run/alloc/prec-{192,288}`): + +| Budget / layer | Memory equivalent | Kept (mean) | at 2 bits | at 1 bit | Overlap of 2-bit set with REAP-75 | +|---:|---|---:|---:|---:|---:| +| 192 | REAP-75 at 2 bits | 101 | 91 | 10 | | +| 288 | REAP-75 at 2 bits + rest at 1 bit | 147 | 141 | 6 | 54% | + +The estimator does not favour a 1-bit floor. Routed mass is concentrated enough that upgrading a +heavy expert to 2 bits buys more estimated error than keeping a light one at 1 bit, so it drops +the light tail and keeps almost everything else at 2 bits. Its 2-bit set also differs from REAP's +(54% overlap), because it ranks by total routed mass (frequency × contribution) where REAP ranks by +contribution per routed token. Step 2b's "all 192 at 1 bit" eval tests the floor directly; the +step-2c evals (`evals/step2c-prec-{192,288}`) are queued behind it. If the floor beats these, +mass × error is the same kind of failed proxy as in step 1 (it ignores rerouting and rare but +critical experts), and allocation has to be learned against KL. + +### SSD churn under admission thresholds (residency replay, CPU) + +Same 24 calibration sessions (49,152 tokens of full-model routing), decayed-count hot set per +layer, 4.6 MB per residual stage. A miss whose gate weight is below the threshold is served at the +base rung and never promoted; misses at or above it are read from the SSD. Gate weights sum to 1.5 +per token-layer (route_scale); median 0.22. Output `~/Models/v41-reap-run/residency-admission.txt`. + +| Hot / layer | Threshold | Promotions / token | SSD MB / token | GB/s at 30 tok/s | Gate mass never promoted | +|---:|---:|---:|---:|---:|---:| +| 96 | none | 33.0 | 152 | 4.55 | 0 | +| 96 | 0.2 | 15.2 | 70 | 2.09 | 5.7% | +| 128 | none | 22.6 | 104 | 3.12 | 0 | +| 128 | 0.2 | 10.4 | 48 | 1.44 | 4.4% | +| 128 | 0.3 | 3.4 | 16 | 0.47 | 16.5% | +| 192 | none | 11.8 | 54 | 1.63 | 0 | +| 192 | 0.2 | 6.2 | 29 | 0.86 | 3.2% | +| 192 | 0.3 | 3.0 | 14 | 0.42 | 15.9% | + +- Promoting every miss is not sustainable below 192 hot: 3–4.5 GB/s of continuous reads against + ~5.4 GB/s sequential, before any prefill traffic. +- A gate-weight threshold of 0.2 halves the traffic at every hot size for 3–6% of gate mass left + at base precision. At hot 128 (23.6 GB of residuals) that is 1.4 GB/s, about a quarter of the + SSD. 0.3 cuts traffic another 3× but leaves ~16% of the mass unpromoted. +- The gate weight is a usable admission signal, and it is the same quantity a learned router + would weigh against a bits penalty: low-weight picks contribute little, so they are the cheap + ones to run at base. Whether 4–6% of gate mass at 1 bit is affordable is what step 2b measures + (all 192 at 1 bit is the 100% case). +- Not yet replayed: next-layer prefetch (needs hidden states, not just routes) and second-hit + admission. + +### Step 2b results: the static frontier + +Same 22 sessions × 1,024 tokens. Units are per-layer memory in 1-bit-expert equivalents +(2-bit expert = 2, FP4 ≈ 4.25). + +| Configuration (inside REAP-50) | Units / layer | KL | top-1 | NLL change | +|---|---:|---:|---:|---:| +| REAP-75's 96 at 2 bits, rest dropped | 192 | **1.216** | **0.689** | +0.313 | +| all 192 at 1 bit | 192 | 1.358 | 0.637 | +0.425 | +| REAP-75's 96 at 1 bit, rest dropped | 96 | 1.826 | 0.565 | +0.943 | +| REAP-75's 96 at 2 bits, other 96 at 1 bit | 288 | 0.954 | 0.738 | +0.010 | +| all 192 at 2 bits | 384 | 0.747 | 0.788 | -0.092 | +| all 192 FP4 (bf16 decode) | ~816 | 0.578 | 0.821 | -0.028 | + +Per domain (KL), all-192 at 1 bit / mixed 288: agentic 1.01 / 0.55, bash 0.82 / 0.45, business +1.80 / 1.06, business-email 1.16 / 0.75, code-diffs 0.93 / 0.47, cyber 1.57 / 1.09, english-chat +1.59 / 0.87, english-prose 0.90 / 0.56, networking 0.93 / 0.60, tool-use 3.22 / 3.56. + +Learnings: + +- **At equal memory, pruning plus 2 bits beats a 1-bit floor**: 1.216 vs 1.358 KL. The static + version of "no REAP, every expert at low bits" loses. The step-2c estimator predicted this + direction (drop the light tail, keep the rest at 2 bits). +- **KEEP's 1-bit tier is expensive**: all of REAP-50 at 1 bit costs +0.78 KL over FP4, against + +0.17 at 2 bits. The 1-bit E8 codebook (weight error 0.31) is too coarse to be a floor that runs + a large share of the routed mass. +- **But a 1-bit rung beats dropping**: adding REAP-50's other 96 experts at 1 bit to REAP-75 at 2 + bits cuts KL from 1.216 to 0.954. Upgrading those same 96 from 1 to 2 bits cuts it further to + 0.747. Per unit of memory the two steps are worth about the same (-0.26 and -0.21 KL per 96 + units), and upgrading REAP-75's 96 from 1 to 2 bits is worth more (-0.61). Marginal value per + unit: hot 1→2 > cold 0→1 ≈ cold 1→2. That is the shape a learned allocator should find, and it + is consistent with the mass × error estimator. +- Frontier so far (units, KL): (96, 1.83), (192, 1.22), (288, 0.95), (384, 0.75), (~816, 0.58). +- For the serve-from-base runtime: running a miss at the 1-bit base is not free. The all-1-bit + pass is the 100%-of-mass case (+0.78 KL); at the 4% never-promoted mass of threshold 0.2 the + cost is unmeasured and should not be assumed linear. A better low rung (1.5 bits, or a 1-bit + base with a small learned correction) would change this picture more than any allocation rule. +- **Tool-use anomaly**: tool-use KL is 2.5–3.6 under every variant and does not order with + memory (1-bit uniform-96 at 3.06 beats 2-bit uniform-96 at 3.48), and its NLL change on REAP-50 + is -1.1 to -1.5 (the pruned model predicts the actual tokens much better than the teacher). That + pattern looks like a measurement problem in those two sessions' teacher logits (alignment or + template), not a model effect. To check before tool-use numbers drive any decision. + +### Tool-use anomaly: checked + +- Teacher logits are aligned with the tokens (argmax matches the next token 56–83% of the time, + the current token 0%). The pod's own REAP-50 run shows the same gap as the Mac, so it is not a + local bug. +- On the pod data the pruned model beats the full model on actual-token NLL in the tool and code + domains: tool-use -1.22 nats, agentic -0.33, english-chat -0.27, bash -0.16; it loses on prose, + business and cyber (+0.3 to +0.4). On tool-use session eval-00018 the full model's NLL is 3.9–5.2 + nats per token in every position band, against 1.4–4.3 for REAP-50. +- The full model has catastrophic positions: 1.4% of all eval positions have teacher NLL > 15 + nats (3.6% in agentic), on trivially predictable tokens (for example `"title": "` then + `Content`, teacher 21.8 nats, REAP-50 0.0). REAP-50 has them too, but fewer (2,199 vs 2,895). +- Those positions carry only ~4% of the measured KL; dropping them moves the domain means by + 0.015–0.02 (REAP-50 0.578 → 0.563, 2-bit 0.747 → 0.725, 1-bit 1.358 → 1.335). The frontier + conclusions do not change. Tool-use KL stays ~2.4 without them, so the tool-use gap is broad, + not just spikes. +- Reading: on DSML tool-call text the full model, run through the reference torch path with the + FP4/FP8 kernels, is far worse than a frontier model should be, and removing experts improves it. + Either some experts misbehave in this numeric path on that text, or the full model really is + weak there. Until that is settled, tool-use KL against this teacher is not a quality measure + (the teacher is the worse model on that domain), and NLL on actual tokens should be reported + next to it. Isolating it needs the full model (pod or base download), so it is noted, not + chased. + +### Step 2c results: static per-expert allocation + +| Configuration | Units / layer | KL | top-1 | NLL change | +|---|---:|---:|---:|---:| +| REAP-75's 96 at 2 bits (prune + KEEP) | 192 | 1.216 | 0.689 | +0.313 | +| all 192 at 1 bit | 192 | 1.358 | 0.637 | +0.425 | +| **prec-192** (≈101 kept: 91 at 2 bits, 10 at 1 bit) | 192 | **1.014** | **0.732** | +0.158 | +| REAP-75's 96 at 2 bits + other 96 at 1 bit | 288 | 0.954 | 0.738 | +0.010 | +| **prec-288** (≈147 kept: 141 at 2 bits, 6 at 1 bit) | 288 | **0.809** | **0.773** | -0.039 | +| all 192 at 2 bits | 384 | 0.747 | 0.788 | -0.092 | + +Per domain (KL), REAP-75@2 → prec-192: agentic 0.73 → 0.55, bash 0.71 → 0.52, business 1.51 → +1.38, business-email 1.07 → 0.78, code-diffs 0.72 → 0.50, cyber 1.53 → 1.09, english-chat 1.19 → +1.19, english-prose 0.82 → 0.67, networking 0.90 → 0.68, tool-use 3.48 → 3.25. + +Learnings: + +- **At REAP-75's memory, the allocator cuts KL by 0.20 (−17%) and lifts top-1 by 4.3 points**, + better or equal on every domain. At 288 units it beats the REAP-ordered mix by 0.145 and comes + within 0.06 of all-192-at-2-bits with 75% of the memory. +- The estimator is a crude proxy (layer-mean errors, mass × error, no training), so this is a + floor on what a learned allocation can do, not a ceiling. +- Open question before crediting "precision allocation": prec-192 is mostly a different choice + of which ~96 experts get 2 bits (58% overlap with REAP-75), plus a small 1-bit tail. Two things + changed at once: the ranking criterion (total routed mass instead of REAP's contribution per + routed token) and the routing it is measured on (the pruned model's own instead of the full + model's). Ablation running (`evals/step2c-selection`, all 96 at 2 bits): full-model total mass, + pruned-model REAP score, pruned-model total mass. + +### Step 2c selection ablation + +All variants keep 96 experts per layer at KEEP 2 bits (`e8p-podimp`, 192 units), so only the +choice of experts differs. Domain-mean KL over the same held-out eval set: + +| 96 experts chosen by | Overlap with REAP-75 | KL | top-1 | +|---|---:|---:|---:| +| REAP score, full-model routing (shipped REAP-75) | 1.000 | 1.216 | 0.689 | +| total routed mass, full-model routing | 0.664 | 1.191 | 0.707 | +| REAP score, pruned-model routing | 0.784 | 1.188 | 0.701 | +| **total routed mass, pruned-model routing** | 0.580 | **1.039** | **0.733** | +| prec-192 (the same ranking, plus a 1-bit tail) | — | 1.014 | 0.732 | + +Learnings: + +- **Selection is most of the step-2c gain.** Of prec-192's 0.202 KL cut, 0.177 comes from + choosing the 96 experts by total routed mass measured on REAP-50's own routing; the 1-bit tail + adds 0.025 (plus 0.023 on tool-use, the domain where it helps most). +- **The two changes only pay together.** Each alone moves KL by about 0.025; combined they move + it by 0.177. Ranking by total mass on the full model's routing does not see how routing + redistributes once half the experts are gone, and REAP's per-route score on the pruned model + still favors experts that are strong when picked but rarely picked. +- So a second prune pass measured on the first pruned model (a cheap form of iterative pruning) + is worth more than KEEP's precision trade at this budget. The 1-bit rung is still a real but + small positive at 192 units, and grows at 288 (prec-288 0.809 vs mixed 0.954, of which the + ranking accounts for an unmeasured share). +- The mass capture came from calibration records (first 1,024 tokens of the 24 calibration + sessions), not the eval set, so this is not leakage. + +### Step 2 closed: pruned-model importance for the KEEP fit + +The 2-bit refit with importance measured on REAP-50 itself (`quant/e8p-prunedimp`) was evaluated +against the pod-importance fit (`e8p-podimp`), same eval set: + +| Variant | Importance | KL | top-1 | NLL change | +|---|---|---:|---:|---:| +| REAP-50, all 192 at 2 bits | pod (full model) | 0.747 | 0.788 | -0.092 | +| REAP-50, all 192 at 2 bits | pruned model | 0.753 | 0.790 | -0.106 | +| prunedmass-96 at 2 bits | pod (full model) | 1.039 | 0.733 | +0.154 | +| prunedmass-96 at 2 bits | pruned model | 1.027 | 0.734 | +0.159 | + +Learnings: + +- **No effect.** The differences (+0.006, -0.012) are within run-to-run noise, and the + prunedmass-96 gain sits almost entirely in tool-use (3.42 to 3.28), the domain whose KL is not + a quality measure against this teacher. Without tool-use the two fits tie. +- So the input importance KEEP weights its fit by barely moves after pruning (hidden-space cosine + 0.976 in step 2 already said so), while the routing mass used for selection moves a lot (step 2c + ablation). Recomputing statistics on the pruned model matters for which experts to keep, not + for how to quantize them. +- Peak RSS 56.6 GB, 48 minutes for two variants. + +### Step 2d: iterative pruning (192 → 144 → 96) + +First pass: the top 144 experts per layer by total routed mass on REAP-50's routing +(`alloc/iter/pm-144.json`). Routing recaptured on REAP-50 masked to those 144 (same 96 +calibration records in 4 chunks, 13–15 minutes each, about 31 GB peak RSS, +`imatrix/pm144-chunk*.npz`), then the top 96 of the 144 by the new mass (`alloc/iter/iter-96.json`). +The second pass changes little: iter-96 overlaps the one-pass 96 by 97.3% (minimum 94.8% in any +layer, about 2.6 experts per layer). Eval at KEEP 2 bits (`e8p-podimp`), `evals/step2d-iterative`: + +| Configuration | Units / layer | KL | KL without tool-use | top-1 | +|---|---:|---:|---:|---:| +| one-pass mass, 96 at 2 bits | 192 | 1.039 | 0.801 | 0.733 | +| **two-pass mass, 96 at 2 bits** | 192 | 1.016 | 0.795 | 0.731 | +| prec-192 (one-pass ranking + 1-bit tail) | 192 | 1.014 | — | 0.732 | +| REAP-75's 96 at 2 bits + other 96 at 1 bit | 288 | 0.954 | 0.693 | 0.738 | +| **one-pass mass, 144 at 2 bits** | 288 | 0.812 | 0.566 | 0.770 | +| prec-288 (≈141 at 2 bits + 6 at 1 bit) | 288 | 0.809 | 0.564 | 0.773 | + +Learnings: + +- **Iterative pruning saturates after one pass.** The second pass gains 0.023 KL, below the + 0.03 bar set before the run, and 0.006 of it outside tool-use. One re-measurement on the pruned + model captures nearly all of the routing redistribution. +- **At 288 units the 1-bit rung adds nothing either:** 144 experts at 2 bits score 0.812 against + prec-288's 0.809. prec-288's lead over the REAP-ordered mix is all selection. +- Net for step 2: under the current codes, the whole gain over REAP-75 + KEEP comes from + re-ranking experts by total routed mass on the pruned model's routing. The static 1-bit rung + helps only marginally (0.025 at 192 units, nothing at 288). A low rung that pays needs to be + better than this 1-bit E8 fit (errors 0.31/0.25), or dynamic, or both. +- Step 2 is done; step 3 (residual sidecars) is next. + +## Precision as a routed dimension (the paper track) + +Working title: *Precision as a Routed Dimension: unifying pruning, quantization and offloading in +MoE inference*. Jack, 2026-11-09: build, run, test, commit; this is the headline to reach. + +### Claim to test + +Static compression picks one precision per expert (or drops it) before inference. The router +already decides, per token, how much each expert matters (its gate weight). Make precision a +routed decision too: every (token, expert) slot runs at one of three tiers chosen at inference +time, from codes nested so that each tier is a prefix of the next: + +| Tier | Weights | Lives in | Cost per use | +|---|---|---|---| +| 0 (drop) | none | nowhere | 0 | +| 1 (base) | 1-bit E8 code | unified memory, always resident | 4.6 MB read | +| 2 (base + residual) | 1-bit base + 1-bit E8 residual of it | residual in an LRU cache, misses on SSD | 9.2 MB read | + +Pruning is tier 0 (per expert or per token), quantization is the tier-1/tier-2 choice, and +offloading is where the tier-2 residual lives. One policy (gate thresholds `tau_drop`, `tau`, and +the cache size K) sets all three, and three costs fall out of it: resident memory (bases + K +residuals per layer), SSD traffic (residual admissions per token), and unified-memory bytes read +per token (the decode speed bound on Apple Silicon, 546 GB/s on the M4 Max). + +The cache never stalls decode: a slot gets tier 2 only if its residual is already resident; a slot +with gate >= tau admits its residual by an async SSD read that serves later tokens. This is the +HOBBIT-style "use a cheaper copy on a miss" (arXiv 2411.01433), with the cheaper copy being the +prefix of the expensive one, so a hit costs only the residual's bytes and memory holds no +duplicate. + +### Evidence so far + +- Nested codes are almost free: on layer 5 (2 experts), base + residual scores 0.103 (gate/up) and + 0.083 (down) weighted relative MSE, against 0.090 / 0.074 for a flat 2-bit E8P fit (all-layer + means). The base alone is 0.31 / 0.25. +- Static per-expert allocation (step 2c) beat REAP-75 + KEEP by 0.20 KL at equal memory, but the + ablation (2c, 2d) traced almost all of it to selection; a static 1-bit tier added 0.025. The + routed version has to beat that static frontier to carry the paper. + +### Tooling + +- `keep_quant.py --residual-of DIR`: the nested residual tier (`quant/e8-nest-r1`, all 192 experts + of REAP-50, running). +- `local_eval.py --dynamic BASE+RESIDUAL --tau NAME=t --tau-drop NAME=t --cache-k NAME=K + --cache-warm NPZ`: per-slot tiers in the forward, with the LRU simulated in token order per + record (`residual_cache`, unit-tested). Writes `tier-counts.npy` per variant: per layer and + expert, slots and gate mass at each tier, and admissions. + +### The experiments that decide it + +1. Endpoints: all tier 2 (tau 0), all tier 1 (tau 1e9). Must land near flat 2-bit (0.747) and + flat 1-bit (1.358). +2. Routed tiers, everything resident (no SSD): tau sweep. Bytes read per token against KL, + compared with static allocations at equal bytes. +3. Routed tiers with offload: K in {48, 96} residuals per layer, tau in {0.2, 0.3}, warm-started + with the pruned-model hot set. Resident units = 192 + K. The bar: beat prec-288 (0.809 at 288 + units, nothing streamed) at 288 units with an SSD rate the M4 Max sustains (about 5 GB/s + sequential; well under that at 30 tok/s). +4. Per-token tier 0 (tau_drop): does skipping low-gate slots cost less KL than static pruning + saves? + +### Observation: two routing regimes by depth + +Share of REAP-50's routed mass (pruned-model capture) held by each layer's top 96 experts: +layers 0–19 keep 0.72–0.83, layers 20–39 keep 0.94–0.98 (layer 32 0.86 is the one dip). Early +layers route diffusely, late layers concentrate. The step-3 sidecar fit sees the same split: in +layers 4–6 the 96-expert 2-bit student misses 25–31% of the teacher MoE's output energy on held-out +tokens, and a rank-256 linear sidecar removes only 9–15% of it, while in layers 0–3 the gap is +3–11% and the sidecar removes 36–76%. A fixed per-layer budget cannot follow this; per-token tiers +can, since a diffuse layer simply routes more of its slots to the base tier. + +### Headline progress rubric + +Every major status update reports progress to the headline with this fixed rubric (weights sum +to 100; partial credit only for measured progress). + +| # | Item | Weight | Score | State (2026-11-09) | +|---|---|---:|---:|---| +| 1 | Nested codes, all layers, validated | 10 | 10 | fit 40/40; error 0.102 gate/up, 0.084 down; end to end +0.053 KL against flat 2-bit | +| 2 | Per-token routed tiers, implemented and tested | 10 | 10 | `--dynamic/--tau/--tau-drop`, smoke-tested | +| 3 | Offload cache simulation and cost model | 10 | 10 | `residual_cache` (unit-tested), `routed_report.py` | +| 4 | Core result: routed beats the static frontier at equal resident memory, SSD rate sustainable | 25 | 20 | E1 (2026-11-10): an expert cache with re-route beats the static set by 0.056 at 144 and 0.231 at 96, intervals clear of zero, 0.45 and 1.72 GB/s at 30 tok/s; inside a pool of 192, 22 sessions, simulated reads. Pool of 384 (2026-11-10): a cache of 144 beats the static 144 by 0.309 (21 of 22 sessions) and the exact REAP-50 checkpoint by 0.12, at 0.7 GB/s for tau 0.3. The rest needs the pod confirmation (E7) | +| 5 | Unification: per-token tier 0 + expert pruning in one policy, frontier figure | 10 | 0 | tier 0 measured free at 7.7% of slots; no joint policy or figure yet | +| 6 | Real Apple Silicon streaming measured, not simulated | 10 | 10 | 2026-11-10: the no-wait cache runs under oMLX on the M4 Max over the full pool of 384 with 144 2-bit lattice residents at 23–27 tok/s (17–19 with thinking on), 80 MB of real reads per token, 63 GiB. Memory pressure (2026-11-12): with another process holding 40 GB of incompressible memory and 20 GB pushed to swap, the server keeps its 83 GiB, answers all 61 requests and decodes at 33.5 tok/s against 36.7 before | +| 7 | Learned tier policy or trained nested codes | 10 | 0 | pod work, after 4 | +| 8 | Generality: a second MoE plus downstream benchmarks | 10 | 0 | | +| 9 | Paper draft | 5 | 1 | framing and experiment plan in this log | +| | **Total** | 100 | **61** | | + +### 2026-11-09: prior art, plan, restart + +- A novelty search found MatMoE (arXiv 2611.01704, 1 Nov 2026: nested codes with per-expert + bit-width chosen from routing scores at run time), DynaExq (2511.15015: hot set at high precision + with async promotion), DyMoE (2603.19172), HOBBIT and MC-MoE. Per-token precision with nested + codes is no longer new by itself. What is left to claim: the nested residual as the offloaded + unit, tier 0 in the same policy, and the regime at 2 bits and below. Details and the full plan to + 100%: [2026-11-09-routed-precision-plan.md](2026-11-09-routed-precision-plan.md). +- The session restart killed both fits (sidecar at layer 9, nested at layer 8). The nested fit + resumed; the sidecar fit is sequential and restarted from layer 0. Jobs now run detached from + `~/Models/v41-reap-run/jobs/` so a session exit cannot stop them. +- 1-bit codes are stored as uint8 from now on (`keep_quant.py`); the existing 48 layer files were + converted in place and verified value-identical, freeing 38 GB. +- The first routed sweep swaps the τ = 0.41 variant for `s96`, a static hot set of 96 residuals with + no admissions: the same codes and memory as the routed cache, so the comparison isolates routing. + +- Headline changed (Jack, 2026-11-09): *Offload the Residual, Not the Expert: Unifying Pruning, Quantization and Offloading in MoE Inference*. The rubric and its weights are unchanged; item 4's + comparison against a separate-copy cache is now the title's own claim. +- Survey of the older KEEP documents, this campaign's findings with bootstrap intervals, and the + candidate levers X1–X7: [2026-11-09-keep-findings-and-levers.md](2026-11-09-keep-findings-and-levers.md). + Bootstrap (positions, paired): selection on pruned routing −0.177 KL [−0.200, −0.148]; second + prune pass −0.024 [−0.034, −0.012]; static allocator at 288 against 144 at 2 bits −0.003 + [−0.009, +0.005]. +- Harness additions for W4: separate-copy high tier (`--dynamic BASE|HI`), oracle cache policy, + projection-level rungs (`--hi-proj`), admission traces (`--trace`), expert masks under + `--dynamic`, and `routed_report.py --compare`. Passes 2, 2b and 3 are queued behind pass 1 + (`~/Models/v41-reap-run/jobs/queue2.sh`, 24 variants, about 11 h). + +### 2026-11-09: X1 on a pod (prune twice at the first cut), and X2 closed + +- X2 (the pruned model's better tool-use NLL) is an artifact, settled from the pod's stored + outputs: the advantage exists only on positions where the teacher is surprised by off-format text + inside the records. Details in the findings note. +- X1 started. Design: mask the full model's router to the top 288 experts by total routed mass + (`x1/fm-288.json`, 97.5% of mass), run calibration on that masked model (`collect.py --mask`, 2 of + 24 sessions, about 256K tokens), keep the top 192 by the new mass (`pt-192`), repeat once to 96 + (`pt-96`), then score five students against the teacher on 54 eval sessions: shipped REAP-50, + one-pass total mass at 192 (`fm-192`, the criterion alone; it overlaps REAP-50 by 76.5%), + `pt-192`, the Mac's 96 inside REAP-50 (`pm-96`), and `pt-96`. Script: `x1_pod.sh`. +- Runpod, US-NE-1 (H100 SXM Low there, RTX PRO 6000 as fallback, CPU pods available): network volume + `yqazs7aqe6` (600 GB) and CPU pod `hd0czk7klgonpi` (4 vCPU, $0.12/h) for the 510 GB fetch, the + wheels and the inputs, so the GPU pod starts on GPU work. Estimated GPU time 6.5–7.5 h at + $3.99/h. The volume stays for W4b. + +### 2026-11-09: step-3 sidecars closed (worse end to end), and an eval-set caveat + +- The 4,096-token sidecar fit (rank 256, 62 records balanced over 13 domains, about 150K tokens) + stopped at layer 15 of 40: peak RSS 81.8 GB passed the 80 GB guard right after layer 14. Layers + 1 and 14 are the two Engram layers; the pages their lookups touch stayed mapped. `sidecar_fit.py` + now unmaps each table after its layer and resumes from saved layers (resume not yet exercised). +- Before resuming, the 15 fitted layers (0–14) were scored alone (`evals/step3-sidecar-l15`; + `local_eval.py --sidecar` now accepts a partial fit). Against the same 96 experts at 2 bits + without sidecars: KL 1.238 against 1.039, top-1 0.705 against 0.733; paired bootstrap + **+0.199 [+0.164, +0.219]**. The median per-position change is 0.0006. The harm is in + english-prose (+0.71, +1.15), english-chat (+0.34, +0.59) and business (+0.37, +0.48); the code, + agentic and tool-use sessions move by −0.11 to +0.12. +- Layer-local holdout numbers said the opposite (holdout residual energy down 8–19% in the + large-gap layers, 79–93% in layers 8, 10 and 14). Cause not established. Energy-weighted error + on the fit's own records is not KL; an application mismatch between fit and eval is not excluded, + though it would be expected to hurt every domain. Same shape as the GLM-4.5-Air sidecar: a proxy + win that failed end to end. The fit is not resumed and the lever is closed. +- Eval-set caveat found on the way: four of the 22 Mac eval sessions (`eval-00000`, `-00001`, + `-00004`, `-00005`; both agentic-python and both agentic-rust) have identical first 1,024 + tokens, a shared system prompt. At 1,024 tokens the Mac set is 19 distinct texts and those two + domains measure the same one. Paired differences between variants are unaffected; domain means + at 1,024 tokens overweight that prompt. The 4,096-token confirmation pass (W4 pass 4) is where + those sessions diverge. + +### 2026-11-09: X1 running on an H100, after 70 lost minutes + +- GPU pod `q8vcl6ob0hf0lk` (H100 SXM, US-NE-1, $3.99/h) created 01:01 UTC once the fetch verified; + the CPU pod was terminated. Calibration on the 288-expert masked model started 02:31 UTC. The 90 + minutes between cost about $6 and did no GPU work: + - 15 minutes: my readiness probe compared the proxy's reply to a prefix, and the reply starts + with a terminal escape sequence, so it waited out all its tries on a pod that was already up. + - One failed launch: the Ubuntu 24.04 image marks the system Python as externally managed + (`PIP_BREAK_SYSTEM_PACKAGES=1` in `x1_pod.sh`). + - 50 minutes: the smoke test sat at 0% GPU. The network volume reads one stream at about + 120 MB/s and a random 256-byte row in 2–3 ms; the loader reads shards through one mmap and + the Engram lookup faults in rows one page at a time. Sixteen parallel streams read 3.2 GB/s + across files and 1.1 GB/s inside one file. +- Fix: container disk raised from 50 to 560 GB (restarts the container), and `x1_pod.sh` copies + the checkpoint to the container's own disk with 16 parallel `cp` and checks the byte total + (`LOCAL_HF=/root/v41-hf`; `pod.sh` and `x1_pod.sh` take `HF` from the environment). The copy took + 7 minutes for 476 GB. Smoke then passed: 16,384 tokens in 447 s, 48.1 GB peak GPU memory. +- For W4b and any later pod on this volume: budget the 7-minute copy, create the pod with a + 560 GB container disk from the start, and probe readiness by substring. + +### 2026-11-09: routed sweep, pass 1 (192 bases, 22 sessions, first 1,024 tokens) + +Nested fit finished: 40 of 40 layers, total error 0.102 (gate, up) and 0.084 (down) against 0.090 +for the flat 2-bit fit. Sweep: 12,415 s forward, peak RSS 64.1 GB. Units are 1-bit experts per +layer; SSD rate is at 30 tok/s. + +| Variant | Units | KL | Top-1 | Tier-2 slots | SSD GB/s | +|---|---:|---:|---:|---:|---:| +| all2: every residual resident, all slots tier 2 | 384 | 0.800 | 0.779 | 100% | 0 | +| t22 (M): all resident, tier 2 at gate ≥ 0.22 | 384 | 0.822 | 0.767 | 46% | 0 | +| d10: t22 plus tier 0 below gate 0.1 | 384 | 0.823 | 0.763 | 46% (7.7% dropped) | 0 | +| c96t2 (R): LRU of 96 residuals, τ 0.2 | 288 | 0.832 | 0.772 | 90% | 0.83 | +| c96t3 (R): LRU of 96, τ 0.3 | 288 | 0.837 | 0.771 | 84% | 0.28 | +| s96 (S): static hot set of 96 | 288 | 0.905 | 0.753 | 75% | 0 | +| c48t2 (R): LRU of 48, τ 0.2 | 240 | 0.876 | 0.758 | 76% | 2.72 | +| all1: base only | 192 | 1.358 | 0.637 | 0% | 0 | + +Static references from step 2, same sessions: 192 experts at flat 2 bits 0.747 (384 units); +144 experts at flat 2 bits 0.812 (288 units); 96 at flat 2 bits 1.039 (192 units). + +Paired block bootstrap (`routed_report.py --compare`): + +| A against B | KL(A) − KL(B) | 95% interval | +|---|---:|---| +| R c96t3 against S s96, both 288 units | −0.068 | −0.077 to −0.060 | +| R c96t2 against S s96 | −0.073 | −0.084 to −0.065 | +| R c48t2 (240 units) against S s96 (288 units) | −0.029 | −0.042 to −0.022 | +| R c96t3 against 144 flat 2-bit experts, both 288 units | **+0.025** | +0.009 to +0.051 | +| R c96t2 against 144 flat 2-bit experts | +0.020 | +0.003 to +0.045 | +| all2 against 192 flat 2-bit experts (the nesting penalty) | +0.053 | +0.039 to +0.071 | +| R c96t2 (288 units) against all2 (384 units) | +0.032 | +0.023 to +0.037 | +| t22 against all2 | +0.022 | +0.015 to +0.031 | +| d10 against t22 (tier 0 on 7.7% of slots) | +0.001 | −0.002 to +0.004 | + +What it says: + +- Routing the cache works. With the same codes and the same memory, an LRU with a gate threshold + beats a fixed hot set by 0.07 KL, and a cache of 48 beats a fixed set of 96. At τ 0.3 the cost + is 9.5 MB per token, 0.28 GB/s, a ninth of the 2.5 GB/s ceiling. +- It does not yet beat plain pruning. At 288 units, 144 experts at flat 2 bits score 0.812 and the + routed cache 0.832–0.837. The nesting penalty (+0.053) is larger than that gap: with a flat-quality + high tier the routed cache would be ahead. This is the argument for a separate-copy high tier, + which pass 2b measures, and for the 2-bit-base ladder (X5). +- Tier 0 is free: dropping the 7.7% of slots with gate below 0.1 changes KL by 0.001 and cuts the + expert bytes read per token by 5%. +- Item 4's 10 points need R to beat D as well, so the score waits for pass 2b. Missing static + point: 120 experts at flat 2 bits (240 units), to compare with c48 directly. + +### 2026-11-09: X1 result, re-base decision, plan revision 2 + +- X1 finished in 1.6 h of GPU work (smoke 02:23 UTC, scores 03:59 UTC). Domain-mean KL on 54 eval + sessions: shipped REAP-50 0.716, one-pass total mass 0.532, prune-twice 0.546; at 96 experts the + Mac's pick 1.163 and prune-twice 1.154. Total mass against REAP-50: −0.184, session-level 95% + interval −0.277 to −0.128, better in 36 of 54 sessions; english-chat is worse (0.448 → 0.572). + Tables, intervals and the reading are in the + [findings note](2026-11-09-keep-findings-and-levers.md), section 5.1. Outputs: + `~/Models/v41-reap-run/x1/pod/` (scores, tables, per-position eval files, 2.4 GB). +- Follow-up on the same pod while it was hot (`x1b_pod.sh`): 96 experts in one pass on full-model + routing (`fm96`, keeps 74.0% of mass) against two passes through `fm192` (`pf96`, 81.2% of the + 192's mass). +- Jack approved re-basing on the total-mass 192 and asked for a disk cleanup. Export running on + the pod's CPU alongside the follow-up: `prune.py --skip-shard` (new, tested) leaves out the two + Engram shards, which are byte-identical to REAP-50's, so the export is 152 GB instead of 341. + Network volume grown 600 → 680 GB to hold it (volumes cannot shrink; $48 a month until deleted). +- Deleted on the Mac with Jack's confirmation, 373 GB: `~/Models/v41-reap75` (on Hugging Face), + `quant/e8p-prunedimp`, `sidecar/codes-prunedmass-96`, `sidecar/mac-4k-r256`, `sidecar/smoke`. +- Plan revised ([revision 2](2026-11-09-routed-precision-plan.md)): re-base first (WR), main + ladder 2-bit base + 1-bit residual, W4b redefined, sidecars and prune-twice closed, X3 demoted, + X8 (domain-balanced selection) added, operating rules for pods and evals. Five predictions are + registered in the findings note, section 6, before the runs that test them. + +### 2026-11-09, late: full volume, cancelled sweeps, plan revision 3 + +- Both pod jobs (the 96-expert follow-up and the export) died silently between 04:35 and 04:43 + UTC. Cause: the network volume filled. I sized it in GB against `du` figures in GiB (476 GiB of + model is 511 GB; the 152 GiB export is 163 GB), so 680 GB was 6 GB short. A full volume here + writes zero-length files and returns success, so neither job logged an error, and my status + check (`pgrep -f prune.py` inside an ssh command that itself contains that string) kept reporting + the export as running. About 15 minutes of idle H100. +- Recovery: partial export deleted, volume grown to 720 GB ($50 a month), follow-up eval resumed + alone at 04:51 UTC from its 8 finished sessions (files written before the volume filled were + checked: complete). The export moves to a CPU pod. +- Jack cut the 1-bit base altogether. Passes 2b and 3 cancelled; pass 2 finishes. +- Where the selection gain sits, and the checks behind it: findings note, section 5.1. +- Plan [revision 3](2026-11-09-routed-precision-plan.md): thesis "pruning is a frozen cache", + target architecture, a critique of the approach, ten falsifiable conjectures, eleven ordered + experiments, and a proposed working title conditional on C1. The first experiment (E1) needs no + new fits and runs after pass 2. + +### 2026-11-09/10: C8 refuted, prior art that changes the plan, build plan written + +- 96-expert follow-up (pod, 54 sessions, domain-mean KL): one pass on full-model routing 1.123; + two passes through the total-mass 192 1.150; three passes 1.154; the Mac's pick inside REAP-50 + 1.163. Re-measuring routing made it worse by 0.027 where conjecture C8 predicted at least 0.05 + better. Same confound as before (the multi-pass tables used 12 times less calibration data); no + interval computed yet. C8 is refuted as stated. +- Prior-art search (2026-11-09): oMLX already runs the full V4.1 on a 128 GB Mac with experts on + SSD (waits on every miss, 4.2–5.6 tok/s, drafter with offload allowed since 2026-09-23); + re-routing a miss is published (Cache-prior, SMoE, MoE-CORE, SwapMoE); a nested bit-slice cache + is published (SliceMoE); 2-bit V4.1 builds exist three times. The "frozen cache" framing is not + stated anywhere found, but is too close to SwapMoE to carry a paper alone. +- Jack: a running model is the bar ("we need the decoder and the rest"); oMLX may be run and the + rule against running Rapid-MLX no longer applies. +- [Build plan](2026-11-10-v41-wow-build-plan.md): answers to the three sizing questions with + sources and dates, three rival plans and why two lost, the chosen route (a pinned oMLX fork with + a compact router that re-routes a miss at no time cost), six kill tests for days 0–1, all 40 + inventory items placed, Jack's seven ideas placed, prior art, decisions. +- Sizing facts worth keeping in view: non-expert weights are 9.85 GB as stored; a 2.03-bit expert + is 8.99 MB, so 192 per layer is 69 GB and 384 is 138 GB; the per-token read floor is 10.7–13.0 GB + (31–45 tok/s on this Mac); the "91 GB thrash" was measured on an M5 Max. +- Pod: with Jack's go, a 33-minute eval of static 250 and 288 of 384 against the full model is + running before the pod is terminated. + +### 2026-11-10: what a fixed resident set costs (K3), pod terminated + +- Static sets by total routed mass against the full model, native precision, no cache, 54 eval + sessions, domain-mean KL (pooled in brackets) and top-1 agreement: + + | Experts per layer | Routed mass kept | KL | Top-1 | + |---:|---:|---:|---:| + | 96 | 74.0% | 1.123 (1.223) | 0.709 | + | 192 | 92.5% | 0.532 (0.650) | 0.814 | + | 250 | 95.2% | 0.323 (0.449) | 0.860 | + | 288 | 97.5% | 0.235 (0.357) | 0.887 | + + Tool-use and the two agentic domains carry most of what remains (tool-use 1.058 at 250 and + 0.837 at 288; code-diffs 0.045 and 0.021). Mass kept is a poor predictor of KL: the last 2.5% of + mass still costs 0.235. +- For the build plan: at the 80 GB cap about 160 experts fit, where a fixed set costs roughly + 0.6–0.7 KL by interpolation (not measured); at 250 it costs 0.32. The cap, more than the code + density, sets how much the cache has to recover. +- H100 pod `q8vcl6ob0hf0lk` terminated 06:21 UTC after about 5.3 h ($21). Volume `yqazs7aqe6` + (720 GB) keeps the full model and all eval outputs (`/workspace/x1/eval`, `eval2`, `eval3`). +- Mac: the last sweep on the 1-bit ladder was stopped at Jack's request (17 CPU-hours in, results + not kept). E1 was cut from eight variants to four and is running. + +### 2026-11-10: E1, an expert cache against a fixed set (C1 holds inside the pool of 192) + +- Setup: REAP-50 pool (192 experts per layer), every resident expert at 2 bits, no low-bit copy. + A token routes among the experts resident when it arrives (`--miss reroute`); experts it asked + for above gate 0.2 are admitted afterwards and serve from the next token on. Warm start from the + static top-N. 22 sessions, first 1,024 tokens, domain-mean KL against the full model, paired + bootstrap over sessions. + + | Variant | Residents | KL | Top-1 | Reads at 30 tok/s | + |---|---:|---:|---:|---:| + | static (`rs144`) | 144 | 0.812 | 0.770 | 0 | + | LRU cache (`rr144`) | 144 | 0.756 | 0.785 | 0.45 GB/s | + | static (`prunedmass-96`, step 2c) | 96 | 1.039 | 0.733 | 0 | + | LRU cache (`rr96`) | 96 | 0.809 | 0.777 | 1.72 GB/s | + | per-session oracle set (`ro96`) | 96 | 0.789 | 0.782 | 0 | + +- Path check: `rs144` reproduces the earlier pruned 144 exactly (0.8123 against 0.812), so + masked routing in the new path equals pruning. +- Cache against static at equal memory: −0.056 [−0.066, −0.046] at 144; −0.231 [−0.248, −0.211] + at 96. A cache of 96 equals a static 144 (−0.004 [−0.020, +0.008]): the same quality for two + thirds of the expert memory. A cache of 144 is within 0.009 [−0.001, +0.012] of all 192 + resident (0.747). +- C1 passes its own test (gain above 0.02 at both sizes, intervals clear of zero, reads under + 2.5 GB/s). Not checked: the "confident positions" form of the gain, because the Mac eval files + do not store the teacher's NLL. +- C5: the oracle beats the LRU by 0.020 [0.005, 0.027], above the predicted 0.015. All of it is + warm-up: the gap is 0.050 in the first 256 tokens and 0.002–0.017 after. So admission policy is + near its ceiling, and the one thing worth building is a better starting set (idea 6, the prompt + as a demand profile). +- C6 holds at 144 (0.45 GB/s; 1.6 admissions per token, 15 MB) and misses its 1 GB/s line at 96 + (1.72 GB/s; 6.2 admissions per token, 57 MB). +- C4 untested: the drop variant was cut from the run. +- Where the gain is: english-chat (0.760 to 0.468 at 144; 1.195 to 0.497 at 96), business (0.971 + to 0.763; 1.391 to 0.786). Agentic and tool-use text does not move (tool-use 3.27 to 3.33 at + 144). The static set was chosen on the mass of the whole calibration mix, so the cache helps + where a session differs from the mix; tool-use loss comes from the pool of 192 itself and needs + the pool of 384 (C2). +- The gain is flat along the session (−0.268, −0.233, −0.209, −0.212 by quarter at 96). +- Limits: pool of 192 and not 384; first 1,024 tokens of 22 sessions; reads are counted, not + measured; a fetch is assumed to land within one token. The log's `admissions_per_token_layer` + is divided by the layer count twice; `routed_report.py` and the counts above are right. +- Rubric item 4 moves from 0 to 12 of 25. Total 43. +- Run cost: 6,757 s forward, peak RSS 62.4 GB. Nothing is running on the Mac or on Runpod. + +### 2026-11-10: K1, first generated tokens on the Mac + +- Stock oMLX (pinned clone, kernels built) loads the 192-expert REAP-50 checkpoint as is. With + its expert offload at 48 of 192 experts it decodes at 3.3–4.5 tok/s, 78% hit rate, about 0.9 GB + read per token, 45.4 GiB after load. +- A fixed 48 experts per layer (top routed mass) with the router cut to those rows decodes at + 30–32 tok/s and prefills at 150 tok/s in the same memory; 64 experts run at the same speed in + 56.6 GiB. No checkpoint copy: a linked model directory plus read-side patches + (`scripts/v41_build/omlx_compact.py`). +- oMLX's offload does not use its fused decode kernels; a fixed set does. The cache will swap + rows of a stacked array in place. +- Text at 48 and 64 experts is degraded (64 loops in a long greedy completion), in line with the + static KL numbers. +- Full table and notes: the build plan's K1 result. The server is stopped; nothing is running. +- Rubric unchanged at 43: item 6 needs our own cache streaming on the Mac, not oMLX's. + +### 2026-11-10: the no-wait cache runs + +- `scripts/v41_build/omlx_routed.py`: demand scored on the whole pool in the token's graph, + admission into spare slots by reader threads writing the stacked arrays in place, the router + switched only around finished writes. 64 FP4 residents of 192, tau 0.2, 4 spares. +- Decode 22–26 tok/s with admissions on (8.3 installs and 156 MB read per token, 75% of demanded + slots resident), 27–28 with them off, 30–31 without demand scoring; oMLX's exact offload 4.5. +- Admissions off reproduces the fixed set's 1,024-token output exactly. Slot bytes matched the + stored experts at every check (20 checks, 12,785 installs). +- Generation at 64 residents is still bad: both the fixed set and the cache loop on a prompt the + exact pool of 192 answers correctly. FP4 allows no more under the cap, so 2-bit experts in the + runtime are next (about 155 residents). +- Rubric item 6 moves from 0 to 5. Total 48. Server stopped; nothing is running. + +### 2026-11-10: K2, the lattice kernel against oMLX's formats + +- At V4.1 shape on the M4 Max, per routed expert: E8P lattice 2-bit 0.037 ms, oMLX affine 2-bit + 0.047, oMLX MXFP4 0.042, oMLX affine 3-bit 0.036 (three runs, within 5%). K2 passes. +- The lattice expert is 9.0 MB against 11.1 for affine 2-bit: about 155 residents under the cap + against about 127. Table and limits in the build plan's K2 result. +- Rubric unchanged at 48. Nothing is running. + +### 2026-11-10: 2-bit lattice experts in the no-wait cache + +- `scripts/v41_build/omlx_lattice.py`: 144 lattice residents per layer of the pool of 192 under + stock oMLX, codes read straight from the quantizer's `.npz` files, the router scoring the + whole pool with residency as a bias mask. Self-check against host-decoded matrices within + 0.3%; installs byte-exact (17 checks, 2,431 installs on the real model). +- Decode 26.0–27.9 tok/s with the cache on (2.2 installs and 20 MB read per token, 96% of + demanded slots resident), 28.6–30.1 with the policy off. 62 GiB steady, 68–69 peak. A + 4,341-token prompt prefills at 111–192 tok/s. +- With thinking off the model writes a duplicate-file script that runs correctly and summarizes + a 4,341-token source file accurately. With thinking on it does not reach an answer in 4,096 + tokens, with the cache on or off, greedy or sampled; the cache is not the cause, and 2-bit + codes are not yet separated from pruning. +- Rubric item 6 moves from 5 to 7. Total 50. Server stopped; nothing is running. + +### 2026-11-10: the exact pruned model also fails with thinking on + +- The pool of 192 at FP4 under stock oMLX offload (exact, 7.1 tok/s), same prompt, greedy, + thinking on: 4,096 reasoning tokens and no answer, repeating itself from about character + 10,500 of 16,235. The lattice cache repeats from about 5,300; the fixed lattice set and the + sampled run do not repeat, and none of the four answers. +- So the failure with thinking on is not caused by 2-bit codes or by the cache. It is in the + pruned checkpoint at full precision, or in the model, template or greedy decoding as served + here; the unpruned model has not been run on this prompt. +- Rubric unchanged at 50. Server stopped; nothing is running. + +### 2026-11-10: thinking finishes under sampling, given 12,000 tokens + +- Lattice cache, 144 residents, thinking on, default sampling, budget 12,000: 7,567 reasoning + tokens with no repeated sentence, then a 620-token answer whose script runs correctly. + 400 s, 20.7 tok/s over 8,187 tokens, 62 GiB. +- So with thinking on the model is slow to decide rather than broken, and greedy decoding is + what loops (on the exact pruned model too). Speed at 8,000 tokens of context is a quarter + below the 1,000-token figure; unattributed. +- Rubric unchanged at 50. Server stopped; nothing is running. + +### 2026-11-10: pool of 384, fit prepared, no GPU in the volume's data center + +- The lattice fit had no CUDA path (it ran on the Mac's Metal search). `keep_quant_torch.py` is + the same fit in torch; on 64 rows of nine projections it reproduces the stored codes + (99.99–100% equal, same weighted error). `pool384_pod.sh` runs it on a pod for the 192 + experts per layer that REAP-50 removed and extracts all 384 router rows. +- Importance for the removed experts is built (`imp-pruned.npz`, 228 MB): 17 experts that + calibration never routed get their layer's mean importance. +- No GPU type had stock in US-NE-1, where the volume is; one create of 1×H100 failed on + capacity. Nothing was created and nothing is billing beyond the volume. +- Rubric unchanged at 50. + +### 2026-11-10: pool of 384 on the Mac + +- No GPU was needed. A CPU pod (4 vCPU, $0.12/h, 76 minutes) streamed the 192 removed experts + per layer and all router rows to the Mac (144 GB); the Metal fit made their 2-bit codes in + 2.5 h (weighted error 0.091 on w1 and w3, 0.079 on w2). Pod terminated. +- The lattice cache now takes a pool wider than the loaded checkpoint (`0098ea8`). Over all 384 + experts with 144 residents: 26.6, 23.7 and 22.6 tok/s with thinking off, 19.1 and 17.3 with + thinking on, 63 GiB, 88–91% of demanded slots resident, 6–9 installs per token (1.0–2.1 GB/s + of reads). +- Text: one thinking-off script correct, one with a wrong token that crashes it; thinking on + finishes under sampling (2,704 reasoning tokens, correct script) and repeats under greedy + decoding like the pruned models. Table in the build plan, section "Pool of 384". +- Rubric item 6 moves from 7 to 8. Item 4 still needs the quality comparison over the pool of + 384, which needs the full FP4 model as reference (pod). Total 51. +- Disk: `pool384/fp4-removed` (144 GB of FP4 experts) is no longer needed by the cache and + waits for Jack's word before deletion. Nothing is running on the Mac or on Runpod. + +### 2026-11-10: tau sweep over the pool of 384, on the Mac + +- Jack asked for the sweep on a pod; US-NE-1 (the volume's data center) had no GPU, and he + chose the Mac. `local_eval.py --router` (`46fff81`) scores the full pool on the REAP-50 + checkpoint from the 2-bit codes alone, against the full model's stored logits. +- KL over 22 sessions, 1,024 tokens: all 384 resident 0.390; cache of 144 at tau 0.1, 0.2, 0.3, + 0.45: 0.483, 0.460, 0.454, 0.573; oracle 144 0.424; static 144 0.763. The exact REAP-50 + checkpoint scores 0.578 on the same tokens. +- The cache beats the static set by 0.309 (standard error 0.052, 21 of 22 sessions). Tau 0.3 + equals 0.2 at a third of the reads (2.6 installs per token, 0.7 GB/s at 30 tok/s). +- Serving at tau 0.3: 27.2, 26.7 and 23.0 tok/s, 1.4–2.5 installs per token, both scripts + correct. Tau 0.3 is the serving default. Tables in the build plan, section "Tau sweep over + the pool of 384". +- Rubric item 4 moves from 12 to 20 of 25; the remaining 5 wait for the pod confirmation (E7). + Total 59. +- Run: 11,991 s, peak RSS 63.8 GB. The pod job `tau384_pod.sh` was not needed. Nothing is + running on the Mac or on Runpod. + +### 2026-11-11: plan for four ideas; page cache dead, drafter works, sampling check + +- Jack picked four ideas (drafter on the no-wait cache, higher precision on the hottest experts, + one download for any Mac, the OS page cache as the cache) and approved the plan. He also + pointed at the upstream model card: generation tests now use its settings (temperature 1.0, + top_p 1.0, large budgets) and say so when they depart. +- Page cache (K5, `scripts/v41_build/k5_pagecache.py`, `d6cf35d`): MLX 0.32.4 aliases a file + mapping, but a kernel reading 1 MiB of a 2 GiB aliased buffer brings every page into memory. + Dead in the one-buffer form. +- Drafter: on through `mtp_enabled` with no code change; +26–28% on scripts, +43% with thinking + on, +9% on a long prompt, +7 GiB, at 120 residents with installs byte-checked. Above the 1.25 + gate on three of four requests. A multi-row kernel would add at most about 12% and is parked. +- Sampling at the card's settings: 4 to 7 of 8 sampled scripts run correctly across the cache + (120 and 144 residents, drafter on and off, temperature 1.0 and 0.6) and the exact REAP-50 + checkpoint (5 of 8). Single-token slips; no row separable at 8 samples; greedy scripts all + correct. More samples queued. +- `local_eval.py --hot` (`76384aa`) runs the hottest residents from the checkpoint's FP4 + weights. Smoke (1 session, 96 tokens): cache of 144 1.52, with 32 hot 1.14; all 384 resident + with 32 hot 0.68. The resident-count curve (C1) and the hot-tier comparison (B1) are running. +- Rubric unchanged at 59: no item's own test was met. + +### 2026-11-11: quality against resident count (C1) and FP4 on the hottest experts (B1) + +- C1, pool of 384, tau 0.3, 22 sessions, KL (without tool-use in brackets): 32 residents 0.942 + (0.732), 48 0.790 (0.582), 64 0.705 (0.496), 96 0.605 (0.390), 144 0.454 (0.320), 200 0.401 + (0.279), 250 0.398 (0.254), all 384 0.390 (0.235). A cache of 48 (28 GiB footprint) is level + with the static 144 (60 GiB); a cache of 96 (44 GiB) is level with the exact REAP-50 + checkpoint. Breadth is spent near 200 residents. +- B1, gate of 0.02 over the equal-memory 2-bit rival: 16 FP4 experts at 53 GB −0.041 (−0.010 + without tool-use); 144 with 32 FP4 against 177 at 65 GB −0.029 (−0.015). Met only with the + tool-use sessions counted, which carry the off-format teacher, so not met on the reading to + trust. Trading 34 residents for 32 FP4 experts loses 0.022. +- Precision alone: all 384 resident with 32 FP4 per layer 0.268 against 0.390 (0.167 against + 0.235 without tool-use, all 20 sessions). 144 with 32 FP4 at 65 GB equals 200 at 74 GB. +- Reading: below about 200 residents breadth buys more than bits; above it only bits are left. + The 128 GB budget (92 GB of experts) is the untested case; `c1b1_mac.sh b1b` (217 with 32 + FP4, 183 with 64 FP4, against 250 at 2 bits) is queued behind the serving runs. +- Tables in the build plan, sections C1 and B1. Rubric unchanged at 59. + +### 2026-11-11: serving at 48 and 96 residents (C2) + +- Card sampling, 8 sampled scripts each: 48 residents (30 GiB, 20–27 tok/s) 0 correct, 96 + residents (46 GiB, 22–25 tok/s) 2 correct. With 120 and 144 from earlier today: 0, 2, 4, 6 of + 8. Text at 48 residents is misspelled throughout and most samples never finish. +- KL understates this: the simulator put a cache of 48 level with the static 144 (0.79 against + 0.76), and that level is unusable at temperature 1.0. "One download, any Mac" does not hold + at the card's sampling below about 144 residents on this evidence. +- Queued: temperature 0.6 at 96 and 48 residents, top_p 0.95 at 144 and 96; before those, 16 + more samples for the cache of 144 and the exact checkpoint, and the 92 GB hot-tier run. + + +### 2026-11-11: hot tier at 92 GB (B1b), more samples, rescue checks + +- B1b, the expert budget of a 128 GB Mac: 217 residents with the hottest 32 at FP4 score 0.325 + (0.219 without tool-use) against 0.398 (0.254) for 250 at 2 bits, both 92 GB. Paired: −0.072 + ± 0.027 (17 of 22), −0.036 ± 0.010 without tool-use (15 of 20). The 0.02 gate is met on both + readings, so B continues. 183 with 64 FP4 scores the same (0.328, 0.221). +- 217 with 32 FP4 at 92 GB is at least level with all 384 resident at 2 bits at 139 GB (−0.064 + ± 0.035; −0.016 ± 0.009 without tool-use). Bits lose to breadth at 53 GB, draw at 65 GB and + win at 92 GB. +- Sampled scripts at the card's settings, 24 each: cache of 144 at 2 bits 17 correct, exact + REAP-50 18 correct; the cache runs at 22–27 tok/s against 5.6–6.3. +- Rescue checks, 8 each: 96 residents at temperature 0.6 1 correct, 48 residents 0; top_p 0.95 + gives 5 at 144 and 2 at 96. Neither setting rescues a small cache; the errors are misspelled + identifiers from re-routed slots. The usable floor on this request is about 120 residents. +- Serving the 92 GB configuration needs about 96 GiB of footprint, above the 80 GB rule. +- Tables in the build plan, sections B1b and "More samples". Rubric unchanged at 59: no item's + own test was met. + +### 2026-11-11: 200 residents served + +- Per-process limit raised to 100 GB for serving (Jack). The cache runs at 200 residents: 81 + GiB at start, 87 GiB peak, 90% of demanded slots resident, 0.40 installs per token. +- Sampled scripts at the card's settings: 19 of 24 correct (144 residents 17, exact REAP-50 + 18). No failure is a misspelled name or syntax error; at 144 residents most were. +- Speed not measured cleanly (the residual fit shared the GPU); queued. Rubric unchanged at 59. + +### 2026-11-12: third bit on the hottest experts (B2) + +- 1-bit residual fitted on the 2-bit base for the 192 survivors: weight error 0.084 to 0.030 + (gate, up), 0.047 to 0.024 (down). `local_eval.py --hot-tier DIR` runs hot experts as base + + residual. +- At 92 GB: 218 residents with the hottest 64 at 3 bits score 0.311 (0.215 without tool-use), + against 0.325 (0.219) for 217 with 32 at FP4 and 0.398 (0.254) for 250 at 2 bits. Paired + against FP4: −0.015 ± 0.008, −0.004 ± 0.003 without tool-use. The residual matches the FP4 + copy at equal memory and needs no second copy of the expert. +- 32 hot at 3 bits equals 32 at FP4; 96 hot adds nothing over 64. +- Rubric unchanged at 59: no item's own test was met. + +### 2026-11-12: 180 residents with the drafter + +- Clean speed at 200 residents, drafter off: 21–28 tok/s, the same as at 144. +- 180 residents with the drafter on: 88 GiB peak, 22–41 tok/s (median 32), 21 of 24 sampled + scripts correct at the card's settings. Exact REAP-50 on stock oMLX: 18 of 24 at about 6 + tok/s. Candidate default for a 128 GB Mac. +- Rubric unchanged at 59: item 6 still needs the memory-pressure run. + +### 2026-11-12: slots from RAM, slim checkpoint, one command (C3, C4) + +- `OMLX_LATTICE_SLOTS=auto` is the default: 70% of RAM less 20 GiB and 8 GiB for the drafter, + at 0.336 GiB per resident (183 with the drafter on this Mac); a warning below 120 residents. +- `slim_checkpoint.py`: 365.8 GB to 225.9 GB, bytes checked. Bundle of 340 GiB with relative + paths; `run_v41.sh BUNDLE` serves it, installing oMLX at the pinned commit on the first run. + Verified from an empty oMLX directory: 183 residents, 83 GiB, 38–40 tok/s. +- `pool384/fp4-removed` deleted (Jack confirmed), 134 GB freed. +- Rubric unchanged at 59: none of this is a rubric item's test. + +### 2026-11-12: memory pressure (D3) + +- `scripts/v41_build/d3_pressure.py`: requests back to back while a second process takes + incompressible memory in 4 GB steps to 40 GB, holds it and lets it go. Server: the bundle + through `run_v41.sh`, 183 residents, drafter on, 83 GiB. +- At tau 0.3 (the shipped setting): 36.7 tok/s before, 29.5 and 32.8 during the steps, 33.5 + while 40 GB are held, 35.9 after. 61 requests, no error. The server's footprint stays at 83 + GiB, swap goes from 5 to 25 GB, the pressure level reaches "warning", never "critical". +- The first run used tau 0.2 by mistake: `run_v41.sh` did not set the threshold and the code's + default is 0.2, six times the SSD reads (33.6 MB per token against 1.5 to 5.5). Fixed in the + launcher. That run: 39.8 before, 32.2 held, 31.8 after. +- Rubric item 6 moves from 8 to 10. Total 61. diff --git a/docs/deepseek-v4-flash/2026-11-08-v41-reap75-local-test-plan.md b/docs/deepseek-v4-flash/2026-11-08-v41-reap75-local-test-plan.md new file mode 100644 index 0000000..eece40e --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-08-v41-reap75-local-test-plan.md @@ -0,0 +1,142 @@ +# DeepSeek-V4.1-Flash-REAP-75: local test plan (M4 Max) + +Goal: run the published [jackmaxwil/DeepSeek-V4.1-Flash-REAP-75](https://huggingface.co/jackmaxwil/DeepSeek-V4.1-Flash-REAP-75) +on this Mac. First prove that the exported checkpoint computes what the Runpod eval measured. Then +produce text from it. Speed is a later question. + +REAP-75 is the testable one. REAP-50's routed experts alone are 148 GiB, more than the machine's +128 GiB, so they cannot stay in memory. + +## What is in the checkpoint + +| Part | Size | How a local run touches it | +|---|---:|---| +| Engram n-gram tables (layers 1 and 14) | 189 GiB | sparse row gathers from the shard files on disk (memory-mapped), never loaded whole | +| Routed experts, 96 per layer, FP4 | 81 GiB | streamed one layer at a time | +| Vision tower | 2 GiB | not used | +| Attention, shared experts, embeddings, head | 1 GiB | per layer | + +Disk is no longer a constraint: 948 GiB was free on 2026-11-08, after the cleanup. + +## Why the torch reference path comes first, not MLX + +V4.1 is not V4-Flash plus Engram. On top of Engram it adds: +- cross-layer KV sharing (`kv_source_layers`); +- a two-level indexer (`index_source_layers`, candidate blocks); +- per-layer compression ratios; +- DSpark heads. + +The MLX runtime in `src/keep/quality/dsv4_*` and `src/keep/convert/dsv4_resident.py` covers none +of this. + +The single-GPU driver already runs DeepSeek's own `inference/model.py` unchanged, one layer at a +time: +- `scripts/runpod/v41_reap/stream.py` already gathers Engram rows from the shards on the host. +- `scripts/runpod/v41_reap/torch_kernel.py` already ports `sparse_attn` and `hc_split_sinkhorn` to + plain torch. + +The only pieces missing on a Mac are the quantized GEMMs and the activation quantizers. Those are +the tilelang kernels, which are GPU-only. + +So the first local run reuses the code that produced the published numbers. An MLX port is worth +it only once that run passes, and only if the speed matters. + +## Steps + +1. **Environment.** + - What: a separate venv with `torch` (CPU/MPS wheel), `safetensors`, `tokenizers`, `sympy`, + `numpy` and `huggingface_hub`. The repo `.venv` has no torch, and should stay that way. + - Check: `python -c "import torch; torch.float8_e8m0fnu"` works. + +2. **Download REAP-75.** + - What: `hf download jackmaxwil/DeepSeek-V4.1-Flash-REAP-75 --local-dir ~/Models/v41-reap75`, then + `shasum -a 256 -c SHA256SUMS` in that directory. That is 273 GiB; time depends on the line. + - Check: every checksum passes. + +3. **Pull the eval pack.** + - What: copy only `pack/v41-pack.jsonl` and `pack/manifest.json` from the stopped pod + `neckuzwv11dy2i`, probably tens of MB. Start the pod with zero GPUs for this and stop it right + after. + - Why: the stats dataset has teacher logits but no input token ids, and the parity check needs + the exact eval inputs. + - After: terminate the pod, with Jack's OK. + - Check: `pack_sha256` in the manifest matches `run-config.json` (`6a354705…`). + +4. **CPU stand-ins for the quantized kernels**, in `torch_kernel.py`: + - `act_quant`: an exact port of the per-32-block FP8 activation quantization (`fp8_block_size = 32` in this release), including the + e8m0 scale rounding. It quantizes and dequantizes, so the numerics match the pod. + - `fp8_gemm`: dequantize the e4m3 weight with its 32x32 e8m0 block scales, then do a bf16/fp32 + matmul. + - `fp4_gemm`: decode the e2m1 nibbles (low nibble first, as in `keep.convert.fp4_expert`) with + per-32 e8m0 scales, then matmul. Dequantize one expert at a time, so a layer never holds more + than one expert in bf16. + - Check: extend `tests/test_runpod_v41_stream.py` with a toy model in fp8 and fp4. It must match + the same model computed in plain bf16 from the dequantized weights to within float tolerance. + +5. **Memory gate.** + - What: run one real layer on one short eval record, under `/usr/bin/time -l`. + - Expected peak: a few GiB of weights plus activations. + - Check: the projected peak for the full run is written down. If it is over 24 GB, stop and ask + Jack. + +6. **Parity: the real test.** + - What: run the exported REAP-75 layer-streamed over 2 or 3 of the shortest eval sessions. Score + it against the teacher top-K in the stats dataset (`eval/sessions/*.npz`). + - Compare with the pod's masked-router numbers for the same positions (`student-96/`). + - Expectation: KL, NLL and top-1 agree per position up to the small difference between the + tilelang kernels and the dequantize-then-matmul stand-ins. + - Pass: |ΔKL| and |ΔNLL| averaged per session are small next to the REAP-75 numbers themselves + (KL 1.35, NLL change −0.13 on all), and top-1 agreement is within a point. + - If it passes, the "not run end to end" caveat comes off both model cards. + - If it fails, the export (`prune.py`: router slicing, expert renumbering) is wrong. Then fix the + export and republish both models. + +7. **Text smoke test.** + - What: greedy decode of about 20 tokens from a short coding prompt, layer-streamed. + - Cost: each token re-reads about 84 GiB of non-Engram weights from the SSD, so expect tens of + seconds per token. That is fine for a smoke test. + - Option: keeping the experts in memory (about 85 GiB) would be faster, but it is over the 24 GB + budget and needs Jack's OK. + +## Speed path after parity: Rapid-MLX + +[raullenchai/Rapid-MLX](https://github.com/raullenchai/Rapid-MLX) (read at `7d2b20c`, 2026-11-08) already +runs DeepSeek-V4.1-Flash-1017 in MLX, from the same source revision (`1bae0e10`). What matters for +REAP-75: + +- **Engram:** tables are re-quantized to MLX affine 4-bit, group 32 (about 61 GB per layer). Rows + are read from mmap with MADV_RANDOM, through a 16,384-row LRU and a one-layer-ahead prefetch + (`rapid_mlx/models/deepseek_v41_native/engram.py`). +- **Experts:** read from `config.json` (`n_routed_experts`, optional per-layer counts), so a + pruned checkpoint fits the model code. +- **Conversion:** their loader reads only converted MLX builds. `convert.py` dequantizes our + FP8/FP4 release format and requantizes it. It asserts expert ids are contiguous from 0, which + `prune.py` already guarantees. But it buffers a whole quantized Engram table (about 61 GB) before + saving, over the 24 GB budget, so it would need a streaming write first. +- **Fit (estimated from their byte counts):** with their 2/3-bit GPTQ recipe, REAP-75 is about + 58 GB resident plus 8 GB of DSpark drafter, with Engram on the SSD. That is over the 24 GB budget + and needs Jack's OK. +- **Kernel gap at 96 experts:** their fused decode-MoE kernel requires `dim < 16 * n_experts` + (`moe_decode.py:557`). At 96 experts that fails (5120 >= 1536), so it silently falls back to the + stock path. +- **Wired memory:** `load.py:181-188` always raises the MLX wired limit, which conflicts with this + repo's wired-memory policy. Their product lane also refuses machines under 224 GiB + (`artifacts.py:57`). +- **Speed:** 34-49 tok/s greedy with DSpark speculative decoding, on an M3 Ultra 256 GB. +- **Correctness evidence:** teacher-forced KL 0.26 and top-1 85% against the BF16 source. + +Never run their scripts or package from the checkout. `scripts/calibrate_deepseek_v41_reap.py` +executes `runtime/runtime.py` from a model directory. + +## Out of scope for now + +- An MLX port of V4.1 (Engram, KV sharing, indexer, compression). +- The vision path and DSpark/MTP speculative heads. +- REAP-50, which has to be tested on a GPU pod or with the experts streamed. +- Speed work. + +## Decisions for Jack + +- Go for the 273 GiB download (step 2). +- Starting the stopped pod for a few minutes with zero GPUs to pull the pack, then terminating it + (step 3). diff --git a/docs/deepseek-v4-flash/2026-11-09-keep-findings-and-levers.md b/docs/deepseek-v4-flash/2026-11-09-keep-findings-and-levers.md new file mode 100644 index 0000000..b45df53 --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-09-keep-findings-and-levers.md @@ -0,0 +1,375 @@ +# What KEEP has learned, and which of it could go beyond the prior art + +Written 2026-11-09 for the paper track (*Offload the Residual, Not the Expert*). Jack asked what the +whole project has measured that may not be in the literature, and whether any win is a bigger lever +than it looked. Three sources: this campaign's own results +([recovery log](2026-11-08-v41-reap-recovery-log.md)), a survey of the older KEEP documents, and the +prior-art search in the [plan](2026-11-09-routed-precision-plan.md). The plan (revision 2) holds +the schedule and the status of each lever; this file holds the evidence and the reasoning. Sections +5 and 6 were added the same evening, after the first results: read them before sections 3 and 4, +several of whose bets they overturn. + +Nothing here has been checked against the literature beyond the six papers in the plan. Each lever +gets its own search before it is claimed. + +## 1. Findings from this campaign + +Differences are domain-mean KL on the Mac eval set (22 sessions, first 1,024 tokens), with a paired +block bootstrap over positions (`routed_report.py --compare`). The interval covers position noise +within the 22 sessions, not which sessions were drawn. + +| Finding | Number | 95% interval | +|---|---:|---| +| Choosing 96 experts by total routed mass on the pruned model's routing, against REAP-75's set, both at 2 bits | −0.177 KL | −0.200 to −0.148 | +| A second prune pass (192 → 144 → 96) over one pass | −0.024 | −0.034 to −0.012 | +| Static allocator at 288 units against 144 experts at 2 bits | −0.003 | −0.009 to +0.005 | + +- **Selection on the pruned model's routing is the lever; precision allocation was not.** Each of the + two changes alone (the ranking rule, or measuring on the pruned model) gave about 0.025. Together + 0.177. A second pass is statistically real but small, most of it on tool-use. +- **Quantization statistics survive pruning; routing does not.** Refitting the 2-bit codes with + importance measured on the pruned model changed nothing (0.753 against 0.747). The importance + cosine was 0.976; the selection overlap with REAP-75 was 58%. +- **Losses are not additive.** KEEP 2-bit costs +0.17 KL on REAP-50 and +0.02 on a 96-expert model. +- **Two regimes by depth.** The top 96 experts hold 72–83% of routed mass in layers 0–19 and 94–98% + in layers 20–39. Yet moving budget toward the diffuse layers lost to a uniform budget (1.046 + against 0.955 at 128 experts): concentration is not sensitivity. +- **A static 1-bit tier is nearly worthless** (0.025 KL at 192 units, nothing at 288). The 1-bit + code's error is 0.31 per projection. +- **Nested lattice codes cost about 12%.** A 1-bit E8 residual on a 1-bit E8 base scores 0.101 + (gate/up) against 0.090 for a flat 2-bit E8P fit. +- **The pruned model beats the full model on tool-use**, by 1.22 nats per token of NLL, on the pod's + own student as well as locally. Not explained. +- **Sidecars (running):** a rank-256 linear correction removes 36–76% of the layer gap where the + gap is small (layers 0–3) and 9–15% where it is large (layers 4–6). + +## 2. Findings from the older KEEP documents + +From a survey of every document under `docs/` and the module docstrings. Rows marked ✓ were +re-read at the cited line for this note; the rest are as the survey reported them. + +### Memory and speed come from different weights + +- ✓ On V4-Flash, decoding one token reads 15.27 GB: attention 67%, routed experts 11% + (`2026-08-21-roadmap.md:116-121`). Residents are dequantized to bf16 at bind. +- ✓ On V4.1, each expert kept in every layer costs 0.359 GB of RAM at 2.031 bits per weight "but adds + no per-token" reads, and fitting the Mac "forces a ~50% prune" + (`2026-11-05-v41-plan-and-phase0-status.md:15-21`). +- GLM-5.3 inverts this: residents are about 85% of decode traffic (`glm-5.3-flash/2026-10-06-plan.md`). + +So on the DeepSeek family, expert precision buys memory and almost no speed. MatMoE argues dynamic +precision from latency; that argument does not carry to this architecture, and the memory argument +is the one this paper makes. + +### The Mac thrashes at a 91 GB resident set + +- ✓ At 90.99 GB of MLX-active memory the bind alone logged 535,792 swapouts (8.4 GB of anonymous + buffers), and byte-identical work spanned 16.5× in time across rows that all passed the quiet + gate (`2026-08-21-roadmap.md:84-89, 208-232`). The document's own conclusion is that the timings + were "memory weather, not decode". +- Policy since then: leave the wired limit at its default (`docs/troubleshooting.md`). + +Anonymous memory at this footprint gets compressed and swapped, which writes to the SSD and stalls +decode. Clean file-backed pages cannot be swapped; they are dropped and read again. That is the +case for lever X4. + +### The quantizer + +- ✓ A 1-bit code gave a 0.56 block cosine, "not a usable model"; the SwiGLU amplifies projection error + about 2.8× at block level; E8P beats affine quantization below about 3.1 bits per weight and + loses above; 2 bits is "at its rate-distortion limit" (`research/wave5-pilot-report.md:12-14, 122, + 215, 305, 339`). +- ✓ Incoherence processing (the random rotation QuIP# is built on) made the importance-aware fit 15% + worse on the down projection, 9.18e-2 against 7.97e-2, because the rotation flattens a per-column + importance that spans 180× (`research/wave5-pilot-report.md:144-149`). +- The codebook step returns 0.71 of fractional error per bit per weight and group-size steps + 0.17–0.24, so group size is the rate dial. Uniform group size beats non-uniform at matched rate. + The belief that the down projection needs extra protection was not supported. +- Only 1 of 11,008 (layer, expert) pairs had no calibration routes. +- An fp32 reduction flips near-tie argmins in about 30% of scale last bits, so exact fits keep + float64 sums on the CPU. + +### Checks that lie + +- ✓ A gate/up-swapped artifact loads, gives finite logits and scores cosine 0.9993 against the correct + one; only a discrimination ratio (4.3e10) caught it (`research/wave5-materializer-report.md:12, + 101`). +- A per-call-sync benchmark harness produced false kills: the E8P kernel read 52.8 GB/s and was + marked KILL, then 122–126 GB/s chained and 201.9–205.5 GB/s tuned (`2026-11-05-phase0-kill-tests.md`). +- Two "slower than source" verdicts were bugs in the measuring path: a Python-version mismatch that + made about 50 kernels unreachable (8.36× once fixed), and a loader that re-globbed the build tree + on every call (4.53× once fixed) (`handoff.md`, `2026-08-21-compounded-speed-verdict.md`). + +### Recovery levers already tried on GLM-4.5-Air + +- ✓ Router distillation was rejected: best global top-1 change "slightly negative" + (`glm45-air/rc-pipeline.md:176`). +- ✓ A rank-4 sidecar raised report top-1 from 0.795 to 0.852 and cut KL to 0.314, but ran 1.718× the + control and failed the speed gate (`glm45-air/rc-pipeline.md:167, 895, 930`). A proxy win that + failed end to end. +- Raising non-expert precision changed nothing once a loader bug was fixed. + +### Streaming and I/O + +- A layer-outer, chunk-inner teacher loop is bit-identical to chunk-major and avoided 58.8 h of I/O. +- Prefill I/O favoured VQ: 6 `preadv` calls per layer against 1,536; the gap to the source model was + compute, not I/O. +- Engram lookups depend only on token ids; the tables are memory-mapped and gathered on the host. + The V4.1 plan already keeps Engram's 203 GB on the SSD. +- ✓ A reference point from outside: the Rapid-MLX V4.1 build reports teacher-forced KL 0.26 and top-1 + 85% against the BF16 source, at 34–49 tok/s on a 256 GB M3 Ultra + (`2026-11-08-v41-reap75-local-test-plan.md:125-126`). + +## 3. What looks like it could go beyond the prior art + +In the order I would bet on them on the morning of 2026-11-09. X numbers refer to the plan. +Outcomes so far: item 1 half right (the ranking rule, not the second pass), item 2 demoted, item 6 +promoted to the main line. See section 6. + +1. **Prune twice, at the first cut (X1).** The one clearly significant win of this campaign has only + been applied to the second halving. REAP-50's own KL (0.578) is the floor under every number + here, and it was chosen the old way. If the effect repeats at 384 → 192 it lifts the whole + frontier, and "measure routing on the pruned model, rank by total mass" is a two-line change to + REAP that anyone can use. Needs the full model: a pod run. +2. **Keep every expert and page residuals instead of pruning half (X3).** The V4.1 plan prunes 50% + because each kept expert costs 0.359 GB and the model has to fit. All 384 one-bit bases cost the + same memory as 192 experts at 2 bits. Whether 384 bases plus a residual cache beats REAP-50 at + 2 bits at equal memory is the title's claim at full scale, and no paper found so far compares + pruning against residual offload. The local sweep (192 bases against 96 or 144 at 2 bits) is its + proxy; the weak 1-bit base is the risk. +3. **The OS page cache as the precision cache (X4).** Wire the bases; map the residuals as clean + file pages; `mincore` before use, `madvise` to fetch. No cache code, and under memory pressure the + model loses precision instead of swapping. KEEP has already measured the failure this avoids + (535,792 swapouts, 16.5× timing spread at 91 GB). This is specific to how unified memory and the + macOS VM behave, and it turns "SSD churn" from writes into reads. +4. **Pruning as specialization (X2): checked the same day, and it is not that.** On the pod's + full-length outputs (27 eval sessions held locally) the student-192 NLL advantage appears in + every agent-style domain (tool-use −1.38 nats/token, agentic-python −0.33, agentic-rust −0.32, + bash −0.16), but only where the teacher is surprised. On positions where the teacher's NLL is at + most 2 the student is worse in every domain (+0.14 on tool-use). On tool-use, positions with + teacher NLL above 5 are 30% of the tokens and carry −1.29 of the −1.38. At the worst of them the + teacher's top prediction is its own tool-call markup (`｜DSML｜` after `<`) or end-of-sentence + after `>`, while the record continues with foreign markup that the source datasets embed in + message text. The full model is confidently wrong about text that is off its own format; the + pruned model is less confident and is penalized less. That is a softer distribution on + off-format data, not a better model. Two consequences: NLL on these domains is not a quality + measure, and KL against the teacher there mostly measures agreement on format expectations. + Lever dropped; the eval caveat stays. +5. **One artifact for every budget.** Because quantization statistics survive pruning and one mass + ranking yields nested expert sets (pm-144 ⊂ pm-160; 97% overlap between passes), a single set of + files can serve any memory budget: nested in expert count and in bits. MatMoE nests bits only. + This is a packaging result that follows from 1–3 rather than an experiment of its own. +6. **Finer rungs and a stronger base (X5).** The 1-bit base is the weak link by three separate + measurements. Projection-level residuals give 1.33- and 1.67-bit rungs from the existing codes, + and a 2-bit base with a 1-bit residual moves the whole ladder to where the base is usable. +7. **Sensitivity-based budgets (X6)** and **a shared base across experts (X7)** are cheaper to test + than to argue about. + +Smaller results that are probably worth a paragraph or a short note each: incoherence processing +hurting importance-aware lattice fits; cosine similarity passing a gate/up-swapped model; and the +measurement lessons (per-call sync, one-sided paging gates). + +## 4. A correction to my own earlier claim + +I first wrote that a nested cache "holds twice as many hot experts" as a separate-copy cache at +equal memory. That is true only if the separate-copy cache keeps the low-precision copy of every hot +expert. One that drops it and reads it back on demotion uses the same memory as the nested design; +what it pays is traffic, three units per promote/demote cycle against one, and a slightly more +accurate high tier in return (0.090 against 0.101). The comparison is therefore run both ways, and +at equal SSD rate as well as equal memory. The title's claim may turn out to be about traffic and +simplicity more than about memory. + +## 5. Results of the evening of 2026-11-09 + +### 5.1 Selection at the first cut (X1, pod) + +Students of the full model, 54 eval sessions at full length, KL against the full model's logits. +Domain mean over 11 domains; pooled over positions in brackets. + +| Student | Experts | KL | Top-1 | +|---|---:|---:|---:| +| Shipped REAP-50 (`reap192`) | 192 | 0.716 (0.950) | 0.800 | +| Total routed mass, full-model routing, one pass (`fm192`) | 192 | 0.532 (0.650) | 0.814 | +| 384 → 288 → 192, routing re-measured on the 288 (`pt192`) | 192 | 0.546 (0.663) | 0.809 | +| The Mac's 96 inside REAP-50 (`pm96`) | 96 | 1.163 (1.443) | 0.722 | +| 384 → 288 → 192 → 96 (`pt96`) | 96 | 1.154 (1.268) | 0.703 | + +| Difference in domain-mean KL | Point | 95%, positions | 95%, sessions | Sessions better | +|---|---:|---|---|---:| +| `fm192` − `reap192` | −0.184 | −0.202 to −0.160 | −0.277 to −0.128 | 36 of 54 | +| `pt192` − `fm192` | +0.014 | +0.006 to +0.020 | +0.002 to +0.043 | 17 of 54 | +| `pt96` − `pm96` | −0.009 | −0.036 to +0.008 | −0.074 to +0.065 | 18 of 54 | + +- The ranking rule carries the gain. REAP's per-route score ignores how often an expert is used; + total routed mass counts it. `fm192` keeps 92.5% of the full model's routed mass and overlaps + REAP-50 by 76.5%. +- By domain the gain is largest where records are long and hard (agentic-python 1.339 → 0.897, + tool-use 1.995 → 1.294) and absent or negative where calibration data is thin: english-chat goes + from 0.448 to 0.572. Calibration tokens by domain run from 43K (business-email) and 52K + (english-chat) to 451K (tool-use). Total mass inherits that imbalance; REAP's per-route mean does + not. This is lever X8. +- Re-measuring routing on the 288-expert model did not help at 192. Confound: the two-pass tables + were ranked from 2 of 24 calibration sessions (257K tokens), the one-pass table from all 24. +- How this fits the Mac result (192 → 96 inside REAP-50: rule alone −0.025, pruned routing alone + −0.025, both −0.177): a cut that keeps 92.5% of mass barely changes routing, so the full + model's statistics describe the pruned model. A cut to 96 keeps about 74% and the routing it + leaves is different. Working rule: re-measure when the cut removes a large share of mass. The + 96-expert follow-up (5.5) tests it. + +Where the gain sits (checked the same evening, from the per-position files; `fm192` − `reap192`): + +| Positions | Share | Domain-mean KL difference | 95%, sessions | +|---|---:|---:|---| +| All | 100% | −0.184 | −0.279 to −0.131 | +| Teacher NLL below 0.5 | 64% | −0.026 | −0.047 to +0.002 | +| Teacher NLL below 2 | 74% | −0.052 | −0.080 to −0.023 | +| Teacher NLL 2 or more | 26% | −0.465 | −0.637 to −0.361 | + +- 59% of the pooled gain comes from the 15% of positions where the teacher's NLL is 5 or more, and + 15% from the 74% where it is below 2. On teacher-confident positions top-1 agreement does not + move (0.947 against 0.944 pooled). On tool-use the gain on confident positions is −0.004: all of + it is on positions the teacher did not expect. +- This is not the X2 artifact. On the surprised positions `fm192` agrees with the teacher's top + token more often (0.42 → 0.51) and its NLL moves toward the teacher's (−2.22 → −0.73 relative to + it): it reproduces the full model's expectations better, including its tool-call markup. It is + fidelity on off-format and agentic text, not a general quality gain. +- English-chat is not established as worse: +0.124 over 4,319 positions in 20 short sessions, + interval −0.050 to +0.206, 8 of 20 sessions better. +- "36 of 54 sessions better" is an artifact of the session mix: 30 sessions are english-chat or + english-prose. Every tool-use (7), agentic (4), code-diffs (3), bash (2) and cyber (2) session + improves. +- The four agentic sessions share their first 3,893–3,917 tokens of 16,383 (a common system + prompt), so agentic-python and agentic-rust overlap by a quarter. With the two merged into one + domain the difference is −0.159 (−0.251 to −0.122). +- The claim that survives: at 192 experts, total-mass selection is much closer to the full model + on agentic and tool-call text and modestly closer on ordinary text (−0.05). Whether that is + better for a user needs a generation check (tool-call validity, code completion), which no + number here provides. + +### 5.2 Routed precision, pass 1 (Mac, inside REAP-50, 1-bit base + 1-bit residual) + +Numbers and intervals are in the recovery log. The decomposition at 288 units per layer: + +| Configuration | KL | Cost over 192 experts at flat 2 bits (0.747) | +|---|---:|---| +| 144 experts at flat 2 bits | 0.812 | +0.065: 48 experts dropped | +| Routed cache, 192 bases + 96 residuals, τ 0.2 | 0.832 | +0.053 nesting, +0.032 misses that fall to 1 bit | +| Fixed hot set, same codes | 0.905 | +0.053 nesting, +0.105 misses | + +- Routing the cache cuts the cost of misses by two thirds against a fixed hot set (0.105 → 0.032) + with 10% of slots still missing. The remaining miss cost is what a 1-bit fallback costs. +- The nesting penalty alone (0.053) is nearly the whole cost of dropping 48 experts (0.065). With + a 1-bit base, nested codes cannot beat pruning by much even with a perfect cache. +- SSD traffic was not binding: 0.28 GB/s at τ 0.3 and 0.83 GB/s at τ 0.2, against a 2.5 GB/s + ceiling. A separate-copy cache moving three times the bytes would still be under it. The + residual's traffic advantage only matters for small caches (K = 48 at τ 0.2 reads 2.72 GB/s). +- Per-token tiering with everything resident (MatMoE's setting) costs 0.022 KL for 54% fewer + tier-2 slots and saves no memory. Expert bytes are about 11% of decode traffic on this family, so + it buys little speed either. +- Per-token tier 0 is free at 7.7% of slots (+0.001 KL). + +### 5.3 Sidecars + +Closed. Fifteen fitted layers made KL worse by 0.199 [+0.164, +0.219] while every layer's holdout +objective improved. Details in the recovery log. + +### 5.4 What a unit buys + +From the static points inside REAP-50 (Mac eval), per unit per layer: + +| Step | KL change | Per unit | +|---|---:|---:| +| 192 experts at 1 bit → 96 experts at 2 bits (192 units either way) | 1.358 → 1.039 | bits beat breadth at equal memory | +| 96 → 144 experts at 2 bits (+96 units) | −0.227 | −0.0024 | +| 144 → 192 experts at 2 bits (+96 units) | −0.065 | −0.0007 | +| 192 experts, 2 bits → native FP4/BF16 (+384 units or more) | −0.169 | at most −0.0004 | + +The second bit is worth far more than the first expert it displaces; the last 48 experts are worth +little; and 0.17 KL sits above 2 bits. If a third bit recovers most of that 0.17 on the experts +that carry most of the routed mass, 96 residual units on the hot experts are worth more than the +96 units that buy experts 145–192. That is the case for the 2 + 1 ladder, and it is a guess until +the residual is fitted and scored. + +### 5.5 96 experts, one pass against two (pod follow-up) + +Pending when this section was written; results are appended to the recovery log. + +## 6. Priors revisited + +What I believed this morning, what the evidence says now, and what I do differently. + +| Prior | Now | Evidence | Consequence | +|---|---|---|---| +| "Prune twice" is the lever, and it will repeat at 384 → 192 | Half wrong. The ranking rule is the lever at the first cut; re-measuring matters only when routing changes a lot | 5.1 | Re-base on `fm192`; rule stated as "rank by total mass on routing that resembles the target" | +| REAP-50 is a sound floor for everything else | Wrong. It gives up 0.18 KL at the same size | 5.1 | WR in the plan; every later number is re-measured | +| A 1-bit base is weak but usable as a fallback tier | Wrong. It is the dominant cost: 1.358 alone, and each missed slot falls to it | 5.2 | Main ladder becomes 2-bit base + 1-bit residual | +| Nested codes cost "about 12%" and that is acceptable | The 12% in weight error is 0.053 KL end to end, as much as dropping a quarter of the experts | 5.2 | Nested codes must be judged end to end; a joint fit is in W7 | +| The residual wins on memory or on traffic | Memory: only against keep-both caches. Traffic: not binding at K = 96 | 5.2 | The title's claim is at risk; decided by the separate-copy sweep and the 2 + 1 ladder | +| Keeping all 384 experts at 1 bit could beat pruning half (X3) | Very unlikely | 5.2, 5.4 | Demoted; W4b no longer fits 384 experts | +| Sidecars recover a useful share of the layer gap | Wrong end to end | 5.3 | Closed; no layer-local metric accepted as evidence again | +| The eval set is 22 independent texts | 19 at 1,024 tokens; two domains share one prompt | recovery log | Eval set version 2 | +| A network volume is a disk | It is a single-stream 120 MB/s device with 3 ms seeks | recovery log | Copy to local disk first; rule in the plan | +| X1 needs 6.5–7.5 GPU hours | 1.6 h | pod logs | Estimates from measured throughput | + +Predictions registered before the runs that test them (2026-11-09, about 20:40 PST): + +1. Separate-copy cache, K = 96, swap accounting, against the routed nested cache at 288 units: the + copy wins on KL by 0.03–0.05 (about 0.78–0.80 against 0.832) and also beats 144 flat experts + (0.812). Its SSD rate stays under 2.5 GB/s at τ 0.3. +2. Oracle cache at K = 96 lands within 0.015 of the routed LRU at τ 0.2: the policy is near its + ceiling and the remaining loss is the base. +3. 144 experts all at nested 2 bits (288 units) scores about 0.86, worse than both the routed + cache (0.832) and flat 144 (0.812). +4. 96 experts in two passes through `fm192` beat one pass on full-model routing by at least 0.05 + domain-mean KL. +5. On the 2 + 1 ladder, 144 experts at 2 bits plus a routed cache of 96 third-bit residuals + (384 units) beats 192 experts at flat 2 bits by at least 0.03 KL. + +Predictions 1 and the old-ladder half of the title question will not be run: Jack cut the 1-bit +base the same evening, and passes 2b and 3 were cancelled. Predictions 2 and 3 are in pass 2, +which finishes. If 5 fails, the paper is about routing a precision cache and about selection, and +the title changes. If 5 holds, the title stands with the qualifier that the base must be a usable +model. + +### 6.1 Late revision: the lesson I drew from the 1-bit result was too narrow + +Two hours after demoting X3 ("keep every expert, page what is missing") I reversed part of that. +The data refuted a 1-bit copy as the answer to a miss. It did not refute keeping a wide pool of +experts on the SSD. Three measurements point the other way: + +- Pruning is the largest loss in the stack: REAP-50 in BF16 is 0.578 KL from the full model, and + 2-bit coding adds 0.169. The residual work addressed the 0.169. +- A cache that follows the session covers far more demand than a fixed set of the same size: 84–90% + of slots against 75% at 96 of 192, and a routed 48 beat a fixed 96. +- Skipping a low-gate slot costs nothing (+0.001 KL at 7.7% of slots), and static pruning already + answers every miss by re-routing to the best survivor. + +Put together: keep residents at 2 bits, let the resident set follow the session, re-route a miss +among residents, and admit what the router asked for. Static pruning is that system with admission +off. This is conjecture C1 in the plan (revision 3) and it is untested. It was testable tonight +with the existing 2-bit codes and I did not see it, because every design I considered kept a +low-precision copy as the fallback. + +Registered before the run (E1, pool = REAP-50's 192, flat 2 bits, same miss rule for both): + +6. LRU cache of 144 against the static top 144 (0.812): better by 0.03–0.05 KL. +7. LRU cache of 96 against the static top 96 (1.039): better by at least 0.10. +8. Re-route beats drop on the same cache by at least 0.02. +9. The LRU of 144 at 2 bits beats the 1-bit-base residual cache at the same 288 units (0.832). + +## 7. Prior art (moved here from the plan, revision 3) + +Found 2026-11-09. Only abstracts were read for all but MatMoE (introduction and method). + +| Paper | What it does | Overlap with us | +|---|---|---| +| MatMoE (arXiv 2611.01704, ICCAD 2026) | Nested integer codes per expert; bit-width per expert at run time from its routing score | Per-token precision with nested codes. All weights resident; no offload, no pruning. | +| DynaExq (arXiv 2511.15015) | Hot experts resident at high precision, low-precision fallback, asynchronous promotion | A no-stall cache. Separate copies. GPU memory, not SSD. | +| DyMoE (arXiv 2603.19172) | Run-time mixed precision by importance, depth-adaptive, look-ahead prefetch | Depth regimes, prefetch. | +| HOBBIT (arXiv 2411.01433) | On a miss, run a low-precision copy instead of waiting | A low tier as the miss path. | +| MC-MoE (arXiv 2410.06270) | Static mixed precision plus online per-token expert pruning | Per-token dropping with static precision. | +| PMQ, GEMQ (ICML 2026), EdgeMoE | Static per-expert bit allocation | The step 2c allocator. | + +To read in full and to search for (W0): expert caches that re-route or skip on a miss; any +comparison of an expert cache against pruning at equal memory; REAP's comparison of its criterion +against frequency-based ones; calibration on the pruned model's routing. diff --git a/docs/deepseek-v4-flash/2026-11-09-routed-precision-plan.md b/docs/deepseek-v4-flash/2026-11-09-routed-precision-plan.md new file mode 100644 index 0000000..8686c52 --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-09-routed-precision-plan.md @@ -0,0 +1,175 @@ +# Routed residency: plan to the paper (revision 3) + +Revised late on 2026-11-09 after a critique of revision 2 (in git at `691f97a`). Score on the fixed +rubric in the [recovery log](2026-11-08-v41-reap-recovery-log.md): **31 / 100**; the rubric and its +weights do not change. Evidence for every number here is in the +[findings note](2026-11-09-keep-findings-and-levers.md); prior art is its section 7. + +## 1. Thesis + +Pruning is a cache that never updates. Keep every resident expert at a usable precision (2 bits), +let the set of resident experts follow the session, and send a token whose expert is absent to the +best resident one instead of to a worse copy. Precision above 2 bits is a second, smaller question. + +C1 held in its first test (E1, 2026-11-10, inside the pool of 192): an expert cache with re-route beats the static set by 0.056 KL at 144 residents and 0.231 at 96, and a cache of 96 equals a static 144. C2 (the pool of 384) is the open half. Measured before that: a routed cache of +residuals beats a fixed one by 0.068 KL; a 1-bit fallback is the dominant cost; dropping low-gate +slots is free; the second bit is worth more than breadth. + +## 2. Target architecture + +```mermaid +flowchart LR + subgraph OFF[Offline, once] + CAL[Calibration: routing statistics] --> SEL[Rank experts by total routed mass] + SEL --> FIT[2-bit lattice code per expert] + FIT --> PACK[One file extent per expert] + end + subgraph SSD[SSD] + POOL[Expert pool at 2 bits: 192 or 384 per layer] + HIGH[High tier: third bit or native FP4] + ENG[Engram tables, 190 GB] + end + subgraph RAM[Unified memory] + CORE[Always resident: attention, shared experts, router rows for the whole pool] + RES[Resident experts: N per layer at 2 bits] + HOT[Hot tier: H slots, optional] + end + subgraph TOK[Per token, per layer] + R1[Router scores the whole pool] --> R2[Top 6 among resident experts] + R2 --> R3[Skip slots below the gate floor] + R3 --> R4[Run experts and sum] + R1 --> DEM[Demand: wanted but not resident] + end + PACK --> POOL + CORE --> R1 + RES --> R2 + DEM -->|gate at least tau, asynchronous| Q[Admission queue] + POOL --> Q + Q -->|pread, usable from the next token| RES + HIGH -.-> HOT +``` + +- Decode never waits on the SSD: a miss is re-routed now and admitted for later tokens. +- Static pruning is this system with admissions switched off. Per-token dropping is the gate floor. + Quantization is the precision of the resident copy. One budget, in resident GB, covers all three. +- Open in this picture: the pool size (192 or all 384), whether a hot tier pays, and the real + resident size of everything that is not a routed expert. + +## 3. Rules + +1. **Two numbers per comparison.** KL on teacher-confident positions (teacher NLL below 2) is + general quality; agreement on the rest is fidelity on agentic and tool-call text. A change that + moves only the second is reported as that. +2. **Spend memory where it pays.** The second bit is worth about three times more per unit than + the last 48 experts, and a 1-bit base loses at every budget. One sweep asks where the next unit + goes: a resident expert, a third bit, or a native copy. +3. **Use the slack.** The routed cache read 0.28–0.83 GB/s against a 2.5 GB/s ceiling, so whole + 2-bit experts (2 units) and native copies (4 units) are affordable to move. +4. **Fix the fallback before the policy.** No work on admission rules unless the oracle gap exceeds + 0.015 KL. +5. **Fail fast.** Score a partial result end to end before finishing it; write the prediction + before the run; estimate from measured throughput. + +## 4. Critique of the approach so far + +| Hole | Why it matters | What closes it | +|---|---|---| +| We never tested the simplest design: a 2-bit expert cache with re-routing on a miss. We were anchored on residuals. | Pruning costs 0.578 KL and 2-bit coding 0.169; the residual work chased the smaller loss with a fallback that could not work. | E1, tonight, no new fits. | +| Every result is teacher-forced KL. No text has been generated. | A 0.03 KL threshold is not tied to anything a user sees; the selection gain sits mostly on positions the teacher did not expect. | Rule 1 now; a generation check (E6) before any public claim. | +| The offload result is a simulation, and V4.1 has no decoder on this Mac. | "On a laptop" rests on a trace replay, an assumed 30 tok/s, and an SSD shared with 190 GB of Engram lookups. | W6 replay states exactly that; real decode only on the second model (E8). A stated limitation. | +| The memory axis counts routed experts only. | Attention, the KV cache and the Engram page cache are resident too, and this Mac thrashed at 91 GB. The true budget may sit below 384 units. | Measure total resident size once; plot the frontier in total GB. | +| τ and K were picked on the eval set. | Selection bias in our own favour. | Tune on the Mac set; confirm once on the pod's untouched 54 sessions (shard 1/2). | +| The eval set is thin: domains with one session, shared prompts, sessions of 11 tokens. | Domain means swing on one text. | Eval set version 2; session-level intervals; all 108 sessions for the final table. | +| The static baseline was REAP-50. | A better static set exists (total mass), so a win over REAP-50 flatters the method. | Static baselines use the total-mass set after the re-base. | +| The third-bit bet has a thin margin. | Everything above 2 bits is worth at most 0.169 KL on this pool. | C3 is tested after C1 and C2, not before. | +| The selection claim may contradict REAP's own comparisons against frequency-based criteria. | A reviewer will know that paper. | W0 reads it; C7 adds a generation test. To verify: this is from memory of the paper. | +| Three avoidable pod failures in one evening, the last a volume filled by mixing GB and GiB. | About $8 and two hours. | Section 9. | + +## 5. Conjectures the paper proves or refutes + +Each is stated so that a single run can refute it. KL is domain-mean; "confident" means positions +with teacher NLL below 2; intervals are session-level bootstrap. + +| # | Conjecture | Test | Refuted if | +|---|---|---|---| +| C1 | **Residency should follow the session.** With the same codes, memory and miss rule, an LRU expert cache with demand admission beats the best static expert set. | E1: pool of 192 at flat 2 bits, N = 144 and 96, static top-N in the same run. | Gain below 0.02 confident KL at both N, or interval includes zero, or reads above 2.5 GB/s at 30 tok/s. | +| C2 | **A wider pool beats a better static set.** At 384 units, 192 residents drawn from all 384 experts beat the best static 192. | E3–E4: 2-bit codes for all 384, cache of 192. | Gain below 0.03 confident KL. | +| C3 | **Above 2 bits, precision should be routed too.** At equal units, a hot tier (third bit or native FP4) on the hottest residents beats spending those units on more 2-bit residents. | E5 on the winner of C1/C2. | Gain below 0.02, or negative. My prior is close to even. | +| C4 | **Re-route beats drop beats a 1-bit copy** as the answer to a miss, at equal memory. | E1: same cache with drop and with re-route; the 1-bit-base cache from pass 1 (0.832 at 288 units). | Any other order. | +| C5 | **The policy is near its ceiling.** A per-session oracle set beats the LRU by at most 0.015 KL. | Pass 2 (residual cache), E1 (expert cache). | Gap above 0.015: then prediction and prefetch are worth building. | +| C6 | **Traffic does not bind.** The best configuration at each budget reads under 1 GB/s at 30 tok/s, and 99% of admissions land within one token interval on the internal SSD. | Simulation in E1/E4; W6 replay. | Either number missed: then the unit moved matters and nested residuals return. | +| C7 | **Rank by sum, not mean.** Total routed mass beats REAP's per-route score (a) on confident KL and (b) in generated tool calls and code. | (a) measured at 192: −0.052 [−0.080, −0.023]. (b) E6. | (b) shows no gain or a loss. | +| C8 | **Re-measure routing when the cut is deep.** 96 experts picked through the 192 beat 96 picked in one pass by at least 0.05. | E2, on the pod now. | Smaller, or reversed. | +| C9 | **The gate floor is free.** Skipping slots with gate below 0.1 costs under 0.005 KL in any configuration. | One point so far (+0.001); added to E4. | Any configuration above 0.005. | +| C10 | **C1 holds on a second MoE.** | E8, Qwen3-30B-A3B, with real generation. | C1's test fails there. | + +C1, C2 and C6 carry the title. C7–C8 are a second, independent result. + +## 6. Experiments, in order + +| # | What | Where, cost | Gate | +|---|---|---|---| +| E0 | Pass 2 on the old ladder (oracle gap for C5). | Mac, finishing about 23:00. | — | +| E1 | **Done 2026-11-10** (trimmed to four variants). Expert cache at 2 bits inside REAP-50 with re-route on a miss: 0.756 against 0.812 static at 144; 0.809 against 1.039 at 96; oracle 0.789 at 96; the static control reproduced 0.812. Details in the recovery log. | Mac, 1.9 h. | C1 holds; C5 gap 0.020, all warm-up; C6 holds at 144, misses at 96; C4 not run. | +| E2 | 96 experts, one pass against two. | Pod, running. | C8. | +| WR | Export and download the total-mass 192: the honest static baseline and a better pool of 192. | CPU pod, overnight. | — | +| E3 | If C1 holds: 2-bit codes for all 384 experts and the full router rows, as a codes-only pool (about 150 GB). Size the fit first; it has only run on the Mac GPU. | Pod. | Needs C1. | +| E4 | Cache of 192 and 144 drawn from 384. | Mac. | C2, C6. | +| E5 | Hot tier on the winner: third bit (needs one residual fit) and native FP4 (no fit). | Mac. | C3. | +| E6 | Generation check: tool-call validity and code completion, total-mass set against REAP-50, then the best cache. Not designed yet; needs the reference decoder on a pod. | Pod. | C7b; gates any public release. | +| W6 | Trace replay on the internal SSD, `pread` pool against `mmap` faults, with and without memory pressure. | Mac. | C6. | +| E7 | One confirmation on the pod's untouched 54 sessions with the routed forward on GPU. | Pod, 2–3 h. | Item 4, 25 points. | +| E8 | Second model. | Mac; download needs Jack's go. | C10. | + +If C1 fails: the static path stands (re-base, C3 on a static set), the title falls back, and the +paper is the unit accounting plus the selection result. + +## 7. Headline + +The current title, *Offload the Residual, Not the Expert*, is not defensible on tonight's data: the +1-bit base it needs is cut, SSD traffic did not bind where it was measured, and the design in +section 2 offloads experts. It survives only if C6 is refuted. + +Proposed working title, conditional on C1: + +> **Pruning Is a Frozen Cache: Unifying Pruning, Quantization and Offloading in MoE Inference** + +It keeps the half of the original headline that has held up, and it states the claim a reader can +check: a static expert set is the special case of a routed one, and it is dominated at equal +memory. Fallback if C1 fails: *Bits Before Breadth: What a Gigabyte Buys in MoE Compression*, an +empirical paper on the unit accounting and the selection rule, which the data already supports. + +Novelty risk: expert caches with a miss fallback exist (HOBBIT, DynaExq, and gating that skips +experts). What they do not appear to do is compare against pruning at equal memory, or treat the +two as one policy. W0 must confirm that before the title is used. + +## 8. Rubric mapping + +| Item | Now | Moves when | +|---|---:|---| +| 1 Nested codes validated | 10 | done (and the 1-bit ladder is retired) | +| 2 Per-token tiers built | 10 | done | +| 3 Cache simulation and cost model | 10 | done | +| 4 Core result (25) | 0 | 10: the routed cache beats the best static set by 0.03 KL (0.02 confident) and the 1-bit-fallback cache, one budget, reads under 2.5 GB/s. 15: two budgets, eval set 2, total-mass baseline. 25: E7. | +| 5 Unification and frontier figure (10) | 0 | one policy sets residency, gate floor and tier; figure in total resident GB | +| 6 Real Apple Silicon streaming (10) | 0 | W6 | +| 7 Learned policy or trained codes (10) | 0 | per-expert thresholds, or a hot tier that passes C3 | +| 8 Second MoE and downstream (10) | 0 | E8, E6 | +| 9 Paper draft (5) | 1 | sections as conjectures resolve | + +## 9. Operating rules and resources + +- Pods: copy the checkpoint to local disk before GPU work; size volumes in GiB with 10% headroom + and check `df` before a large write (a full network volume writes empty files without an error); + never run a second large job beside a GPU job; test launchers before the GPU clock matters. +- Evals: session-level intervals; state domain mean or pooled; no layer-local metric as evidence; + settle the expert pool before fitting codes to it. +- Spend tonight: about $20 (H100 4.5 h, CPU pod, volume). Volume `yqazs7aqe6` is 720 GB, $50 a + month, deleted after E7. +- Mac: 80 GB cap; 469 GB free and rising after the cleanup. + +## 10. Decisions for Jack + +1. The title in section 7, after E1. +2. A replacement REAP-50 on Hugging Face: not before E6. +3. Qwen3-30B-A3B download for E8. diff --git a/docs/deepseek-v4-flash/2026-11-10-v41-wow-build-plan.md b/docs/deepseek-v4-flash/2026-11-10-v41-wow-build-plan.md new file mode 100644 index 0000000..e7ad151 --- /dev/null +++ b/docs/deepseek-v4-flash/2026-11-10-v41-wow-build-plan.md @@ -0,0 +1,941 @@ +# V4.1 on a 128 GB Mac, generating: build plan + +Written 2026-11-09/10 from three read-only surveys (repo inventory, sizing, prior art), three rival +plans drafted from different premises, and a red-team of each against the code. Everything marked +"estimate" is a planner's estimate and has not been measured. Days are working days. + +**The bar.** A stranger with a 128 GB Mac runs one command and watches the full +DeepSeek-V4.1-Flash, with all 384 experts per layer reachable, do an agentic coding task at a speed +they would use. Nothing generates V4.1 in this repo today. + +**The claim to build toward.** The full model already runs on 128 GB (oMLX), at 4.2–5.6 tok/s, +because every cache miss waits for the SSD. The fully resident build (Rapid-MLX) does 28 tok/s +without the drafter and needs 256 GB. Our claim is the trade in between: a miss that costs no +time, at a measured and bounded loss of exactness. It is a speed and evidence claim, not a new +idea; section 7 says which ideas are already shipped. + +## 1. The three questions + +### 1.1 How much of a V4.1 decoder already exists here? + +Source: repo inventory, 2026-11-09, read file by file against the reference +(`~/Models/v41-reap50/inference/model.py`). + +- No V4.1 decoder. A complete MLX decoder for the older V4-Flash exists + (`src/ramp/models/deepseek_v4_flash_adapter.py`, `src/keep/quality/dsv4_mtp_runtime.py`): + prefill, KV cache, greedy decode, drafter generation with rollback. Correct on toys; never shown + fast (its one decode figure, 0.24 tok/s, was paging on an M5 Max). +- Reusable for V4.1: router and SwiGLU (`adapter:913-1078`), RoPE, sinkhorn math and its fused + Metal kernel (`src/ramp/kernels/hc_sinkhorn.py`), expert compute from 2-bit codes + (`src/ramp/kernels/e8p_m1.py`, 202–206 GB/s on this Mac at V4-Flash shapes), FP8 residents as + mxfp8 (`src/keep/convert/fp8_block.py`), native FP4 experts (`mx.gather_qmm`), `preadv` I/O + (`src/keep/quality/dsv4_teacher_runner.py:640-760`), and a numeric oracle: exact CPU ports of all + six reference kernels (`scripts/runpod/v41_reap/torch_kernel.py`). +- Missing: Engram in MLX entirely; V4.1 attention (compressed KV shared across layers, shared index + keys, two-level candidate blocks); hyper-connection wiring (mix weights used one sublayer late); + a cache layout with rollback; drafter changes; a loader; a sampler; the chat format; any server. +- The reference cannot run on a Mac as shipped (its kernels are CUDA-only) and decodes one token + per call after the prompt, so chunked prefill and drafter verification have no reference to copy. +- Outside this repo the whole decoder exists twice, under permissive licences (section 7). + +### 1.2 Where does decode time go? + +Source: sizing survey, 2026-11-09, from the safetensors headers and the reference; no V4.1 decode +has been timed on any Mac in this repo. + +| Read per decoded token, full model | GB | +|---|---:| +| Attention | 5.15 | +| Shared experts | 1.42 | +| Output head | 1.32 | +| Router, hyper-connections, norms, Engram projection | 0.63 | +| **Non-expert subtotal** | **8.52** | +| 240 routed experts at native FP4 / at 2.03 bits | 4.51 / 2.16 | +| **Total** | **13.0 / 10.7** | + +- Ceiling at this Mac's measured read rate (409 GB/s on 2026-10-06, 474–477 on 2026-11-05): + 31–45 tok/s. Routed experts are 35% of bytes at native precision, not the 11% seen on V4-Flash + in August; that lead does not transfer. +- With the drafter, one verify pass serves several tokens, so non-expert reads are shared and + experts become the larger share. Break-even is about 2.2 accepted tokens per cycle (estimate); + Rapid-MLX reports about 3.1. +- oMLX's 4–6 tok/s is consistent with waiting on about 74 expert loads per token (estimate from + its hit rates). Removing the waits alone caps near 20 tok/s; the rest depends on staying on its + fused kernels. +- Unknown until measured here: the time split by component, sinkhorn cost (80 calls per token), + the indexer scan at long context. + +### 1.3 What is resident besides routed experts? + +Source: sizing survey, 2026-11-09. + +| Item | GB | +|---|---:| +| All non-expert weights as stored (FP8) | 9.85 | +| The same expanded to bf16 | 16.7 | +| Drafter as stored | 7.9 | +| Context state at 64K tokens (3,200 bytes per token) | 0.21 | +| One routed expert: native FP4 / our 2.03-bit code / oMLX 3-bit | 0.0188 / 0.0090 / about 0.0176 | +| 192 experts per layer at 2.03 bits | 69 | +| 384 experts per layer at 2.03 bits | 138 | + +- Memory goes almost entirely to routed experts. The Engram page cache is the open term: up to + 0.77 MB per new token, so as much as 50 GB at 64K if n-grams never repeat (upper bound). +- The "91 GB thrash" was measured on an M5 Max 128 GB on 2026-08-21, not on this M4 Max. Nobody + has measured the ceiling here. +- At the 80 GB process cap about 160 experts per layer fit; about 250 need roughly 108 GB with + residents and the drafter, which one planner puts at Metal's limit on this machine. + The top 192 / 250 / 288 experts by routed mass cover 92.5% / 95.2% / 97.5% of routed mass. + +## 2. Chosen architecture + +**Build on oMLX as a pinned fork; plug in our cache policy, selection tables and (if it passes its +kill test) our 2-bit kernels.** + +- oMLX (`~/Development/_refs/omlx`, commit `4213c328`, Apache-2.0 with MIT parts) already has the + V4.1 decoder, the Engram reader, the drafter, an expert store with `pread` into slots, a server + with tool-call parsing, and it reads the FP4/FP8 checkpoint already on this Mac. +- The change is a **compact router**: each layer keeps C experts in slot stacks and the router's + weight and bias are sliced to those C rows. oMLX's fused kernels then pick the top 6 among + residents and renormalise. That is the re-route, with no wait, no read and no extra sync. Demand + for absent experts is read once per token from the full router and drives background admission. +- New code (estimate, about 300 lines plus four edits): `routed_cache.py` beside + `omlx/patches/deepseek_v41/moe_offload.py`; later `e8p_moe.py` calling + `src/ramp/kernels/e8p_m1.py`. +- Stock oMLX in its exact, waiting mode is both the day-1 measurement tool and the baseline. + +Why the rivals lost: + +- **Extend our own V4-Flash runtime.** Its own planner called it "a rewrite, and a fourth + implementation": the adapter differs from V4.1 in attention, hyper-connection wiring, caches + and drafter, and has no Engram. Twelve days to a first MLX token and about 47 to a demo, spent + mostly on what two projects already ship, on a runtime with no clean speed measurement. Its + E8P kernels and oracle survive as plug-ins. +- **Port the reference directly.** A torch decoder on CPU gives a first token in 3.5 days at + 20–80 s per token, and none of it carries into a fast runtime. Its one unique product, a + generation oracle, is covered by stock oMLX from day 1 at about 5 tok/s. Kept from it: parity + gates must be statistical (rounding alone moves KL by 0.1–0.3 per position between runs). +- **Rapid-MLX as the base.** The cleaner kernel host (9.4k lines against about 60k) and the faster + resident path, but it has no expert store, refuses to start under 224 GiB, and needs a download + of about 320 GB. It is the fallback if oMLX's fused path proves slow. + +Red-team of the chosen plan: + +- **Speed is unverified until day 4.** oMLX publishes no fully resident V4.1 number. The + planner's estimate is 15–20 tok/s without the drafter. +- **Their codebase moves** (40 commits to the V4.1 directory since October; exact `mlx` pin; + compiled kernels). Pin the commit and do not rebase before the demo. +- **Little of the system is ours.** What is: the zero-latency miss on a Mac, the policies in + section 5, a denser code, and the measurements. +- **oMLX chose exactness on purpose**, and cites a result that pinning the hot half and zeroing the + rest costs about 91% accuracy on a maths benchmark (source could not be verified; zeroing is not + re-routing). Teacher-forced KL will not settle this. Generation has to. + +## 3. Kill tests before any building (days 0–1) + +| # | Test | Kills | +|---|---|---| +| K1 | Stock oMLX on this Mac, local pruned checkpoint: tok/s, hit rate, resident memory, and its no-miss floor. **Passed 2026-11-10:** the loader takes the 192-expert checkpoint; with oMLX's own offload at 48 of 192 experts it decodes at 3.3–4.5 tok/s (hit rate 78%, about 0.9 GB read per token) and is not on the fused path; with a fixed 48 or 64 experts and a sliced router it decodes at 30–32 tok/s in the same memory. | The base, if the no-miss floor is under about 10 tok/s: switch to Rapid-MLX's decode kernels. | +| K2 | `bench_e8p_m1` at 5120 × 2304 against oMLX's 2-bit affine kernel. **Passed 2026-11-10:** per routed expert the lattice kernel takes 0.037 ms, oMLX's affine 2-bit 0.047 and its FP4 0.042; the lattice expert is 9.0 MB against 11.1. | Our lattice kernel in the runtime: ship affine 2-bit instead. | +| K3 | Static 250 and 288 of 384 against the full model, teacher-forced, on the pod. **Done 2026-11-10:** KL 0.323 at 250 and 0.235 at 288 (0.532 at 192, 1.123 at 96), native precision, no cache. | Nothing; it sets how much the cache has to recover. At the 80 GB cap (about 160 experts) a fixed set costs roughly 0.6–0.7 by interpolation, so the cap matters more than the code density. | +| K4 | E1: expert cache with re-route against static sets, teacher-forced. **Passed 2026-11-10:** the cache beats the static set by 0.056 KL at 144 residents and 0.231 at 96 (pool of 192, 2 bits); a cache of 96 equals a static 144; the oracle gap (0.020) is all in the first 256 tokens. | The quality case for the cache, if it is not better than a static set. | +| K5 | 20 lines: an MLX array over a 1 GB memory-mapped file through the E8P kernel; check footprint and `mincore`. **Failed 2026-11-11** (`scripts/v41_build/k5_pagecache.py`, mlx 0.32.4): `mx.asarray(np.memmap, copy=False)` does alias the mapping (same address, no page touched), but a GPU kernel that reads 1 MiB of a 2 GiB aliased buffer brings all 131,072 pages into memory (0.79 s). One buffer over all experts would page in the whole pool. | Idea 3 (page cache), if it copies. Dead in the one-buffer form; one buffer per expert would need the routed experts known on the host before dispatch, which costs a sync per layer. | +| K6 | Offline, from stored routes: how well token ids predict experts; how well a prompt's demand predicts the rest of the session; recall of the drafter's 256-d head. | Ideas 2, 6 and 7 respectively. | + +## 4. Critical path + +**To the first generated token on this Mac: day 1.** Stock oMLX on `~/Models/v41-reap50` (already +on disk, no download). That token comes from the pruned 192-expert model in exact mode. It is a +measurement, not the result. + +**To the first no-wait token: day 4** (estimate; I would plan 4–6). + +| Day | Work | Done when | +|---|---|---| +| 1 | K1, K2, K5; parity of oMLX against `torch_kernel.py` on the local model | Top-1 match on 200 positions and KL under 0.01 against the oracle | +| 2–4 | Compact router with a static resident set | Bit-identical to stock with C = all; all 40 layers stay on the fused path; at least 12 tok/s without the drafter | + +**To the recorded demo: about 19 days with no slack; plan 19–26.** + +| Days | Work | +|---|---| +| 5–7 | Background admission, warm tables from our selection, drafter on | +| in parallel | Pod or download: the other 192 experts and the full router rows (decision 2) | +| 8–11 | Lattice-code expert path and the pool format (only if K2 passed) | +| 12–13 | The 384 pool end to end | +| 14–16 | Generation benchmarks, baseline (stock oMLX at equal memory), two ablations | +| 17–19 | Wrapper script, model card, recording | + +## 5. The seven ideas, placed + +None is on the path to the first token. Ideas 4 and 5 share one mechanism. Two planners placed +them independently; where they disagreed the more sceptical reading is kept. + +| # | Idea | Attaches at (oMLX fork) | Cheapest kill | When | +|---|---|---|---|---| +| 1 | Drafter with a no-wait cache | `mtp.py` and the verify path in `decode_fused.py`; admission only between draft-verify cycles so one cycle sees one resident set. Needs a multi-row lattice kernel. | Day 7: tokens per cycle over cycle time under 1.5× single-token decode. | Demo | +| 2 | Drafted tokens as lookahead | Admission queue in `routed_cache.py`; needs a token-id → expert prior. | K6: recall of the next four tokens' experts against plain decayed counts. | After | +| 3 | OS page cache as the cache | A native module with a no-copy Metal buffer over a mapping. | K5, as a standalone Metal test: dead if Metal wires the pages or a fault stalls the GPU. | After | +| 4 | Exactness inside tool calls | `output_parser.py` already tracks the tool-call state; the exact path is the existing load-and-wait. Drafter depth 1 inside a call. | 50 tool-call prompts at C = 160: cut if re-routed validity equals exact. | Demo, conditional (1–2 d) | +| 5 | Exact-to-fast dial | The routing weight that missed, scored in `routed_cache.py`; same exact path as 4. | Offline now: the harness already takes a threshold, so plot KL against the share of tokens that wait. A flat curve kills the runtime dial. | Curve before the demo; runtime after | +| 6 | Prompt as profile | Prefill branch of the MoE call: accumulate the full router's weight per expert, bulk-load before decode. | K6: mass covered over the next 64 generated tokens, against the static set plus counts. | Demo if it passes (1 d) | +| 7 | Shortlist head | `head.py` plus the drafter's 256-d head. | Day-4 profile: cut if the head is under 5% of token time. | After | + +What the code says about them: + +- **Idea 1 is the strongest, with a measurement behind it.** oMLX has allowed the drafter with + offload since 2026-09-23, and it gains only 29% (6.6 → 8.5 tok/s at 81–89% acceptance), because a + four-token verify asks for four times the experts and waits for them (333 ms against 151 ms). + That wait is exactly what a no-wait cache removes. +- **Idea 3 is harder than it looked.** Neither runtime computes on mapped pages; both copy into + allocated buffers, and Rapid-MLX's loader warns that a cold page-in inside a GPU command trips + Metal's watchdog. A safe partial version: use `mincore` only to rank absent experts that are + already in the page cache, so they are admitted first. +- **Ideas 2 and 7 are weak as stated.** A verify pass already computes the drafted tokens, so + Engram lookahead saves under 1% of wall time, and oMLX's own router prefetch ran 2–15% slower. + The drafter's 256-d head is a bias on the previous token and never sees the hidden state, so a + shortlist would need a newly trained projection. + +## 6. The inventory, placed + +Exists = in this repo. oMLX = reuse from the fork (Apache-2.0/MIT). Build = new work. Cut = not for +the demo, with the reason. + +| # | Item | Placement | +|---|---|---| +| 1 | Block port | oMLX `language.py` | +| 2 | Engram | oMLX `engram.py`, `storage.py` | +| 3 | FP8/FP4 loaders | oMLX `loading.py`, `convert.py` | +| 4 | Incremental decode, KV, prefill | oMLX `decode_fused.py`, `cache.py` | +| 5 | Tokenizer, chat, tool calls, sampling | oMLX `encoding.py`, `tool_parser.py` | +| 6 | Drafter | oMLX `dspark.py`, `mtp.py`; build the verify check on the compact router | +| 7 | Parity harness | Build on `torch_kernel.py`, `local_parity.py`; statistical gates | +| 8 | Speed: experts from 2-bit codes, attention kernels | Build `e8p_moe.py` on `e8p_m1.py` if K2 passes; attention kernels from oMLX | +| 9 | On-disk pool format | Build: stacked safetensors readable by oMLX's row reads | +| 10 | Resident-set manager | Build on oMLX `_ExpertSlots` | +| 11 | Admission path | Build: background, applied between tokens | +| 12 | Router changes | Build: compact router. Gate floor cut until measured in the runtime | +| 13 | Prefill behaviour | Comes with the compact router | +| 14 | Memory-pressure response | oMLX's memory guard; shrink C | +| 15 | I/O sharing with Engram | oMLX: the expert pool is already separate | +| 16 | Warm start, persistence | Build: save and load slot ids | +| 17 | Hot tier above 2 bits | Cut: unproven, and one tier is enough for the demo | +| 18 | 2-bit codes for all 384 | Build: fit on the Mac after fetching the other 192 experts, or on a pod (decision 2) | +| 19 | Importance statistics, router rows | Exists for all 384 (statistics); router rows from the pod volume | +| 20 | Non-expert weight format | oMLX: source FP8 as stored (9.85 GB) | +| 21 | Selection and warm tables | Exists (`rerank.py`, total mass); regenerate for the pool | +| 22 | Packaging, download | Build after the demo: codes pool on Hugging Face (public: decision 4) | +| 23 | Licence review | Half a day: MIT model, Apache-2.0 fork notice | +| 24 | Eval set version 2 | Build (paper track, not the demo) | +| 25 | Generation benchmarks | Build: tool-call validity, code tasks that execute, one agentic task | +| 26 | Baselines | Stock oMLX at equal memory, run here; Rapid-MLX numbers quoted (216 GB does not fit) | +| 27 | Hardware measurements | Build: decode ms per token by component, memory, SSD rate, swap count | +| 28 | Smaller-RAM emulation | Cut until after the demo; the slot count gives it later | +| 29 | Ablations | Two: static set against cache, wait against re-route | +| 30 | Second model | Cut from the demo; stays in the paper plan | +| 31 | Reproducibility | Pinned fork commit, checksums, one script | +| 32 | CLI and server | oMLX `omlx serve` | +| 33 | One-command install | Build: wrapper script | +| 34 | Support matrix, CI | Cut to "128 GB M-series" | +| 35 | Recorded demo | Build | +| 36 | Cards, paper | Model card for the demo; paper separately | +| 37 | Gate: cache beats static | K4, passed 2026-11-10 | +| 38 | Gate: cache over 384 recovers the pruning loss | Teacher-forced in our harness once item 18 exists; then in generation | +| 39 | Gate: speed bound | K1, then the day-4 measurement | +| 40 | Gate: compute from 2-bit codes fast enough | K2 | + +## 7. Prior art + +| Project or paper | What it does | Reuse | Licence | +|---|---|---|---| +| DeepSeek V4.1-Flash | Model, torch reference, Engram, drafter | Reference and oracle | MIT, weights included | +| oMLX | V4.1 on MLX; expert store on SSD, waits on a miss; drafter with offload; 4.2–5.6 tok/s on 128 GB | The base | Apache-2.0, MIT parts | +| Rapid-MLX | V4.1 fully resident: 216 GB, KL 0.26, 28 tok/s plain, 33–49 with the drafter; Engram prefetch | Fallback kernels | Apache-2.0 | +| ds4 | C/Metal engine with a slab expert cache; waits on a miss | Design reading | MIT | +| llama.cpp | V4.1 merged 2026-10-25; IQ2_XXS is an E8-lattice 2-bit format | — | MIT | +| Cache-prior (2412.00099), SMoE (2508.18983), MoE-CORE (2610.01950) | Re-route or substitute a miss using router scores | The rule | papers | +| SwapMoE (2308.15030) | Dynamic resident set against static pruning at equal memory, small models | The comparison | paper | +| SliceMoE (2512.12990), MatMoE (2611.01704) | Nested bit slices; nested codes chosen by routing score | — | papers | +| HOBBIT, DynaExq, DyMoE, Fate, AdapMoE, MoE-Infinity | Low-precision fallback, asynchronous promotion, prefetch | Baselines | papers | +| REAP (2510.13999) | Per-route pruning score; deliberately avoids frequency weighting | Already implemented here | Apache-2.0 | + +Which of our ideas someone has already shipped: + +- **Shipped:** the full model on 128 GB with experts on SSD (oMLX, ds4); drafter with offload + (oMLX); 2-bit V4.1 builds (three of them); Engram on SSD with prefetch (Rapid-MLX). +- **Published, not shipped on a Mac:** re-routing a miss; nested precision caches; a dynamic set + against static pruning. +- **Not found:** a miss that costs no time in a Mac runtime at this scale; exactness switched by + grammar state; the exact-to-fast curve; a residency check on the page cache with re-route; + drafted tokens as lookahead; quality against pruning at V4 scale; selection by total routed mass. + +## 8. Decisions for Jack + +1. **Memory cap. Decided 2026-11-09: stay at 80 GB for now** (the Mac is under load from other + sessions; it can be raised later, and pods are the route for anything that needs more). The + build therefore targets C = 160 experts per layer on the Mac. Anything that needs about 250 + resident (roughly 108 GB) is measured teacher-forced on a pod until the cap changes. +2. **Source of the other 192 experts.** (a) Fetch their FP4 weights from the pod volume (about + 144 GiB) and fit our codes on the Mac overnight; (b) fit on a pod (spend; the fit has only run + on the Mac GPU); (c) download oMLX's 3-bit build (about 331 GiB), which gives a 384-expert + token soonest. Free disk is about 430 GiB, so not all of them. +3. **Fork.** Whether a pinned oMLX fork is acceptable as the shipped form, and whether it is public. +4. **Claim.** A model that is not exact is the product. The demo should not be published before + the generation check. +5. **Housekeeping.** A planner's Python call created `~/Development/_refs/omlx/.venv` (1.8 GB), + `uv.lock` and an egg-info directory in the oMLX clone. Remove or keep. +6. **Paper title.** Still open; this plan does not depend on it. + +## 9. Runbook for K1 (first generated token, stock oMLX) + +Checked 2026-11-09 by reading the clone and running only `omlx --help`: the environment at +`~/Development/_refs/omlx/.venv` imports (`mlx` 0.32.4, `omlx` 0.7.4.dev1). Not yet done: the +native kernels are not built, and no model has been loaded. + +Run when no other heavy job is on the Mac (E1 holds it tonight): + +1. `make kernels` in the clone (oMLX's fused paths need them; a plain install skips them). +2. A model directory containing a link to `~/Models/v41-reap50`. oMLX's V4.1 loader accepts the + original checkpoint format; whether it accepts a 192-expert count is the first thing this run + tests. +3. Model settings: `moe_expert_offload_enabled`, `moe_expert_offload_resident_fraction` 0.375 + (72 of 192 experts at native FP4, about 54 GB), `deepseek_v41_engram_ssd_offload`; memory guard + below 80 GB. oMLX's own figures show peak footprint about 12 GiB above its Metal-active size, so + start lower (0.25) if the first load comes close to the cap. +4. `omlx serve --model-dir …`, then two 64-token chat completions on one prompt (cold, then warm) + and the same prompt replayed for the no-miss floor. +5. Record: decode tok/s, prefill tok/s, hit rate (oMLX logs it per request), peak footprint, and + whether every layer stayed on the fused path. + +Reference points from oMLX's own document: on a smaller MoE model it measures 122.5 tok/s fully +resident, 59.4 at 50% residency and 40.1 at 25%, so waiting on misses costs it a factor of two to +three there; on V4.1 it says decode "is bound by miss latency at one to two misses per layer per +token". + +Not runnable yet: the offline checks for ideas 2 and 6 (K6) need per-token router demand, which +the harness does not dump; the stored routes are a 2,048-token sample without token ids. + +### K1 result (2026-11-10) + +Run on the M4 Max from the pinned clone, native FP4 experts, Engram on SSD, greedy decoding, one +request at a time. Scripts: `scripts/v41_build/omlx_counters.py` (reads oMLX's expert-slot +counters, which its server does not log for V4.1) and `scripts/v41_build/omlx_compact.py` (a +fixed resident set through a linked model directory and three read-side patches; the clone is not +modified). + +| Configuration | Resident experts | Decode tok/s | Prefill tok/s | Memory after load (peak) | +|---|---:|---:|---:|---:| +| oMLX expert offload, waits on a miss | 48 of 192 | 3.3–4.5 | 6 | 45.4 GiB (54) | +| fixed set, sliced router | 48 | 30–32 | 150 | 45.4 GiB (55) | +| fixed set, sliced router | 64 | 30–31 | 145 | 56.6 GiB (67) | + +- oMLX's offload path: 78% hit rate on a 256-token completion, about 47 misses per token at + 18.8 MB each, so roughly 0.9 GB read per token; an exact replay of a 45-token exchange still + misses 39% of lookups at 48 slots. Its offloaded experts are slot wrappers that the fused decode + kernels do not accept (`decode_moe.supported`), so this path is slow even on a hit. +- The fixed set runs on the fused kernels because the experts are plain stacked arrays and + `n_routed_experts` is the resident count. This is the compact router's static case, and it is + seven times the speed of waiting on a miss at equal memory. The estimate of 15–20 tok/s was low. +- Consequence for the cache: keep the experts as one stacked array of C rows per projection and + overwrite a row in place on admission, with the router row rewritten in the same step. Do not + build on oMLX's slot wrapper. +- Memory: about 11.8 GiB for everything but routed experts, 0.70 GiB per expert at FP4, and about + 10 GiB of extra peak during load. Under the 80 GB cap that is about 72 experts at FP4 and about + 155 at 2 bits. +- Quality at these sizes is poor, as the teacher-forced numbers predict: 48 experts give clipped + but on-topic text, and 64 fall into a loop within a 1,024-token greedy completion. Speed is + settled; quality needs the cache and 2-bit experts. +- Side effect on the machine: the server wrote `~/.omlx/bin/omlx-cluster-python`, a link to the + clone's Python. + +### No-wait cache, first version (2026-11-10) + +`scripts/v41_build/omlx_routed.py`, on the fixed-set base above, clone unmodified. 64 resident +experts per layer at native FP4, drawn from the pruned pool of 192; admission at gate weight 0.2, +4 spare slots per layer, least recently demanded evicted. + +How a miss costs no time: + +- Every layer also scores the whole pool's router rows on its input. The demanded experts ride in + the token's own graph, so reading them two calls later waits for nothing. +- The router is the switch. An evicted slot's bias is set so nothing routes to it; reader threads + overwrite its rows in place only after every call built before that has finished; its router + row becomes the new expert only after the bytes are in. A token never waits on a read and never + runs a half-written expert. An admitted expert serves from about three tokens after the demand. +- Stored FP4 bytes are the array's bytes, so an install is six positional reads straight into + the slot, 18.8 MB, with no copy. + +| Run | Decode tok/s | Notes | +|---|---:|---| +| fixed 64, no demand tracking | 30–31 | from K1 | +| cache code, admissions off | 26.6–28.2 | output identical to fixed 64 over 1,024 tokens; demand scoring costs 10–14% | +| cache on | 22.0–26.1 | 8.3 installs per token (156 MB), 75% of demanded slots resident | +| oMLX offload, exact, 48 slots | 4.5 | reference text below | + +- Checks: with admissions off the output equals the fixed set's, token for token. With them on, + every routed slot in the first and last layer matched its expert's stored bytes at 20 + checkpoints across 12,785 installs. A torn read in flight is excluded by the ordering rule, not + by a measurement. +- Quality is not there at this size. On a 1,024-token coding prompt the exact pool of 192 writes a + correct script; fixed 64 and cache 64 both leave the reference within a few words and end in a + loop. The cache's text is more fluent early on, which is not evidence. 64 experts is what FP4 + allows under the 80 GB cap; E1 puts the useful range at 96–144 residents. +- So the next build item is 2-bit experts in this runtime (K2, then lattice codes as slot arrays): + about 155 residents under the same cap. Also open: demand scoring through oMLX's own gate kernel + to recover the 10–14%; 23% of wanted admissions found no free slot at 4 spares; prefill runs + through the same code but was only exercised on short prompts. + +### K2 result (2026-11-10) + +`scripts/v41_build/bench_expert_formats.py`, M4 Max, one process, MLX 0.32.4, 64 experts, 40 +chained MoEs per eval, three runs (medians agree within 5%). oMLX rows time its fused decode MoE +(router, routed experts, shared expert) on random weights; the lattice row times this repo's +`e8p_m1` kernel on layer 20's real 2-bit codes, routed experts only. Cost per expert is the +difference between top-6 and top-1, over five. + +| Routed-expert format | Bits per weight | Expert, MB | MoE at top-6, ms | Per expert, ms | Fixed, ms | +|---|---:|---:|---:|---:|---:| +| oMLX MXFP4 (checkpoint format) | 4.25 | 18.8 | 0.39–0.41 | 0.042–0.045 | 0.145 | +| oMLX affine 2-bit, group 64 | 2.5 | 11.1 | 0.41 | 0.046–0.048 | 0.124–0.134 | +| oMLX affine 3-bit, group 64 | 3.5 | 15.5 | 0.35–0.36 | 0.035–0.036 | 0.140–0.145 | +| E8P lattice 2-bit (routed only) | 2.04 | 9.0 | 0.29 | 0.036–0.038 | 0.065–0.069 | + +- The lattice kernel is not slower per expert than any oMLX format, so K2 does not kill it. It + is 3.3 times this repo's generic kernel at this shape (0.95 ms). +- What it buys is residents, not speed: under the 80 GB cap about 155 lattice experts fit + against about 127 at affine 2-bit. E1 puts a cache of 144 within 0.009 KL of the whole pool of + 192, so the lattice format clears that line and affine 2-bit does not. +- Not measured: the lattice kernel inside oMLX's decode step. The router and shared expert then + run as their own dispatches; from the fixed costs above the MoE would take about 0.43 ms + against 0.39–0.41, roughly 3–5% of a token. +- Side note: oMLX's 3-bit affine MoE is its fastest, ahead of FP4. +- Next build item: lattice experts as slot arrays in the cache. That needs the codes stored as + raw per-expert bytes (they are `.npz` today), the routed half of the fused MoE replaced by + `e8p_m1`, and codes for the experts outside the pruned pool. + +### Lattice experts in the cache (2026-11-10) + +`scripts/v41_build/omlx_lattice.py`. oMLX (pinned clone, unmodified) loads a six-expert model +made by `omlx_compact.py make --experts 6`; after the load each layer's routed experts are +replaced by 144 lattice experts held as stacked code and scale arrays. + +- Storage needed no conversion. The quantizer's `layer-NN.npz` files are stored uncompressed, so + an expert's codes are a fixed byte range of the file and an install is a positional read into + the slot's rows. +- The router keeps every row of the pool and scores all 192 once per token with oMLX's gate + kernel. One small kernel takes two top-6's of those scores: among the residents (every other + expert's bias is -1e30), which is what runs, and over the pool, which is the demand the policy + reads two calls later. An install or eviction changes one bias and one entry of an + expert-to-slot table. This replaced the separate demand matmul of the FP4 cache. +- Decode step: gate kernel, top-k kernel, `e8p_m1` for w1 and w3, oMLX's SwiGLU and act-quant + kernel, `e8p_m1` for w2, then oMLX's shared-expert kernels, which add the routed sum. The + shared up kernel runs with no routed column (`TOPK` 0). +- Prefill: oMLX's unfused MoE with the routed experts on this repo's sorted-route kernels. +- `omlx_routed.py` is now the policy with the FP4 storage as one subclass; the lattice cache is + the other. + +Self-check (`omlx_lattice.py check`, two layers of real codes, random routers and shared +experts): the decode step is within 0.3% of the same MoE computed from matrices decoded on the +host, the prefill path within 0.1% of the decode path, and after 83 installs no slot differs +from its stored bytes. The decode-step MoE takes 0.40–0.42 ms against 0.38–0.41 for oMLX's +fused FP4 MoE in the same harness. The lattice kernel is 1.7 times slower on fp32 activation +rows than on 16-bit ones (0.51 against 0.29 ms for the routed half), so it gets bf16. + +Serving, M4 Max, pool of 192, 144 residents per layer, tau 0.2, 4 spares: + +| | Decode, tok/s | | +|---|---:|---| +| Policy off (fixed 144) | 28.6–30.1 | text reproducible run to run | +| Cache on | 26.0–27.9 | 2.2 installs and 20 MB read per token; 96% of demanded slots resident | +| oMLX FP4, fixed 64 (K1) | 30–31 | | +| FP4 cache, 64 residents | 22–26 | 8.3 installs and 156 MB per token; 75% resident | + +- Memory: 62 GiB steady, 68–69 GiB peak during the load, which takes about 25 s. +- Prefill: a 4,341-token prompt took 22.6 s and 39.1 s (192 and 111 tok/s, two runs, no KV + reuse). Prompts of 38–64 tokens reach the first token in 1.4–2.3 s. A prefill is one forward + call, so the cache cannot adapt inside it; over the runs with the long prompt 83% of demanded + slots were resident. +- Policy cost: 1.5 ms of Python per token, and about 2 ms more that is not attributed yet. + Reading the demand through a lazy reshape forced a GPU sync per layer and cost another 10% + (22.6–25.8 tok/s) until the raw kernel outputs were kept instead. +- Installs on the real model: 17 checks over 2,431 installs, no slot differing from its bytes. +- 5–11% of wanted admissions found no free slot at 4 spares. + +Text, same prompt as before (a script that finds duplicate files by SHA-256): + +- Thinking off, temperature 0: a complete script in 1,035–1,182 tokens, and it runs correctly + on a test tree. A 4,341-token prompt holding `omlx_routed.py` got an accurate five-bullet + summary and a plausible bug risk. +- Thinking on: no answer within 4,096 tokens in any of three runs. Greedy with the cache on is + coherent for about 1,000 tokens and then repeats itself; default sampling does not repeat but + keeps reviewing code it drafted in its reasoning; greedy with the policy off (fixed 144) + also runs out. So the cache is not the cause. +- The exact pool of 192 at FP4 (stock oMLX offload, 7.1 tok/s, greedy) fails the same way: 4,096 + reasoning tokens, no answer, and from about character 10,500 of 16,235 it repeats the same + list of options it will not add. The lattice cache starts repeating at about character 5,300; + the fixed lattice set and the sampled run do not repeat. So neither the 2-bit codes nor the + cache causes the failure with thinking on: the pruned checkpoint has it at full precision. + Whether the unpruned model has it too is not known; that model is not on the Mac. +- Given room, thinking does finish under default sampling: the lattice cache with a budget of + 12,000 tokens reasoned for 7,567 tokens without repeating itself, then answered in 620 with + a script that runs correctly on the test tree. 400 s in all, 20.7 tok/s over the 8,187 + tokens (1.5 installs per token, policy 2.4 ms per token). The drop from 26–28 tok/s at + 1,000 tokens is not attributed; context length is the first suspect and the fixed set was + never timed this far out. +- With the cache on, greedy output is not reproducible: install timing changes routing. Two + runs diverged after 18 characters; a third matched the second. + +Not done: more than 144 residents (160 would be about 67 GiB steady), memory pressure, and +multi-token verify windows in serving (the self-check covers 8 rows). + +### Pool of 384 (2026-11-10) + +Every expert of the full checkpoint is now reachable on the Mac at 2 bits. + +- Codes. The 192 experts per layer that REAP-50 removed were copied from the Runpod volume + through a CPU pod (`extract_tensors.py`, 144 GB, 76 minutes, about $0.16) and fitted on the Mac + with the unchanged Metal search: 2.5 h for 40 layers, weighted error 0.091 on w1 and w3 and + 0.079 on w2 (the first 192: 0.090 and 0.074). The torch port refits 64 rows of nine + projections to the same codes. All 384 router rows per layer came with them; the rows of the + 192 survivors equal the REAP-50 checkpoint's. +- Cache. `OMLX_LATTICE` may name a pool file: code directories with the pool number of each + expert, a router file and routed-mass statistics for the starting residents. oMLX still loads + the REAP-50 checkpoint (everything outside the routed experts is shared), and the cache + replaces router and experts. Self-check over 384: same gaps as over 192; the decode-step MoE + takes 0.444 ms (0.40–0.42 over 192). +- Starting residents are the 144 experts with the most routed mass in calibration: 83% of the + mass on average, 66% in the worst layer. On average 118 of them are REAP-50 survivors. + +Serving, 144 residents, tau 0.2, 4 spares, 63 GiB (the pool of 192: 62 GiB): + +| Request | Tokens | tok/s | Demanded slots resident | Installs per token | Read | +|---|---:|---:|---:|---:|---:| +| Script, thinking off, greedy (first request) | 1,124 | 26.6 | 88% | 8.9 | 80 MB per token | +| Same again | 748 | 23.7 | 88% | 9.1 | 82 MB per token | +| 4,341-token prompt, summary | 282 | 22.6 | | | prefill 23.6 s (184 tok/s) | +| Script, thinking on, greedy, 4,096 budget | 4,096 | 19.1 | 91% | 6.3 | 1.05 GB/s | +| Script, thinking on, sampled, 12,000 budget | 3,256 | 17.3 | | 8.7 | 1.34 GB/s | + +- Against the pool of 192 at the same 144 residents: four times the installs (2.2 per token + there), 88–91% of demanded slots resident against 96%, and 0–15% lower speed with thinking + off (26.0–27.9 there). +- Text. The first thinking-off script runs correctly on the test tree. The second, from a + different cache state, has one wrong token (`sha.hexdigdigest()`) and crashes. One sample; + the pool of 192 produced no such error in its runs. +- Thinking on, greedy: no answer in 4,096 tokens, repeating from about character 12,800 of + 16,235, as the pruned models do. Thinking on, sampled: 2,704 reasoning tokens, then a + 552-token answer whose script runs correctly, 192 s in all (the pool of 192 needed 7,567 + reasoning tokens; one sample each). +- 7% of wanted admissions found no free slot in the first request, 5% over the thinking runs. + +Not measured: memory pressure. + +### Tau sweep over the pool of 384 (2026-11-10) + +Quality against the full FP4 model needs only its stored logits, not the model: the pod wrote +them for the Mac's 22 eval sessions. `local_eval.py --router` puts the 384-row router on the +REAP-50 checkpoint and takes every routed expert from the 2-bit codes the Mac serves, so the +simulated cache (144 residents, re-route on a miss, one-token lag, starting residents by routed +mass) runs on the Mac. First 1,024 tokens of each session, KL against the full model: + +| Variant | KL | Top-1 | Installs per token | Read at 30 tok/s | +|---|---:|---:|---:|---:| +| All 384 resident (2-bit coding alone) | 0.390 | 84.4% | 0 | 0 | +| Cache of 144, best set per record (oracle) | 0.424 | 84.0% | | | +| Cache of 144, tau 0.1 | 0.483 | 82.5% | 17.1 | 4.6 GB/s | +| Cache of 144, tau 0.2 | 0.460 | 82.7% | 7.6 | 2.1 GB/s | +| Cache of 144, tau 0.3 | 0.454 | 82.5% | 2.6 | 0.7 GB/s | +| Cache of 144, tau 0.45 | 0.573 | 80.6% | 1.1 | 0.3 GB/s | +| Static 144 (hottest by routed mass) | 0.763 | 76.2% | 0 | 0 | + +For reference, same sessions and tokens: the exact REAP-50 checkpoint (192 FP4 experts) 0.578, +REAP-50 at 2 bits 0.747, the cache of 144 inside the pool of 192 0.756. + +- The cache of 144 over the full pool beats the static 144 by 0.309 (standard error 0.052, 21 + of 22 sessions) and the exact REAP-50 checkpoint by 0.12, at 53 GB of resident experts. +- Tau 0.2 and 0.3 are equal (difference 0.005, standard error 0.017) and 0.3 reads a third as + much. Tau 0.1 is worse than 0.2 (0.024, standard error 0.013): admitting weak slots evicts + useful experts. Tau 0.45 loses 0.12 in every session. +- At tau 0.3 the cache is 0.030 from the oracle and 0.064 from having every expert resident. +- Tool-use text stays at 1.8–2.0 in every variant, all 384 resident included (1.94). That is + the off-format teacher found on 2026-11-09 (the eval text carries foreign tool markup), not a + loss from the cache or the coding. Without that domain: 0.23 with all 384 resident, 0.32 for + the cache of 144 at tau 0.3, 0.66 for the static 144. +- Run: 11,991 s, peak RSS 63.8 GB (`~/Models/v41-reap-run/jobs/tau384_mac.sh`). + +Serving at tau 0.3 (144 residents, 4 spares, 63 GiB), against the tau 0.2 table above: + +| Request | Tokens | tok/s | Demanded slots resident | Installs per token | Read | +|---|---:|---:|---:|---:|---:| +| Script, thinking off, greedy (first request) | 1,030 | 27.2 | 80% | 2.5 | 22 MB per token | +| Same again | 1,085 | 26.7 | 83% | 1.4 | 13 MB per token | +| 4,341-token prompt, summary | 310 | 23.0 | 71% | | prefill 24.8 s (175 tok/s) | + +Both scripts run correctly on the test tree. Installs match the simulation (2.6 per token). +Tau 0.3 is the serving default from here. + +Limits: the cache is simulated with a fetch that lands within one token; 22 sessions, first +1,024 tokens; two generated scripts at tau 0.3; no pod confirmation on longer sessions (E7). + +### Drafter on the no-wait cache (2026-11-11) + +The checkpoint's drafter (DSpark: 3 stages, 128 experts each, 7.4 GiB) is switched on by +`mtp_enabled` in `/model_settings.json`; no code change was needed. A cycle drafts 3 +tokens and verifies them in one backbone forward of about 4 rows through the same hook the +single-token path uses. 120 residents (so the footprint stays under the 80 GB cap), tau 0.3, +pool of 384, installed slots byte-checked (0 mismatches): + +| Request | Drafter off | Drafter on | Gain | +|---|---:|---:|---:| +| Script, thinking off, greedy | 24.3 tok/s | 31.2 | +28% | +| Same, temperature 1.0, top_p 1.0 | 24.2 | 30.5 | +26% | +| 4,341-token prompt, summary | 21.8 | 23.9 | +9% | +| Script, thinking on, temperature 1.0 | 17.45 (3,872 tokens) | 25.0 (8,492 tokens) | +43% | +| Footprint, steady / peak | 55 / 61 GiB | 62 / 68 GiB | +7 GiB | + +- About 2.9 tokens are kept per cycle; a cycle takes 93 ms against 41 ms for a single token. +- Installs per forward row are unchanged; per emitted token they are 1.38 times the no-drafter + figure (0.67 GB/s read against 0.41 with thinking on). +- oMLX's own offload gains 29% from the drafter (6.6 to 8.5 tok/s) because verification waits + on experts; here nothing waits. +- A multi-row lattice kernel (one decode per distinct expert in the block) is not worth + building yet: over the full model's routes, distinct experts per route are 0.84 at 2 rows, + 0.70 at 4 and 0.63 at 6, so it saves at most about 12% of a cycle. +- Job: `~/Models/v41-reap-run/jobs/k1_mtp.sh`. + +### Sampling at the card's settings (2026-11-11) + +The same script request, 8 samples per row, 4,096-token budget, thinking off. "Correct" means +the script runs and lists the three duplicate files of the test tree. + +| Model | Residents | Temperature | Drafter | Correct of 8 | +|---|---:|---:|---|---:| +| Cache over the pool of 384, 2-bit | 120 | 1.0 | off | 4 | +| Same | 120 | 1.0 | on | 4 | +| Same | 144 | 1.0 | off | 6 | +| Same | 120 | 0.6 | off | 7 | +| Exact REAP-50 checkpoint (192 FP4 experts, stock oMLX offload, 4–5 tok/s) | 192 | 1.0 | off | 5 | + +- Failures are single-token slips (`hexdigist`, `mun` for `min`, a constant misspelled at its + use, a missing colon) and two samples with wrong output. Every greedy script was correct. +- With 8 samples no row is separable from another (4 of 8 against 7 of 8 is p = 0.28, Fisher). + The exact pruned checkpoint at its original precision fails 3 of 8, so the slips are not + evidence against the 2-bit coding, the cache or the drafter. +- Not known: how often the full model slips at temperature 1.0 (it does not fit on the Mac), + and whether more samples separate the rows. 16 more per row for the cache of 144 and the + exact checkpoint are queued behind the simulator run. +- Jobs: `~/Models/v41-reap-run/jobs/k1_mtp_samples.sh`; scorer in the session scratchpad. + +### Quality against resident count (C1, 2026-11-11) + +Same simulator, sessions and tokens as the tau sweep: pool of 384, tau 0.3, re-route on a miss, +22 sessions, first 1,024 tokens, KL against the full FP4 model. Footprint is the measured +serving model (11.8 GiB outside routed experts plus 0.336 GiB per resident slot). + +| Residents per layer | Resident experts | Footprint | KL | KL without tool-use | Top-1 | Read per token | +|---:|---:|---:|---:|---:|---:|---:| +| 32 | 11.8 GB | 23 GiB | 0.942 | 0.732 | 73.5% | 118 MB | +| 48 | 17.7 GB | 28 GiB | 0.790 | 0.582 | 77.1% | 82 MB | +| 64 | 23.6 GB | 33 GiB | 0.705 | 0.496 | 78.7% | 61 MB | +| 96 | 35.4 GB | 44 GiB | 0.605 | 0.390 | 81.1% | 39 MB | +| 144 | 53.1 GB | 60 GiB | 0.454 | 0.320 | 82.5% | 24 MB | +| 200 | 73.7 GB | 79 GiB | 0.401 | 0.279 | 83.8% | 13 MB | +| 250 | 92.2 GB | 96 GiB | 0.398 | 0.254 | 84.4% | 5.5 MB | +| 384 (all) | 141.6 GB | | 0.390 | 0.235 | 84.4% | 0 | +| Static 144, for reference | 53.1 GB | 60 GiB | 0.763 | 0.654 | 76.2% | 0 | + +- 144 reproduces the tau sweep (0.4542, every session equal). +- A cache of 48 at 28 GiB is level with the static 144 at 60 GiB (0.790 against 0.763; 0.582 + against 0.654 without tool-use), and a cache of 96 at 44 GiB is level with the exact REAP-50 + checkpoint (0.605 against 0.578), which needs 341 GB of FP4 experts on disk. +- 144 to 200 gains 0.053 (standard error 0.009, 22 of 22 sessions). 200 to 250 gains 0.003 + over all sessions and 0.024 without tool-use: breadth is spent near 200 residents. +- Reads grow as the cache shrinks: 2.5 GB/s at 30 tok/s for 48 residents, 3.5 GB/s for 32. +- Run: 12,400 s (`~/Models/v41-reap-run/jobs/c1b1_mac.sh`). + +### FP4 on the hottest experts (B1, 2026-11-11) + +`local_eval.py --hot NAME=H`: the H hottest experts per layer by routed mass, among those the +REAP-50 checkpoint holds, run from its FP4 weights whenever they are resident (18.8 MB against +9.2 MB at 2 bits). This is the ceiling for any extra bits on those experts. + +| Variant | Resident experts | KL | Without tool-use | Top-1 | +|---|---:|---:|---:|---:| +| 144 at 2 bits | 53.1 GB | 0.454 | 0.320 | 82.5% | +| 127 residents, 16 of them FP4 | 53.0 GB | 0.413 | 0.310 | 83.3% | +| 110 residents, 32 of them FP4 | 52.8 GB | 0.476 | 0.331 | 82.9% | +| 177 at 2 bits | 65.2 GB | 0.421 | 0.294 | 83.2% | +| 144 residents, 32 of them FP4 | 65.4 GB | 0.392 | 0.280 | 83.9% | +| 200 at 2 bits | 73.7 GB | 0.401 | 0.279 | 83.8% | +| All 384 at 2 bits | 141.6 GB | 0.390 | 0.235 | 84.4% | +| All 384, 32 of them FP4 | 153.8 GB | 0.268 | 0.167 | 86.9% | + +Paired by session (difference, standard error, sessions better): + +| Comparison | All 22 sessions | 20 without tool-use | +|---|---|---| +| 16 FP4 against 144 at 2 bits, 53 GB | −0.041 ± 0.023 (12) | −0.010 ± 0.006 (10) | +| 32 FP4 against 144 at 2 bits, 53 GB | +0.022 ± 0.010 (5) | +0.011 ± 0.006 (5) | +| 144 with 32 FP4 against 177 at 2 bits, 65 GB | −0.029 ± 0.013 (14) | −0.015 ± 0.009 (12) | +| 144 with 32 FP4 against 144 at 2 bits (12 GB more) | −0.062 ± 0.016 (19) | −0.041 ± 0.008 (17) | +| All 384 with 32 FP4 against all 384 at 2 bits | −0.122 ± 0.039 (22) | −0.067 ± 0.013 (20) | + +- Gate (0.02 over the equal-memory 2-bit rival): met with tool-use counted (0.041 and 0.029), + not met without it (0.010 and 0.015). The tool-use sessions carry the off-format teacher, so + the reading without them is the one to trust. At 53 and 65 GB, FP4 on the hot set is a + small gain at best, and trading 34 residents for 32 FP4 experts is a loss. +- Precision alone is large: with every expert resident, FP4 on 32 experts per layer (8% of + them) removes 0.067 of the 0.235 left, in all 20 sessions. Those experts take 22% of routed + slots. +- 144 residents with 32 at FP4 (65 GB) equals 200 residents at 2 bits (74 GB). +- So extra bits pay where breadth is spent, above about 200 residents, which is the budget of + a 128 GB Mac. That budget was not tested here; 217 residents with 32 FP4 and 183 with 64 + FP4 against 250 at 2 bits (all 92 GB) is the deciding run. +- Run: 8,549 s, peak RSS 61 GB. + +### Serving at 48 and 96 residents (C2, 2026-11-11) + +Same sampled script request as above, card sampling (temperature 1.0, top_p 1.0), 8 samples, +drafter off, tau 0.3. Fewer slots stand in for a smaller Mac's expert memory, not for its +memory pressure. + +| Residents | Footprint | tok/s | Demanded slots resident | Installs per token | Read per token | Correct of 8 | +|---:|---:|---:|---:|---:|---:|---:| +| 48 | 30 GiB | 20–27 | 65% | 5.8 | 53 MB | 0 | +| 96 | 46 GiB | 22–25 | 75% | 3.2 | 28 MB | 2 | +| 120 (above) | 55 GiB | 24 | | | | 4 | +| 144 (above) | 63 GiB | 17–27 | | | | 6 | + +- Speed holds at small sizes; the text does not. At 48 residents the output is misspelled + throughout (`hashllib`, `OSEError`, `o a file`) and 6 of 8 samples run to the 4,096-token + limit. At 96, 3 of 8 run to the limit. +- Correct scripts rise with residents: 0, 2, 4, 6 of 8 at 48, 96, 120, 144. A KL near 0.8 is + unusable under the card's sampling and 0.45 is marginal, so the KL curve above overstates + what a small Mac gets. The exact REAP-50 checkpoint (KL 0.578) gives 5 of 8. +- Open: whether a lower temperature or a trimmed tail (top_p 0.95) makes 96 residents usable, + and how many of 8 the full model gets. The first is queued (`jobs/after_b1b.sh`). + + +### FP4 on the hottest experts at 92 GB (B1b, 2026-11-11) + +The deciding run named above: the expert budget of a 128 GB Mac (92 GB), pool of 384, tau 0.3, +re-route, 22 sessions, 1,024 tokens. KL against the full model; "no tool-use" drops the two +off-format sessions. + +| Variant | Residents | Expert memory | KL | KL, no tool-use | Top-1 | +|---|---|---:|---:|---:|---:| +| 250 at 2 bits (C1) | 250 | 92 GB | 0.398 | 0.254 | 84.4% | +| 217, hottest 32 at FP4 | 217 | 92.3 GB | 0.325 | 0.219 | 85.6% | +| 183, hottest 64 at FP4 | 183 | 92.0 GB | 0.328 | 0.221 | 86.1% | +| All 384 at 2 bits (C1) | 384 | 139 GB | 0.390 | 0.235 | | +| All 384, hottest 32 at FP4 (B1) | 384 | 154 GB | 0.268 | 0.167 | 86.9% | + +Paired differences, mean ± standard error (sessions better): + +| Pair | All 22 | Without tool-use (20) | +|---|---|---| +| 217 with 32 FP4 against 250 at 2 bits | −0.072 ± 0.027 (17) | −0.036 ± 0.010 (15) | +| 183 with 64 FP4 against 250 at 2 bits | −0.070 ± 0.030 (16) | −0.034 ± 0.011 (14) | +| 217 with 32 FP4 against 200 at 2 bits (74 GB) | −0.075 ± 0.016 (21) | −0.060 ± 0.009 (19) | +| 217 with 32 FP4 against all 384 at 2 bits (139 GB) | −0.064 ± 0.035 (15) | −0.016 ± 0.009 (13) | +| 183 with 64 FP4 against 217 with 32 FP4 | +0.002 ± 0.015 | +0.002 ± 0.003 | + +- Gate (0.02 over the equal-memory 2-bit rival): met on both readings, 0.072 with tool-use and + 0.036 without. B continues. +- At 92 GB, 217 residents with 32 at FP4 are at least level with every expert resident at 2 + bits, which takes 139 GB. The gain is in the agentic, networking and cyber sessions (0.32 to + 0.25, 0.25 to 0.16, 0.48 to 0.38); prose and business sessions do not move. +- 32 and 64 FP4 experts per layer score the same, so the hot set is small: 32 is enough, and + the choice between them can be made on serving cost. +- With B1: bits lose to breadth at 53 GB, draw at 65 GB, win at 92 GB. The crossover is near + 200 residents. +- Serving this needs about 96 GiB of footprint (185 slots at 2 bits, 32 at FP4, 11.8 GiB + outside the routed experts), above the 80 GB per-process rule on this Mac. Under the rule + the hot tier can be served at about 130 residents with 32 at FP4. +- Run: 3,339 s, peak RSS 60 GB. + +### More samples and two rescue checks (2026-11-11) + +Same sampled script request, 8 samples per run, 4,096-token budget, drafter off. + +| Model | Residents | Temperature | top_p | Correct | tok/s | +|---|---:|---:|---:|---|---:| +| Cache, pool of 384 at 2 bits | 144 | 1.0 | 1.0 | 17 of 24 (6, 5, 6) | 22–27 | +| Exact REAP-50, FP4, stock offload | 192 | 1.0 | 1.0 | 18 of 24 (5, 6, 7) | 5.6–6.3 | +| Cache | 144 | 1.0 | 0.95 | 5 of 8 | 24–27 | +| Cache | 120 | 0.6 | 1.0 | 7 of 8 (earlier) | | +| Cache | 96 | 1.0 | 1.0 | 2 of 8 (C2) | 22–25 | +| Cache | 96 | 1.0 | 0.95 | 2 of 8 | 24–26 | +| Cache | 96 | 0.6 | 1.0 | 1 of 8 | 20–28 | +| Cache | 48 | 0.6 | 1.0 | 0 of 8 | 23–25 | + +- With 24 samples each, the cache of 144 at 2 bits and the exact REAP-50 checkpoint are level + (71% and 75%) at the card's sampling, and the cache is about four times as fast. Their + failures are of the same kind: one misspelled name or a wrong output. +- Neither a lower temperature nor top_p 0.95 rescues 96 or 48 residents. The failures at 96 + with temperature 0.6 are still misspelled identifiers (`hexdigdigest`, `diggest`, + `filnames`), so the damage comes from the re-routed slots (25% of demand at 96 residents, + 35% at 48), not from the sampling tail. top_p 0.95 does nothing at 144 either. +- Between 96 and 120 residents at temperature 0.6 the count goes from 1 to 7 of 8. The usable + floor on this request is about 120 residents (55 GiB), which is a 96 GB Mac, not a 64 GB one. +- Still unknown: how many of 24 the full model gets (it does not fit on this Mac). + +### Serving 200 residents (2026-11-11) + +Jack raised the per-process limit to 100 GB for serving. Same sampled script request, card +sampling, 24 samples, drafter off, tau 0.3, 200 residents at 2 bits over the pool of 384. + +| Residents | Footprint, peak | Demanded slots resident | Installs per token | Read per token | Correct of 24 | +|---:|---:|---:|---:|---:|---:| +| 144 (above) | 63 GiB | | | | 17 | +| 200 | 87 GiB | 90% | 0.40 | 3.6 MB | 19 | +| Exact REAP-50 (above) | | | | | 18 | + +- 19 of 24, and none of the 5 failures is a misspelled name or a syntax error: 4 scripts run + and print the wrong list, 1 answer has no code block. At 144 residents 5 of 7 failures were + misspelled names or syntax errors. Fewer re-routed slots (10% of demand against 25% at 96) + remove the spelling damage. +- Footprint starts at 81 GiB and peaks at 87 GiB over three server runs, 6 to 8 GiB above + the 11.8 + 0.336 per resident estimate, so the estimate needs that margin at this size. +- Speed was 11–18 tok/s, but the residual fit shared the GPU during these runs; a clean + speed run is queued (`jobs/after_b2.sh`). + +### Third bit on the hottest experts (B2, 2026-11-12) + +A 1-bit E8 residual was fitted on the 2-bit base for the 192 survivors +(`keep_quant.py --residual-of quant/e8p-podimp --code-bits 8`, 310 s per layer, 33 GB). +Weighted weight error falls from 0.084 to 0.030 (gate and up) and from 0.047 to 0.024 (down). +A hot expert then runs base + residual, 13.8 MB instead of 18.8 MB for FP4 +(`local_eval.py --hot-tier DIR`). Same 92 GB budget and settings as B1b. + +| Variant | Residents | KL | KL, no tool-use | Top-1 | +|---|---:|---:|---:|---:| +| 250 at 2 bits (C1) | 250 | 0.398 | 0.254 | 84.4% | +| Hottest 32 at FP4 (B1b) | 217 | 0.325 | 0.219 | 85.6% | +| Hottest 32 at 3 bits | 234 | 0.327 | 0.223 | 85.3% | +| Hottest 64 at 3 bits | 218 | 0.311 | 0.215 | 85.9% | +| Hottest 96 at 3 bits | 202 | 0.308 | 0.222 | 85.8% | + +Paired differences, mean ± standard error (sessions better): + +| Pair | All 22 | Without tool-use (20) | +|---|---|---| +| 64 at 3 bits against 32 at FP4 | −0.015 ± 0.008 (16) | −0.004 ± 0.003 (14) | +| 32 at 3 bits against 32 at FP4 | +0.002 ± 0.004 (10) | +0.004 ± 0.004 (9) | +| 96 at 3 bits against 64 at 3 bits | −0.002 ± 0.011 (8) | +0.007 ± 0.004 (7) | +| 64 at 3 bits against 250 at 2 bits | −0.087 ± 0.034 (19) | −0.039 ± 0.009 (17) | +| 64 at 3 bits against all 384 at 2 bits (139 GB) | −0.079 ± 0.042 (14) | −0.020 ± 0.008 (12) | + +- The residual does the work of the FP4 copy at equal memory: level with it on the reading + without tool-use, slightly ahead with it. One more bit on a hot expert recovers what FP4 on + it recovers, because the memory saved goes to more hot experts or more residents. +- 64 hot experts per layer is the best of the three; 96 gives nothing more. +- Unlike an FP4 copy, the residual adds to the base the cache already holds. A hot expert + needs no second format in the slot array, and losing the residual leaves the expert running + at 2 bits. +- Run: 5,854 s, peak RSS 63 GB. Smoke (one session, 96 tokens, 144 residents): 1.521 at 2 + bits, 1.309 with 32 hot experts at 3 bits. + +### Clean speed at 200 residents, and 180 residents with the drafter (2026-11-12) + +Same sampled script request, card sampling, tau 0.3, nothing else on the GPU. + +| Configuration | Footprint, peak | tok/s | Correct | +|---|---:|---:|---| +| 200 residents, drafter off | 87 GiB | 21–28 (3 samples) | 19 of 24 (above), 3 of 3 here | +| 180 residents, drafter on | 88 GiB | 22–41, median 32 (24 samples) | 21 of 24 (7, 8, 6) | +| 144 residents, drafter off (above) | 63 GiB | 22–27 | 17 of 24 | +| Exact REAP-50, stock offload (above) | | 5.6–6.3 | 18 of 24 | + +- More residents cost no speed: 200 residents decode as fast as 144. +- 180 residents with the drafter is the candidate default for a 128 GB Mac: every expert of + the full model reachable, 88 GiB, a median of 32 tok/s, and 21 of 24 sampled scripts + correct, against 18 of 24 at about 6 tok/s for the exact REAP-50 checkpoint on stock oMLX. + The three failures are a wrong listing, a wrong output format and one wrong keyword + argument; none is a misspelling. +- 21 against 18 of 24 is not a separable difference (Fisher p about 0.5); the claim is "no + worse, five times faster", not "better". + +### Slots from RAM (C3, 2026-11-12) + +`OMLX_LATTICE_SLOTS` now defaults to `auto` (`auto_slots` in `scripts/v41_build/omlx_lattice.py`): +a footprint budget of 70% of RAM (`OMLX_LATTICE_BUDGET_GIB` overrides it), less 20 GiB outside +the routed experts and 8 GiB when `model_settings.json` switches the drafter on, at 0.336 GiB +per resident. The constants are this M4 Max's measured peaks; other chips are unmeasured. + +| RAM | Residents, drafter off | Residents, drafter on | +|---:|---:|---:| +| 64 GB | 73 | 49 | +| 96 GB | 140 | 116 | +| 128 GB | 207 | 183 | + +- Below 120 residents the server prints that generated code was unreliable in our tests + (C2), so a 64 GB Mac starts with a warning and a 96 GB Mac should leave the drafter off. +- Checked on this Mac with the drafter on: 183 residents chosen, 83 GiB peak over two + requests, 39–41 tok/s, both scripts correct. + +### One directory, one command (C4, 2026-11-12) + +- `scripts/v41_build/slim_checkpoint.py SRC DST` writes a compact model directory as real + files without the routed experts its loader never reads: 365.8 GB to 225.9 GB for the + lattice model (the two 101.5 GB Engram shards and three others lose nothing and are cloned, + 43 shards are rewritten and read back against the source by CRC). The six stub experts per + layer, the drafter and its experts stay. +- Bundle layout, 340 GiB (`~/Models/v41-reap-run/bundle`): `models/v41-lattice` (slim + checkpoint), `codes/kept` and `codes/removed` (the 2-bit codes of all 384 experts, unchanged), + `ids-kept.json`, `ids-removed.json`, `router.safetensors`, `mass.npz`, and `pool.json`, whose + paths are now relative to it. +- `scripts/v41_build/run_v41.sh BUNDLE [PORT]` checks the machine, clones oMLX at the pinned + commit and builds its venv with uv on the first run, writes `base/model_settings.json` + (drafter on from 128 GB of RAM) and starts the server; the resident count comes from RAM. +- Checked on this Mac: with the existing oMLX install, and again from an empty oMLX directory + (fresh clone at `4213c328`, `uv sync`): 183 residents, 83 GiB, 38–40 tok/s, 2 of 3 sampled + scripts correct in each run. +- `pool384/fp4-removed` (134 GB) was deleted with Jack's confirmation; the hot tier uses + survivors only. Not done: publishing (needs Jack's OK; the Mac's uplink is 3 MB/s, so an + upload would run from a pod) and a model card for the bundle. + +### Residual kernel for the hot tier (B3, first part, 2026-11-12) + +- `src/ramp/kernels/e8_1bit_m1.metal` and `e8_1bit_m1.py`: the token and per-route layouts of + `e8p_m1` over 8-bit E8 codes, for a residual added to the 2-bit output. A route whose index + is `NO_RESIDUAL` writes zeros, so one call serves a token whose routes are partly hot. +- `tests/test_e8_1bit_m1.py` (19 cases): both layouts against a float64 decode of the same + codes, fp32 within 1e-5 and bf16 within 2e-2 of the largest output, with routes that have + no residual, a ragged last tile and the model's group size; also passes under + `MTL_SHADER_VALIDATION=1`. +- Timed alone at the model's sizes (one token, six routes, 200 interleaved runs, M4 Max): gate + 0.217 ms against 0.227 ms for the 2-bit kernel, down 0.248 against 0.252; with two of six + routes hot 0.185 and 0.194. Single calls are launch-bound, so the cost inside the decode step + is still to be measured. +- Not done: the hot residual arrays, the per-slot hot index and the memory accounting in + `omlx_lattice.py`. + +### Memory pressure (D3, 2026-11-12) + +`scripts/v41_build/d3_pressure.py` sends the sampled script request (500 tokens, card sampling) +back to back while a second process fills incompressible memory in 4 GB steps every 45 s up to +40 GB, holds it for 2 minutes, releases it and measures 2 minutes more. Server: the bundle +through `run_v41.sh`, 183 residents, drafter on. + +| Phase | tok/s, median (range), tau 0.3 | tok/s, median (range), tau 0.2 | +|---|---|---| +| Before | 36.7 (30.2–42.6) | 39.8 (37.1–42.4) | +| Taking 1–20 GB | 29.5 (25.9–32.2) | 31.4 (27.4–35.8) | +| Taking 21–40 GB | 32.8 (29.0–34.9) | 32.7 (20.7–35.4) | +| Holding 40 GB | 33.5 (30.5–34.5) | 32.2 (20.5–34.8) | +| After release | 35.9 (34.0–37.2) | 31.8 (23.5–34.6) | + +- The server keeps its 83 GiB in both runs and answers every request (61 and 62). macOS + compresses up to 24 GB and swaps 20 GB out (5 to 25 GB of swap); the pressure level reaches + "warning", never "critical". With 40 GB held, decoding is 9% slower at tau 0.3; the slowest + phase is the first allocation steps (20%), when the other process is filling memory. +- At tau 0.3 the cache installs 0.17 experts per token (1.5 MB read per token) and 90% of + demanded slots are resident; at tau 0.2, 3.7 installs and 33.6 MB per token. The heavier + reads make single requests drop to 20 tok/s under pressure and slow the recovery. +- The tau 0.2 run was a launcher bug: `run_v41.sh` did not set the admission threshold and + the code's default is 0.2, while every measurement here used 0.3. The launcher now sets 0.3. + The bundle checks under C4 above (38–40 tok/s) ran at 0.2. +- What this does not show: the server's own pages being taken. Its footprint counts + compressed and swapped pages too, so whether any of its 83 GiB left RAM is not measured; + the speed says little of it did. diff --git a/docs/research/2026-11-05-v41-reap-calibration-data.md b/docs/research/2026-11-05-v41-reap-calibration-data.md new file mode 100644 index 0000000..eb5f5cd --- /dev/null +++ b/docs/research/2026-11-05-v41-reap-calibration-data.md @@ -0,0 +1,327 @@ +# Calibration and held-out data for REAP pruning of DeepSeek-V4.1-Flash-1017 + +Date: 2026-11-05. Scope: web research only. Every Hugging Face dataset id below +was checked against `https://huggingface.co/api/datasets/` on 2026-11-05/06 +unless it is marked **UNVERIFIED** or is hosted elsewhere. + +## 0. Target model facts (from the HF config, verified) + +`deepseek-ai/DeepSeek-V4.1-Flash-1017` (sha `1bae0e10...`, MIT, modified +2026-10-17): 40 layers, `hidden_size` 5120, `moe_intermediate_size` 2304, 384 +routed experts plus 1 shared expert, top-6, `sqrtsoftplus` scoring with +`noaux_tc`, FP8 dense weights with FP4 experts (block 32x32, ue8m0 scales), and +a 64K native context YaRN-scaled to 1M. The config also carries 3 MTP/"dspark" +layers that have their own 128-expert MoE. + +Implications for calibration: + +- An expert gets about 6/384 = 1.56% of tokens on average, so 3M calibration + tokens give about 47K routed tokens per expert per layer. Rare experts can get + less than 10% of that. +- The second moment of an expert's gate/up input is a 5120x5120 matrix, and + down_proj's input is 2304-dim. A well-conditioned estimate needs many more + tokens than d_in for each expert. Dense 2-bit methods use 25-35M tokens per + model (QuIP#: 6144 x 4096; QTIP Llama-3: 4096 x 8192) [Q1, Q2]. + +## 1. State of the art: calibration data for MoE expert pruning and PTQ + +### 1.1 What REAP itself did (arXiv 2510.13999, ICLR 2026) [R1, R2] + +- **Models up to about 110B:** 1,024 samples packed to 2,048 tokens. +- **Models of about 110B and up** (Qwen3-Coder-480B, Kimi-K2): 12,228 samples + with max length 16,384, no truncation and no packing. The mix includes + evol-codealpaca, `Salesforce/xlam-function-calling-60k`, and + `SWE-bench/SWE-smith-trajectories`. +- **Reasoning models** add 12,288 samples from `open-r1/Mixture-of-Thoughts` + (4,096 each from code, math and science). +- **Repo composite recipe** (24,576 samples, max 16K): 4,096 evol-codealpaca, + 12,288 Mixture-of-Thoughts, 4,096 xLAM, and 4,096 from the SWE-smith `tool` + split. It uses `dataset:samples` syntax. +- **Chat formatting:** the default chat template is applied, and `` is + closed for hybrid models. +- **Domain matters a lot for generative tasks.** Calibrating on C4 made code + generation "collapse": many pruned models emitted incoherent text and scored + 0%. On multiple-choice tasks C4 was equal or slightly better, which is why MC + benchmarks hide the damage. Domain-matched calibration helps most on + fine-grained MoEs (Qwen3-30B, ERNIE). A combined code + writing + math set gave + the best all-round retention for REAP. +- Cerebras GLM-4.6 REAP cards list the same three sources (evol-codealpaca, + xLAM, SWE-smith-trajectories) [R3]. + +### 1.2 Later work (2025-2026) + +| Paper | Calibration finding relevant to us | +|---|---| +| **EAN / MC-Suite**, Jaiswal et al., arXiv 2504.05586 [P1] | Among 16 criteria, the expert activation norm (EAN) was strongest. Expert dropping mainly hurts instruction following, and k-shot examples or SFT recover it. | +| **EASY-EP**, NeurIPS 2025, arXiv 2504.06792 [P2] | On DeepSeek-R1 and V3: 25 in-domain demonstrations recover about 99% of critical experts, and 5 recover over 90%. Input plus the model's **own generated output** beat inputs only, outputs only, inputs plus gold answer, and pretraining/CommonCrawl text (CommonCrawl scored about 0 on AIME). Mixed-domain pruning averages normalized per-domain scores. | +| **MoEQuant**, ICML 2025, arXiv 2505.03804 [P3] | Standard calibration sets give long-tailed expert coverage, so under-routed experts get too few calibration tokens and large quantization error. Expert-balanced self-sampling, where the model generates its own data, gave more than 10 points on HumanEval for DeepSeekMoE-16B at 4-bit. | +| **AIMER**, arXiv 2603.18492 (Apr 2026) [P4] | For REAP on Qwen3-30B at 50%, results were "highly sensitive to calibration set size even with the same corpus" across 0.5M to 2.1M C4 tokens. Baselines used 4.2M C4 tokens. REAP calibration took 0.75 to 3 hours on 7-30B models. | +| **HOPE**, Tseng et al. (AWS), arXiv 2609.18916 (Sep 2026) [P5] | Second-order expert pruning (REAP is its first-order special case). Calibrated on all 111K Evol-CodeAlpaca prompts, or on full **SWE-bench Verified agent trajectories** (tool calls plus observations, truncated at 100K tokens) generated by the unpruned model. REAP needed about 8K prompts to reach 99% prune-set stability. Calibration-set Jaccard between the two sets was about 0.6. On SWE-Bench Pro, the agentic calibration scored higher for REAP **at every pruning rate**, and for HOPE at 50%. | +| **Robust calibration for domain-general pruning**, arXiv 2610.15578 (Oct 2026) [P6] | On gpt-oss-20b, a default of 64 sequences was enough, and scaling beyond small sizes gave little gain. English-only calibration transferred across 7 languages. Diverse pretraining-like data matched test-distribution data. Caveat: code was **not** evaluated and results are preliminary. | +| **Generic TB-Coverage**, arXiv 2607.01710 [P7] | Profiles experts per corpus and keeps a coverage budget per corpus before the global mask. Generic corpora only; it beat REAP on zero-shot MC. | +| **OptEP**, arXiv 2610.27123 [P8] | Reconstruction-optimal subset selection with task-domain calibration. Its gains are largest under aggressive pruning. | +| **Post-compression adjustment**, arXiv 2609.06076 [P9] | 3,000 C4 examples and 1 epoch of fine-tuning recover 37% of the pruning gap. This is relevant if we fine-tune after REAP. | +| **Self-calibration**, Williams et al., NAACL 2025, arXiv 2410.17170 [P10] | Model-generated calibration (BOS-conditioned, temperature 1.0) is competitive with real data and often better for pruning. Quantization was less sensitive to calibration source. Self-generated text is less diverse. | +| **Beware of Calibration Data**, Ji et al., ICLR 2025, arXiv 2410.17711 [P11] | At high sparsity, calibration data can matter more than the pruning method. Self-generated data with a perplexity filter helps. | +| **On the Impact of Calibration Data in PTQ and Pruning**, Williams & Aletras, ACL 2024, arXiv 2311.09755 [P12] | Calibration source shifts downstream results substantially, more so for pruning than for quantization. | + +### 1.3 Synthesis + +1. **Domain-specific calibration is mandatory for code and agentic use.** REAP + saw C4 collapse code output, and EASY-EP saw CommonCrawl collapse reasoning. + MC-only evals hide this, so held-out evals must be generative or logit-based + on the target domains. +2. **Agentic multi-turn trajectories should be the largest slice.** HOPE is the + only direct comparison. There, SWE-trajectory calibration beat + Evol-CodeAlpaca for REAP on SWE-Bench Pro at every rate tested, and the two + sets choose about 40% different experts. REAP's own large-model and agentic + checkpoints added SWE-smith tool trajectories. The evidence is suggestive, not + overwhelming: some lower-rate cells favored plain code for other criteria. So + keep a plain-code slice too. +3. **Self-generated data beats human data when it is "prompt + the model's own + response".** EASY-EP (on DeepSeek-V3 and R1, our closest architecture + family), MoEQuant, Self-calibration and Ji et al. all agree. The reason is + that the router sees the distribution it will see at inference. HOPE also + generated its agentic traces with the unpruned model. For agent traces, where + observations come from an environment, the practical proxy is traces from a + sibling model. `nvidia/Open-SWE-Traces` v1.1 contains about 50K OpenHands + trajectories from **DeepSeek-V4-Flash-0731**. +4. **Size.** Pruning decisions stabilize quickly (64 sequences in [P6], 25 shots + per domain in [P2]), but REAP with fine-grained experts shows size + sensitivity between 0.5M and 2.1M tokens [P4]. Production REAP runs use + 12K-25K samples at 16K max, which is tens of millions of tokens. With 384 + experts, the binding constraint is **per-expert coverage**, both for REAP + saliency means and especially for the 2-bit second moments. A 2-4M token + budget is fine for saliency ranking but is at the low end for per-expert + Hessians. +5. **Sequence length.** REAP uses 16,384 for large models with no packing, and + HOPE uses up to 100K for agent traces. Agentic routing depends on deep + context (tool outputs and file contents late in the trajectory), so + 2,048-token packing under-represents it. + +## 2. Candidate datasets per domain + +Legend: **V** = verified via the HF API on this date. Size is approximate. A +"messages" format means role-tagged multi-turn chat. Datasets marked as +benchmarks are avoided as calibration on purpose. + +### 2.1 Agentic coding: Python, TypeScript, Rust (SWE-style tool loops) + +| Repo id | V | License | Size | Format | Updated | Notes and contamination | +|---|---|---|---|---|---|---| +| `nvidia/Open-SWE-Traces` (config v1.1, split `openhands`, path `data/openhands/deepseek_v4_flash_0731/`) | V | CC-BY-4.0 | v1.0 151K, v1.1 263K, plus v1.2/v2.0 trajectories; DSV4-Flash subset ~3.8 GB parquet | `messages` (system/user/assistant/tool) + `tools` JSON + `language` + `resolved` + patches | 2026-10-28 | **Top pick.** Tasks come from SWE-rebench-V2 and Scale-SWE. v1.0 counts: Python 34K, TypeScript 27K, Rust 17K. Includes **DeepSeek-V4-Flash-0731** trajectories, the closest public data to self-generated. SWE-rebench-V2 excludes SWE-bench repos but overlaps SWE-bench Multilingual (12 tasks) and Multi-SWE-bench (83 tasks); exclude those IDs ([C1], agentleak). | +| `SWE-bench/SWE-smith-trajectories` | V | MIT | ~50K | `messages`, SWE-agent tool calls | 2025-07-19 | Used by REAP. Python-only synthetic bugs. Moderately used; good as a secondary harness for diversity. | +| `nebius/SWE-rebench-openhands-trajectories` | V | CC-BY-4.0 | 10K-100K | `messages` (OpenHands) | 2025-12-27 | Python. Overlaps SWE-rebench tasks in Open-SWE-Traces, so dedupe by `instance_id`. | +| `togethercomputer/CoderForge-Preview` | V | **no repo license** (per-row repo license only) | 258K trajectories, ~93 GB | OpenHands `messages` + `tools` | 2026-02-26 | See the section 4 vetting. Python-only, Qwen3-Coder-480B teacher. Better as eval. | +| `nvidia/SWE-Hero-openhands-trajectories` | V | CC-BY-4.0 | 10K-100K | OpenHands `messages` | 2026-05-08 | Alternative or overflow. | + +Plain code, real repository files and diffs (to cover code reading and writing +outside agent loops): + +| Repo id | V | License | Size | Format | Updated | Notes | +|---|---|---|---|---|---|---| +| `nebius/SWE-rebench-V2-PRs` | V | CC-BY-4.0 (per-repo SPDX column) | 126K real PRs, multi-language (Python, TS, Rust, Go...) | issue + PR description + gold `patch` + `test_patch` | 2026-08-18 | Real human commits, 2025-2026 heavy. Render as "issue -> diff". Same benchmark-overlap caveat as above. | +| `common-pile/stackv2_edu_filtered` | V | per-file open licenses (Common Pile) | 10M-100M files | single `text` + rich `metadata` (licenses, extension) | 2025-06-06 | Openly licensed, edu-filtered Stack v2 files. Filter by extension `.py/.rs/.ts/.tsx` and prefer `committer_date` >= 2023. | +| `HuggingFaceTB/stack-edu` (configs `Python`, `Rust`, `TypeScript`) | V | per-file (Stack v2 terms) | 100M+ files | ids/blobs; content fetched from Software Heritage S3 | 2025-03-20 | Use only if the Common Pile filter is too thin for Rust. Needs a SWH download step. | +| `SWE-bench/SWE-smith-rs`, `SWE-bench/SWE-smith-ts` | V | not stated in card | 1K-10K tasks each | task instances (no trajectories) | 2026-02 | Good **eval** seeds for Rust/TS agent rollouts by the teacher. | + +### 2.2 Bash / shell / terminal agents + +| Repo id | V | License | Size | Format | Updated | Notes | +|---|---|---|---|---|---|---| +| `nvidia/Nemotron-Terminal-Corpus` | V | CC-BY-4.0 | ~366K trajectories (`dataset_adapters` ~226K, `skill_based_*` ~140K) | terminus-2 `messages` (keystrokes / command JSON + terminal output) | 2026-02-27 | **Top pick.** Synthetic tasks from a primitive-skill taxonomy; the target benchmark is Terminal-Bench 2.0. Prefer `skill_based_*` for shell skills. | +| `open-thoughts/AgentTrove` | V | Apache-2.0 | 1.7M rows from 219 sources | terminus-2 `messages` + `original_source` + `original_teacher` + `reward` | 2026-05-07 | Filter `original_source` (e.g. `nl2bash`, `staqc`, `freelancer`) and `reward==1`. Teachers include GLM-4.6/4.7, Kimi K2, MiniMax, and GPT-5.1-nano/5-mini. Flag the OpenAI-generated rows (ToS) and drop them if that is a concern. | +| `open-thoughts/OpenThoughts-Agent-v1-SFT` | V | Apache-2.0 | 10K-100K | terminus-2 `messages` | 2026-01-27 | Smaller curated subset. Overlaps AgentTrove, so pick one. | +| `nvidia/Nemotron-RL-Agentic-Terminal-Pivot-v1` | V | CC-BY-4.0 | 10K-100K | RL prompts and environments | 2026-08-28 | Useful as **eval seeds**: roll out with the teacher. | + +### 2.3 General tool use (function calling, multi-turn tool loops) + +| Repo id | V | License | Size | Format | Updated | Notes | +|---|---|---|---|---|---|---| +| `Agent-Ark/Toucan-1.5M` | V | Apache-2.0 | 1.5M trajectories over 495 real MCP servers and 2,000+ tools | `messages` (Hermes tool format) + `subset` (`multi-turn`, `irrelevant`, ...) + LLM-judge quality | 2025-10-04 | **Top pick.** Real tool executions with parallel and sequential calls. Configs `Kimi-K2`, `Qwen3`, `OSS`, `SFT`. Use `multi-turn` with high judge scores. Targets BFCL/MCP-Universe. | +| `nvidia/Nemotron-Agentic-v1` | V | CC-BY-4.0 | ~? (multi-subset) | multi-turn user/agent/tool simulation | 2025-12-15 | Qwen3-235B / GPT-OSS generated and judge-filtered. Persona-seeded. | +| `nvidia/Nemotron-SFT-Agentic-v4` | V | OpenMDW-1.1 | multi-subset | `tool_calling` + `knowledge_agent` traces | 2026-10-15 | Newest; generated by Qwen3.5-397B and GLM-5.2. Check that OpenMDW is acceptable. | +| `Nanbeige/ToolMind` | V | Apache-2.0 | 100K-1M | multi-turn tool-use `messages` | 2026-01-09 | Alternative. | +| `Salesforce/APIGen-MT-5k` | V | **CC-BY-NC-4.0** | 5K | multi-turn (tau-bench-like) | 2025-10-10 | Non-commercial; **eval-only** at most. | +| `Salesforce/xlam-function-calling-60k` | V | CC-BY-4.0 (gated auto) | 60K single-turn | query + tools + answer | 2025-01-24 | Used by REAP. Single-turn and widely used, so give it a small share at most. | + +### 2.4 Networking + +Hugging Face has no strong, permissive networking corpus. Recommended sources: + +| Source | V | License | Size | Format | Notes | +|---|---|---|---|---|---| +| **IETF RFCs** from rfc-editor.org (`https://www.rfc-editor.org/rfc/rfcNNNN.txt`, index `rfc-index.xml`) | URL verified (HTTP 200); not HF | IETF Trust Legal Provisions (RFCs may be copied and redistributed in full) | ~9,900 RFCs | plain text | **Top pick.** Use recent RFCs (RFC 9000+, 2021-2026: QUIC, HTTP/3, TCP 9293, TLS, BGP, DNS) and hold out by RFC number. The bulk `RFC-all.tar.gz` path returned 404, so fetch per-RFC. | +| `common-pile/stackexchange_filtered` (filter `metadata.site` in `networkengineering.stackexchange.com`, `serverfault.com`, `unix.stackexchange.com`, `superuser.com`) | V | CC-BY-SA (per-doc `license`) | 31M docs total | question + answers + comments as one `text` | Real human operational Q&A. Filter `created` >= 2022 to reduce memorized content. | +| `jiebi/RFCAlign` | V | MIT | 1M-10M | synthetic Q/A from IETF mailing lists (Llama/Qwen) | Synthetic retrieval data. Low priority. | +| `nvidia/Nemotron-Terminal-Corpus` (network tasks inside) | V | CC-BY-4.0 | n/a | n/a | Adds hands-on `ip`/`ss`/`curl`/`dig` tool loops. | + +### 2.5 Cybersecurity + +| Repo id | V | License | Size | Format | Updated | Notes | +|---|---|---|---|---|---|---| +| `trendmicro-ailab/Primus-FineWeb` | V | ODC-BY (gated: accept terms) | 2.57B tokens | single `text` web documents | 2025-08-09 | **Top pick** for prose. Classifier-filtered FineWeb; snapshots through CC-MAIN-2024-10. | +| `CIRCL/vulnerability-scores` | V | CC-BY-4.0 | 837K records: CVE 5, GHSA, CSAF Red Hat/CISA/Cisco, PySec | title + description + CPEs + CVSS + patch URL | 2026-10-29 | Real advisories, continuously updated. Has **group-aware train/val/test splits** (no CVE/GHSA leakage), which makes it ideal for a disjoint held-out set. | +| `github/advisory-database` (GitHub, not HF) | V (GitHub API) | CC-BY-4.0 | ~300K+ OSV JSON | JSON advisories | pushed 2026-11-06 | Overlaps CIRCL's `github` source; use one. | +| `PentesterFlow/OffsecCyber` | V (page) | Apache-2.0 | 60.5K | `messages` (some with `tools`) | 2026 | Mixed: about 51% Qwen-synthetic, plus NVD, HackerOne and PortSwigger. Web-app focus. Secondary. | +| `common-pile/stackexchange_filtered` (`security.stackexchange.com`, `reverseengineering`) | V | CC-BY-SA | n/a | Q&A text | n/a | Real human Q&A. | + +Avoid: `Vyber07/cyber-security` (see section 4), `own2pwn-fr/ctfhoard-corpus` +(245 GB, no card, no license), and CTF benchmark sets (`penfever/*-harbor`, +`cvebench`, intercode-ctf), which are evals. + +### 2.6 Business documents + +| Repo id | V | License | Size | Format | Updated | Notes | +|---|---|---|---|---|---|---| +| `HuggingFaceFW/finepdfs` (config `eng_Latn`) | V | ODC-BY | 100M+ PDFs (all langs) | `text` + `url` + quality scores | 2026-04-03 | **Top pick.** Real PDFs (agendas, reports, policies, contracts). Filter by URL/domain and keywords (invoice, policy, minutes, proposal, SOW, quarterly) and by `ocr_quality_scores`. Hold out by URL domain. | +| SEC EDGAR full text (sec.gov, public domain), or `juand-r/ai-sec-10k-filings-since-2020` | V | HF card has no license; EDGAR itself is public domain | 1B-10B rows | raw SEC SGML text | 2026-05-31 | Take MD&A and risk-factor sections from 2023-2026 10-K/10-Q filings. Strip SGML. | +| `corbt/enron-emails` | V | no license field (Enron corpus released via FERC) | ~500K emails | structured email fields | 2025-04-07 | Real business correspondence, but dated (2001). Small share. | +| `common-pile/regulations_filtered` / `common-pile/usgpo_filtered` | V | public domain / open | 100K-10M | `text` | 2025-06-06 | Government and regulatory documents. Optional. | +| `nvidia/Nemotron-SFT-Finance-Document-Grounded-Reasoning-v1` | V | OpenMDW-1.1 | n/a | document-grounded `messages` | 2026-10-14 | Synthetic. Optional for document-QA style. | + +### 2.7 English (general instruction and prose) + +| Repo id | V | License | Size | Format | Updated | Notes | +|---|---|---|---|---|---|---| +| `allenai/WildChat-4.8M` | V | ODC-BY | 4.8M real user conversations | `conversation` messages + `language` | 2025-08-11 | **Top pick** for real prompts. Filter `language=="English"`, drop toxic and redacted rows, **use the user prompts and regenerate responses with V4.1-Flash** (EASY-EP style). Not a benchmark. | +| `HuggingFaceFW/finewiki` (config `en`) | V | CC-BY-SA-4.0 / GFDL | 10M-100M | `text` | 2025-10-22 | Clean 2025 Wikipedia. | +| `allenai/dolma3.5_pool` | V | ODC-BY | huge | `text` | 2026-07-10 | Pretraining-like diversity ([P6] finds this is enough for domain-general retention). | +| `nvidia/Nemotron-SFT-Instruction-Following-Chat-v3` | V | mixed (ODC-BY / CC-BY / Apache / other) | 100K-1M | `messages` | 2026-06-04 | Synthetic. Alternative prompt source. | + +## 3. Recommended mix + +### 3.1 Principles + +- **Format.** Render everything with V4.1-Flash's own chat template and tool-call + format. Convert Hermes/OpenHands/terminus tool calls into the model's native + tool schema; otherwise the router sees foreign control tokens. Match the + thinking mode to deployment. If both modes are used, take about 60% thinking + and 40% non-thinking. +- **Self-generation.** For prompt-type domains (English, tool use, cyber Q&A, + networking Q&A, business-doc QA), keep the human or source prompts and + **regenerate the assistant turns with V4.1-Flash itself**. You already run the + teacher forward pass, so the cost is generation time only. For agentic traces, + prefer the DeepSeek-V4-Flash-0731 trajectories in Open-SWE-Traces. Keep raw + documents (RFCs, code files, advisories, filings) human-written: in agent use + they arrive as tool observations the router must handle. +- **Sequence length.** Use **max 16,384** tokens, with no packing across + conversations, as REAP does for its large models. For agent trajectories + longer than 16K, take message-aligned windows that keep the system prompt and + tools header plus a random contiguous span of later turns, not just the first + 16K. Add a **long-context slice of about 10% of tokens at 32K-64K** (the + model's native 64K pre-YaRN) from Open-SWE-Traces, so that long-horizon routing + is represented, as HOPE did with traces up to 100K. Pack short documents + (advisories, Q&A) to 4K-8K with EOS separators. +- **Coverage gate.** After the forward pass, log routed-token counts for each + (layer, expert). For REAP, require at least 1K routed tokens for every expert + you might keep. For the 2-bit second moments, require at least about 2 x d_in + (about 10K tokens for gate/up, about 5K for down) for every **kept** expert. + If any kept expert falls short, top up with more agentic traces, or with + MoEQuant-style self-sampling steered toward under-covered experts, before + computing Hessians. +- **Seeds.** As REAP did, check the stability of the prune set across 2 + calibration seeds or subsamples (Jaccard). HOPE found about 0.6 between + different domain sets. + +### 3.2 Calibration split (~3.2M tokens) + +| Domain | Tokens | Sources (in priority order) | +|---|---|---| +| Agentic coding (Python, TS, Rust) | 1,000K | Open-SWE-Traces v1.1 DeepSeek-V4-Flash-0731 OpenHands: Python 400K / TS 300K / Rust 300K. Fill Rust/TS gaps from v1.0/v1.2 or SWE-smith-trajectories. Include the 32K-64K long windows here. | +| Plain code + real diffs | 400K | SWE-rebench-V2-PRs (issue -> patch) 200K, plus stackv2_edu_filtered `.py/.rs/.ts` files 200K, split evenly across the three languages. | +| Bash / terminal agent | 400K | Nemotron-Terminal-Corpus `skill_based_*` 250K, plus AgentTrove (`nl2bash`/`staqc`, reward 1, non-OpenAI teachers) 150K. | +| General tool use | 450K | Toucan-1.5M `multi-turn` 250K, plus Nemotron-SFT-Agentic-v4 / Nemotron-Agentic-v1 150K, plus xLAM single-turn 50K. | +| Cybersecurity | 300K | Primus-FineWeb 150K, CIRCL vulnerability-scores (train split) 75K, and security.stackexchange 75K (prompts with self-generated answers). | +| Networking | 250K | Recent RFCs 150K, plus networkengineering/serverfault/unix StackExchange 100K. | +| Business documents | 250K | finepdfs eng_Latn business-filtered 150K, SEC 10-K/10-Q sections 2023-2026 60K, Enron 40K. | +| English general | 150K | WildChat English prompts with V4.1-Flash responses 100K, plus finewiki en 50K. | +| **Total** | **~3.2M** | About 60% agentic or tool-loop tokens, which matches the target use. | + +Rationale: + +- Agentic and tool-loop data gets the largest share because HOPE showed + agentic calibration helps REAP on SWE-Bench Pro at every rate, and REAP's + agentic checkpoints use it. +- Every target domain gets at least 150K tokens, because REAP and EASY-EP show + that a domain absent from calibration loses its experts. EASY-EP's 25-shot + result suggests 150K tokens is far above the minimum for ranking. +- The total stays within 2-4M because [P6] and [P2] show ranking saturates + early. AIMER's 0.5M-2.1M sensitivity argues against going below about 2M. +- Before quantization, consider growing the agentic slice to about 8-10M tokens + for Hessian accuracy alone (QTIP/QuIP# scale). That data need not enter the + REAP saliency pass. + +### 3.3 Held-out eval split (disjoint), for teacher top-k logits + +Use about 40-60K tokens per domain (~450K total, 100-300 sequences per domain), +with the same template and lengths. Make each split disjoint as follows: + +- **Agentic coding:** disjoint **by repository** from calibration, not just by + instance. Use Open-SWE-Traces repos held out, plus CoderForge `filtered_reward1` + Python repos not in calibration. CoderForge's different teacher and harness + also give a robustness check. Seed Rust/TS agent rollouts from SWE-smith-rs/ts + tasks generated by the teacher. +- **Bash:** a different `skill` taxonomy branch, or a different AgentTrove + `original_source`. +- **Tool use:** held-out MCP servers in Toucan (split by `target_tools` server). +- **Cyber:** CIRCL's own `test` split, which is group-split by CVE/alias. +- **Networking:** a held-out set of RFC numbers, plus StackExchange questions + created after the calibration cutoff date. +- **Business:** held-out URL domains in finepdfs, plus different SEC filers. +- **English:** held-out WildChat conversation ids. + +Also report per-domain KL against the teacher's top-k, not only accuracy. The +REAP finding that MC accuracy hides generative collapse applies here. + +## 4. Vetting of five candidates Jack suggested + +Revision SHAs are pinned as of 2026-11-05/06. + +| Field | `togethercomputer/CoderForge-Preview` | `nvidia/Open-SWE-Traces` | `Vyber07/cyber-security` | `MoreThought/Fable-5.1-Max-Reasoning-Filtered-10000x` | `Axion1-0/axion-training-data` | +|---|---|---|---|---|---| +| Pinned SHA | `060fca96cf723b2ebab3181e9e59fafd273df3cb` | `4fa4c2ae9afc6b397bc50391d43450c79993c138` | `bc7a780571cb67d482250649ab573908b9557873` | `4904ca0c09044f447aff16e4f66b6f2f53c9f702` | `668831c96e7c3d974bc5c414329d508277f23663` | +| License / gated | **No repo-level license** in the card or tags. Each row has a `license` field holding the source repo's SPDX id (e.g. Apache-2.0). Not gated. | CC-BY-4.0, not gated, "commercial/non-commercial use" | Card says Apache-2.0. **Gated (auto)**: files need you to share HF account info. We did not accept the gate, so contents are unverified. | Card says Apache-2.0, not gated. The uploader's license cannot override the generating model's terms (see ToS). | Apache-2.0, not gated | +| Size / format | 258K trajectories (155K with reward 1), ~93 GB raw. Avg **41.4K tokens** per trajectory (P99 91K). `messages` JSON string + `tools` JSON + `reward` + `finish_reason` + Docker `image`. Also a Qwen-tokenized config. | v1.0 151K, v1.1 263K, v1.2/v2.0 more. Parquet with `messages` (system/user/assistant/tool), `tools`, `language`, `resolved`, gold and model patches. | Card claims 16 GB / 507 files in one place and ~95 GB in another; the Hub tag says <1K rows. The top level of the tree shows 1,000 entries, ~13 GB. ShareGPT, instruction/IO and "metadata-only" JSON schemas mixed; card says only 394/507 files are valid JSON. Includes a stray `.openclaude__settings.local.json`. | 10K traces (5K `train`, 10K `full`), 1.97 GB, claimed ~500M tokens (~50K per trace). `messages` + `tools` (`bash`, `read_file`, `write_file`, `search`, `list_dir`). | 18,363 SFT trajectories (stage1 10,819, stage2 7,544) + 484 RL tasks, ~533 MB. `messages`; tool responses are in `user` turns. | +| Multi-turn / tool calls | **Yes, agentic.** OpenHands v0.52.1 tools: `execute_bash`, `str_replace_editor` (file edits), `think`, `finish`. Up to 100 steps in Docker; real executions. | **Yes, agentic.** OpenHands (`execute_bash`, editor), SWE-agent, mini-swe-agent, Claude-Code and OpenCode harnesses. Real executions in SWE-rebench-V2 / Scale-SWE containers. | Some ShareGPT multi-turn chat, per the card. No tool calls described. | Superficially yes (avg 40-100 messages), but **1,370 of 1,372 tool-call argument strings in our sample are placeholders** (`{"path": ""}`, `{"command": ""}`). Tool names look remapped: a `search` call returns "Command running in background..." output. | Yes: search, web visit and Python tool loops (stage2); Python only (stage1). | +| Language | English (prompts are GitHub issues; a few non-English issues are possible) | English, with the same caveat. Run langid on `problem_statement` and drop non-English. | Card says English (unverified) | English | English | +| Jack's domains covered | Python agentic coding and bash (inside the loop). No Rust or TS (blog and prompt template are Python-repo only). | **Python, TypeScript, Rust** agentic coding (plus Go/JS/Java/PHP), bash inside the loop | Cybersecurity, possibly networking (claimed) | Claimed coding and bash, but the visible prompts are the uploader's own data-pipeline "seed authoring" tasks | General tool use, but in scientific research (biology, chemistry, KEGG/UniProt). Not coding. | +| Provenance | Model-generated by **Qwen3-Coder-480B** (open weights, Apache-2.0). No commercial-model ToS issue. Tasks: R2E-Gym (4.2K), SWE-Smith (37.2K), SWE-Rebench (9.8K), 1,655 repos. | Model-generated by **MiniMax-M2.5, Qwen3.5-122B, Qwen3.6-27B, DeepSeek-V4-Flash-0731, Qwen3.8-27B, Qwen3.8-Flash-Next**, all open-weight. No commercial-model ToS issue. The "claude-code" split uses the Claude Code harness driven by Qwen3.8 models, not Claude outputs. | Undocumented: no statement of human vs model generation, no model named, sources (HackerOne, HTB, TryHackMe, PortSwigger) with no licensing provenance. | **Distilled from "Fable 5.1", an Anthropic Claude-family model** ("Claude Fable 5" appears in Axion's comparison table). Traces were captured from **Claude Code sessions**: paths `~/.claude/projects/...` and `/tmp/claude-1000/...` appear in tool outputs. | Model-generated by an **unnamed** LLM solver. GPT-5.5 was used only to verify reference answers. The student is Qwen3.5-122B. If the solver was a commercial API, the ToS status is unknown. | +| Contamination risk | Low for SWE-bench Verified: the authors excluded tasks matching SWE-bench Verified by repo + base commit or problem statement. Not checked against SWE-bench Multilingual or Multi-SWE-bench. SWE-Smith repos are synthetic-bug repos. | SWE-rebench-V2 excludes SWE-bench repos but overlaps SWE-bench Multilingual (12 tasks, 15 repos) and Multi-SWE-bench (83 tasks, 16 repos) per agentleak; exclude those ids. Scale-SWE (Python, CC-BY-4.0) reports SWE-bench Verified results; check its repos against the 12 Verified repos. | Unknown; cannot inspect. PortSwigger/HTB lab writeups may mirror CTF benchmarks. | Unknown task sources. Not relevant given the reject. | No explicit SFT decontamination statement. RL tasks exclude held-out eval overlap. Science QA could overlap GPQA/HLE-style evals (not our targets). | +| Dedup / overlap with other picks | **High overlap** with Open-SWE-Traces (both draw on SWE-rebench) and with SWE-smith-trajectories (SWE-Smith tasks). If CoderForge supplies eval, the calibration split must exclude its repos (dedupe on repo + base commit, or `instance_id` parsed from `trajectory_id`). | Is the primary pick. Overlaps nebius/SWE-rebench-openhands-trajectories, SWE-rebench-V2-PRs and CoderForge at the task and repo level. | Likely overlaps PentesterFlow/OffsecCyber and NVD/CVE sources (unverifiable) | n/a | None with the other picks (different domain) | +| Recommendation | **Eval-only** (primary): a held-out Python agentic split from a different teacher and harness, measuring generalization beyond the DeepSeek-distribution calibration. Use `filtered_reward1` with repos disjoint from calibration, cut to message-aligned 16K-32K windows. Calibration use is acceptable only as overflow, after the repo-level license gap is accepted. | **Calibration: top pick** for Python/TS/Rust agentic coding. Use the v1.1 DeepSeek-V4-Flash-0731 OpenHands subset first, as the nearest-to-self data. Pin the SHA, filter by `language`, prefer `resolved==1` (keep some failures for realism), and exclude agentleak overlap ids. | **Reject.** Gated content we have not inspected; undocumented provenance and licensing of platform material; self-inconsistent card; an explicit "no restrictions" offensive scope with no safety review; tooling artifacts in the repo. Use Primus-FineWeb + CIRCL instead. | **Reject.** (1) Commercial-model **ToS problem**: Anthropic's terms restrict using outputs to build competing models, and an uploader's Apache-2.0 label cannot grant rights it does not hold. Calibrating and distilling a pruned model on it is at least a gray area. (2) Data defects: tool arguments replaced by placeholders, so the tokens are unrealistic for routing statistics. (3) Heavy duplication: only 5 distinct first-user prompts across the first 23 traces, and several near-identical traces, despite the dedup claim. (4) Prompts are meta "seed-authoring" tasks for the uploader's pipeline. | **Reject for calibration** (off-domain science research; non-native tool-response format in `user` turns; unnamed teacher). Optional **eval-only** as a small out-of-domain tool-loop probe, to watch for collateral damage to general tool use. Low priority. | + +Notes: + +- The Vyber07 card text above was read from the public page only; the files were + not downloaded because that needs the gate accepted. +- Fable statistics come from the first ~3 MB (23 traces) of `5k/train.jsonl`, + fetched with an HTTP range request and parsed as data. + +## 5. Sources + +- [R1] REAP paper: https://arxiv.org/abs/2510.13999 (HTML https://arxiv.org/html/2510.13999) +- [R2] REAP code and composite calibration recipe: https://github.com/CerebrasResearch/reap +- [R3] Cerebras REAP card: https://huggingface.co/cerebras/GLM-4.6-REAP-252B-A12B-FP8 +- [P1] EAN / MC-Suite: https://arxiv.org/abs/2504.05586 +- [P2] EASY-EP: https://arxiv.org/abs/2504.06792 (NeurIPS 2025) +- [P3] MoEQuant: https://arxiv.org/abs/2505.03804 +- [P4] AIMER: https://arxiv.org/html/2603.18492v2 +- [P5] HOPE: https://arxiv.org/abs/2609.18916 +- [P6] Robust calibration for domain-general MoE pruning: https://arxiv.org/abs/2610.15578 +- [P7] Generic TB-Coverage: https://arxiv.org/abs/2607.01710 +- [P8] OptEP: https://arxiv.org/abs/2610.27123 +- [P9] Post-compression adjustment: https://arxiv.org/abs/2609.06076 +- [P10] Self-calibration: https://arxiv.org/abs/2410.17170 +- [P11] Beware of Calibration Data: https://arxiv.org/abs/2410.17711 +- [P12] Calibration data in PTQ and pruning: https://arxiv.org/abs/2311.09755 +- [Q1] QuIP#: https://arxiv.org/abs/2402.04396 +- [Q2] QTIP: https://arxiv.org/abs/2406.11235 +- [C1] SWE-rebench-V2 benchmark overlap (agentleak): https://github.com/raimondasl/agentleak +- CoderForge blog: https://www.together.ai/blog/coderforge-preview +- Axion 1.0 paper: https://huggingface.co/papers/2611.03233 +- SWE-rebench-V2: https://arxiv.org/abs/2602.23866 +- Terminal-Corpus paper: https://arxiv.org/abs/2602.21193 +- Toucan: https://arxiv.org/abs/2510.01179 +- Primus: https://arxiv.org/abs/2502.11191 +- Dataset pages: `https://huggingface.co/datasets/` for every id above. diff --git a/docs/research/2607.05147v1.pdf b/docs/research/2607.05147v1.pdf new file mode 100644 index 0000000..2bea013 Binary files /dev/null and b/docs/research/2607.05147v1.pdf differ 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/scripts/runpod/v41_reap/build_pack.py b/scripts/runpod/v41_reap/build_pack.py new file mode 100644 index 0000000..82e2a39 --- /dev/null +++ b/scripts/runpod/v41_reap/build_pack.py @@ -0,0 +1,388 @@ +"""Build the V4.1 REAP calibration + held-out eval pack from public datasets. + +Mix and rationale: docs/research/2026-11-05-v41-reap-calibration-data.md +(section 3). English only, by policy: any record containing CJK, Hangul or +Kana characters, or more than 2% non-Latin letters, is dropped. + +Every chat source is rendered with V4.1's own ``encoding/encoding.py`` (tools +on the system message, DSML tool calls) and tokenized with the V4.1 +tokenizer. Raw documents go in as BOS + text. Records are capped at +``--max-tokens`` (hard truncation). Short documents are packed to about +``--pack-tokens`` with EOS separators. + +Split: each record's group key (repo, task, RFC number, URL host, CVE id, +...) is hashed, and one bucket in ten feeds the eval split until that +domain's eval budget is met, so calibration and eval never share a group. +Calibration records are grouped into sessions of up to ``--session-tokens``; +each eval record is its own session (one logits stream each). + +Run on the pod (a CPU pod is enough), not the Mac:: + + python scripts/runpod/v41_reap/build_pack.py --model-dir /workspace/v41-hf \\ + --out /workspace/pack + +Output: ``v41-pack.jsonl`` (collect.py's input) and ``manifest.json`` +(resolved dataset revisions, rows kept/dropped, tokens per domain and split). +""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import re +import os +import sys +import urllib.request +from collections import Counter, defaultdict +from collections.abc import Callable, Iterator +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +K = 1000 + +# fmt: off +#: name, dataset, load kwargs, adapter, domain (str, or a callable on the row), group-key field, {domain: (calib, eval) tokens} +SOURCES: list[dict[str, Any]] = [ + dict(name="open_swe_traces_dsv4", dataset="nvidia/Open-SWE-Traces", + load=dict(data_files="data/openhands/deepseek_v4_flash_0731/*/*.parquet", split="train"), + adapter="openai", domain=lambda r: f"agentic-{r['language']}", group="repo", + budget={"agentic-python": (400 * K, 50 * K), "agentic-typescript": (300 * K, 40 * K), "agentic-rust": (300 * K, 40 * K)}), + dict(name="swe_rebench_v2_prs", dataset="nebius/SWE-rebench-V2-PRs", load=dict(split="train"), + adapter="issue_patch", domain="code-diffs", group="repo", budget={"code-diffs": (400 * K, 50 * K)}), + # ponytail: no plain-file source. stackv2_edu_filtered's shards are grouped by language with no index, + # so finding .py/.rs/.ts means scanning most of it; the agent traces already show files through tool output. + dict(name="nemotron_terminal", dataset="nvidia/Nemotron-Terminal-Corpus", load=dict(name="skill_based_mixed", split="train"), + adapter="terminus", domain="bash", group="task", budget={"bash": (400 * K, 50 * K)}), + dict(name="toucan", dataset="Agent-Ark/Toucan-1.5M", load=dict(name="SFT", split="train"), + adapter="toucan", domain="tool-use", group="tools", budget={"tool-use": (450 * K, 60 * K)}), + dict(name="circl_vulns", dataset="CIRCL/vulnerability-scores", load=dict(split="train"), + adapter="advisory", domain="cyber", group="id", pack=True, budget={"cyber": (150 * K, 25 * K)}), + dict(name="stackexchange", dataset="common-pile/stackexchange_filtered", load=dict(split="train"), + adapter="stackexchange", domain=lambda r: _se_domain(r), group="url", pack=True, + budget={"cyber-qa": (150 * K, 25 * K), "networking-qa": (100 * K, 20 * K)}), + dict(name="rfcs", dataset="rfc-editor.org", load=dict(first=9000, last=9800), + adapter="rfc", domain="networking", group="rfc", budget={"networking": (150 * K, 25 * K)}), + dict(name="finepdfs_business", dataset="HuggingFaceFW/finepdfs", load=dict(name="eng_Latn", split="train"), + adapter="business_pdf", domain="business", group="url", group_host=True, budget={"business": (210 * K, 30 * K)}), + dict(name="enron", dataset="corbt/enron-emails", load=dict(split="train"), + adapter="email", domain="business-email", group="from", pack=True, budget={"business-email": (40 * K, 10 * K)}), + dict(name="wildchat_prompts", dataset="allenai/WildChat-4.8M", load=dict(split="train"), + adapter="wildchat", domain="english-chat", group="conversation_hash", budget={"english-chat": (50 * K, 15 * K)}), + dict(name="finewiki", dataset="HuggingFaceFW/finewiki", load=dict(name="en", split="train"), + adapter="text", domain="english-prose", group="id", budget={"english-prose": (100 * K, 15 * K)}), +] +# fmt: on + +SE_SITES = { + "cyber-qa": ("security.stackexchange.com", "reverseengineering.stackexchange.com"), + "networking-qa": ("networkengineering.stackexchange.com", "serverfault.com", "unix.stackexchange.com"), +} +BUSINESS_WORDS = re.compile( + r"\b(invoice|purchase order|statement of work|quarterly report|annual report|board minutes|meeting minutes|" + r"business plan|proposal|policy|procedure|memorandum|agreement|contract|budget|financial statements?)\b", + re.IGNORECASE, +) +NON_ENGLISH_SCRIPT = re.compile(r"[぀-ヿ㐀-䶿一-鿿가-힯豈-﫿ｦ-ﾟ]") + + +def is_english(text: str) -> bool: + """No CJK / Hangul / Kana at all, and at most 2% of letters outside Latin.""" + + if NON_ENGLISH_SCRIPT.search(text): + return False + letters = [c for c in text if c.isalpha()] + foreign = sum(1 for c in letters if ord(c) > 0x024F) + return not letters or foreign <= 0.02 * len(letters) + + +def _se_domain(row: dict[str, Any]) -> str | None: + site = (row.get("metadata") or {}).get("site") + return next((domain for domain, sites in SE_SITES.items() if site in sites), None) + + +# --------------------------------------------------------------------------- +# Adapters: row -> ("chat", messages) | ("doc", text) | None +# --------------------------------------------------------------------------- + + +def _json(value: Any) -> Any: + return json.loads(value) if isinstance(value, str) else value + + +def _split_think(content: str) -> tuple[str, str]: + match = re.match(r"\s*(.*?)\s*(.*)", content or "", re.DOTALL) + return (match.group(1).strip(), match.group(2)) if match else ("", content or "") + + +def adapt_openai(row: dict[str, Any]) -> tuple[str, Any] | None: + messages = [] + for m in _json(row["messages"]): + msg = {"role": m["role"], "content": m.get("content") or ""} + if m.get("reasoning_content"): + msg["reasoning_content"] = m["reasoning_content"] + if m.get("tool_calls"): + msg["tool_calls"] = _json(m["tool_calls"]) + if m["role"] == "tool" and m.get("tool_call_id"): + msg["tool_call_id"] = m["tool_call_id"] + messages.append(msg) + tools = [_json(t) for t in _json(row.get("tools") or [])] + if tools and messages and messages[0]["role"] == "system": + messages[0]["tools"] = tools + elif tools: + messages.insert(0, {"role": "system", "content": "", "tools": tools}) + return "chat", messages + + +def adapt_terminus(row: dict[str, Any]) -> tuple[str, Any] | None: + messages = [] + for m in row["conversations"]: + if m["role"] == "assistant": + reasoning, content = _split_think(m["content"]) + messages.append({"role": "assistant", "content": content, **({"reasoning_content": reasoning} if reasoning else {})}) + else: + messages.append({"role": m["role"], "content": m["content"]}) + return "chat", messages + + +def adapt_toucan(row: dict[str, Any]) -> tuple[str, Any] | None: + if "multi-turn" not in (row.get("subset_name") or ""): + return None + messages: list[dict[str, Any]] = [{"role": "system", "content": "", "tools": _json(row["tools"])}] + calls = 0 + for m in _json(row["messages"]): + role, content = m["role"], m.get("content") or "" + if role == "tool_call": + call = ast.literal_eval(content) if isinstance(content, str) else content + calls += 1 + if messages[-1]["role"] != "assistant": + messages.append({"role": "assistant", "content": ""}) + messages[-1].setdefault("tool_calls", []).append( + {"id": f"call_{calls}", "type": "function", "function": {"name": call["name"], "arguments": call["arguments"]}} + ) + elif role == "tool_response": + messages.append({"role": "tool", "content": content, "tool_call_id": f"call_{calls}"}) + else: + messages.append({"role": role, "content": content}) + return "chat", messages + + +def adapt_issue_patch(row: dict[str, Any]) -> tuple[str, Any] | None: + issue = row.get("problem_statement") or row.get("pr_description") or "" + if not issue or not row.get("patch"): + return None + return "chat", [ + {"role": "user", "content": f"Repository: {row['repo']}\n\n{issue}\n\nWrite the patch."}, + {"role": "assistant", "content": row["patch"]}, + ] + + +def adapt_advisory(row: dict[str, Any]) -> tuple[str, Any] | None: + if not row.get("description"): + return None + head = row["id"] + (f": {row['title']}" if row.get("title") else "") + return "doc", f"{head}\n{row['description']}" + + +def adapt_stackexchange(row: dict[str, Any]) -> tuple[str, Any] | None: + return "doc", row["text"] + + +def adapt_business_pdf(row: dict[str, Any]) -> tuple[str, Any] | None: + ocr = row.get("ocr_quality_scores") or [0] + if min(ocr) < 2.0 or not BUSINESS_WORDS.search(row["text"][:4000]): + return None + return "doc", row["text"] + + +def adapt_email(row: dict[str, Any]) -> tuple[str, Any] | None: + if len(row.get("body") or "") < 200: + return None + return "doc", f"From: {row['from']}\nTo: {', '.join(row.get('to') or [])}\nSubject: {row.get('subject', '')}\n\n{row['body']}" + + +def adapt_wildchat(row: dict[str, Any]) -> tuple[str, Any] | None: + # Human prompts only: the responses are OpenAI-model output, so they stay out. + if row.get("language") != "English" or row.get("toxic") or row.get("redacted"): + return None + first = next((m["content"] for m in row["conversation"] if m["role"] == "user"), "") + return ("chat", [{"role": "user", "content": first}]) if len(first) > 40 else None + + +def adapt_text(row: dict[str, Any]) -> tuple[str, Any] | None: + return "doc", row["text"] + + +ADAPTERS: dict[str, Callable[[dict[str, Any]], tuple[str, Any] | None]] = { + "openai": adapt_openai, "terminus": adapt_terminus, "toucan": adapt_toucan, "issue_patch": adapt_issue_patch, + "advisory": adapt_advisory, "stackexchange": adapt_stackexchange, + "business_pdf": adapt_business_pdf, "email": adapt_email, "wildchat": adapt_wildchat, "text": adapt_text, +} # fmt: skip + + +def rfc_rows(first: int, last: int) -> Iterator[dict[str, Any]]: + for number in range(first, last + 1): + try: + with urllib.request.urlopen(f"https://www.rfc-editor.org/rfc/rfc{number}.txt", timeout=30) as response: + yield {"rfc": number, "text": response.read().decode("utf-8", "replace")} + except OSError: + continue # unassigned or not-yet-published numbers + + +# --------------------------------------------------------------------------- +# Rendering, splitting, packing +# --------------------------------------------------------------------------- + + +class Renderer: + def __init__(self, model_dir: Path): + from tokenizers import Tokenizer + + sys.path.insert(0, str(model_dir / "encoding")) + import encoding # V4.1's own chat format (the pinned revision's encoding/encoding.py) + + self.encoding = encoding + self.tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json")) + self.bos = self.tokenizer.token_to_id(encoding.bos_token) + self.eos = self.tokenizer.token_to_id(encoding.eos_token) + + def __call__(self, kind: str, payload: Any) -> list[int]: + if kind == "chat": + thinking = any(m.get("reasoning_content") for m in payload) + text = self.encoding.encode_messages(payload, thinking_mode="thinking" if thinking else "chat") + return self.tokenizer.encode(text, add_special_tokens=False).ids + return [self.bos, *self.tokenizer.encode(payload, add_special_tokens=False).ids] + + +def is_eval_group(key: str) -> bool: + return int(hashlib.sha256(key.encode()).hexdigest(), 16) % 10 == 0 + + +def iter_rows(source: dict[str, Any], revision: str | None) -> Iterator[dict[str, Any]]: + if source["adapter"] == "rfc": + yield from rfc_rows(**source["load"]) + return + from datasets import load_dataset + + load = dict(source["load"]) + yield from load_dataset(source["dataset"], streaming=True, revision=revision, **load) + + +def build(sources: list[dict[str, Any]], render: Callable[[str, Any], list[int]], *, eos: int, max_tokens: int, + pack_tokens: int, session_tokens: int, revisions: dict[str, str | None], rows: Callable = iter_rows) -> tuple[list[dict], dict]: # fmt: skip + records: list[dict[str, Any]] = [] + report: dict[str, Any] = {"sources": {}, "tokens": defaultdict(lambda: Counter())} + for source in sources: + adapter = ADAPTERS.get(source["adapter"]) + need = {d: {"calibration": c, "eval": e} for d, (c, e) in source["budget"].items()} + got = report["tokens"] + stats = Counter() + buffers: dict[tuple[str, str], list[int]] = {} + groups: dict[tuple[str, str], list[str]] = defaultdict(list) + + def emit(domain: str, split: str, ids: list[int], keys: list[str]) -> None: + records.append({"id": f"{source['name']}-{len(records)}", "split": split, "domain": domain, + "source": source["name"], "groups": keys, "token_ids": ids[:max_tokens]}) # fmt: skip + got[domain][split] += min(len(ids), max_tokens) + + for row in rows(source, revisions.get(source["name"])): + if all(got[d][s] >= n for d, v in need.items() for s, n in v.items()): + break + stats["rows"] += 1 + domain = source["domain"](row) if callable(source["domain"]) else source["domain"] + if domain not in need: + stats["off_domain"] += 1 + continue + converted = adapt_text({"text": row["text"]}) if source["adapter"] == "rfc" else adapter(row) + if converted is None: + stats["filtered"] += 1 + continue + if not is_english(json.dumps(converted[1], ensure_ascii=False)): + stats["non_english"] += 1 + continue + key = str(row.get(source["group"])) + if source.get("group_host"): # hold out whole sites, not single URLs + key = urlparse(key).netloc or key + split = "eval" if is_eval_group(key) else "calibration" + if got[domain][split] >= need[domain][split]: + stats["over_budget"] += 1 + continue + ids = render(*converted) + stats["kept"] += 1 + if not source.get("pack"): + emit(domain, split, ids, [key]) + continue + buf = buffers.setdefault((domain, split), []) + buf.extend(ids + [eos]) + groups[(domain, split)].append(key) + if len(buf) >= pack_tokens: + emit(domain, split, buf, groups.pop((domain, split))) + buffers[(domain, split)] = [] + for (domain, split), buf in buffers.items(): + if buf: + emit(domain, split, buf, groups[(domain, split)]) + report["sources"][source["name"]] = {"dataset": source["dataset"], "revision": revisions.get(source["name"]), **stats} + + # Sessions: each eval record alone; calibration interleaved across domains, cut at session_tokens. + calibration = sorted((r for r in records if r["split"] == "calibration"), key=lambda r: hashlib.sha256(r["id"].encode()).hexdigest()) + session, size = 0, 0 + for record in calibration: + if size and size + len(record["token_ids"]) > session_tokens: + session, size = session + 1, 0 + record["session"] = f"cal-{session:04d}" + record["target_spans"] = [] + size += len(record["token_ids"]) + for i, record in enumerate(r for r in records if r["split"] == "eval"): + record["session"] = f"eval-{i:05d}" + record["target_spans"] = [[1, len(record["token_ids"])]] + report["tokens"] = {d: dict(c) for d, c in sorted(report["tokens"].items())} + return calibration + [r for r in records if r["split"] == "eval"], report + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-dir", required=True, type=Path, help="V4.1 HF snapshot (tokenizer.json, encoding/)") + parser.add_argument("--out", required=True, type=Path) + parser.add_argument("--max-tokens", type=int, default=16384) + parser.add_argument("--pack-tokens", type=int, default=6144) + parser.add_argument("--session-tokens", type=int, default=131072) + parser.add_argument("--only", help="comma-separated source names (debugging)") + args = parser.parse_args(argv) + + from huggingface_hub import HfApi + + sources = [s for s in SOURCES if not args.only or s["name"] in args.only.split(",")] + api = HfApi() + revisions = {s["name"]: (api.dataset_info(s["dataset"]).sha if "/" in s["dataset"] else None) for s in sources} + render = Renderer(args.model_dir) + records, report = build(sources, render, eos=render.eos, max_tokens=args.max_tokens, pack_tokens=args.pack_tokens, + session_tokens=args.session_tokens, revisions=revisions) # fmt: skip + args.out.mkdir(parents=True, exist_ok=True) + pack = args.out / "v41-pack.jsonl" + with pack.open("w") as handle: + for record in records: + handle.write(json.dumps(record) + "\n") + report.update( + record_type="keep_v41_reap_pack_v1", + pack_sha256=hashlib.sha256(pack.read_bytes()).hexdigest(), + tokenizer_sha256=hashlib.sha256((args.model_dir / "tokenizer.json").read_bytes()).hexdigest(), + encoding_sha256=hashlib.sha256((args.model_dir / "encoding" / "encoding.py").read_bytes()).hexdigest(), + records=len(records), + sessions=len({r["session"] for r in records}), + settings={"max_tokens": args.max_tokens, "pack_tokens": args.pack_tokens, "session_tokens": args.session_tokens}, + ) + (args.out / "manifest.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps({k: report[k] for k in ("records", "sessions", "tokens")}, indent=2)) + return 0 + + +if __name__ == "__main__": + status = main() + sys.stdout.flush() + # Streaming `datasets` leaves native worker threads behind, and interpreter + # finalization then aborts with PyGILState_Release (exit 134). That happened + # on 2026-11-06 after the pack was fully written, and it stopped the pod. + # Skip finalization: every output is already closed. + os._exit(status) diff --git a/scripts/runpod/v41_reap/collect.py b/scripts/runpod/v41_reap/collect.py new file mode 100644 index 0000000..a7ab730 --- /dev/null +++ b/scripts/runpod/v41_reap/collect.py @@ -0,0 +1,634 @@ +"""REAP, imatrix, routes and teacher logits for DeepSeek-V4.1-Flash, on one GPU. + +DeepSeek's own reference code (``inference/model.py`` at the pinned revision) +runs one layer at a time through ``stream.LayerStreamer``. Weights come straight +from the HF shards (no ``convert.py``), and any GPU of 80 GB or more works. +Design: ``docs/deepseek-v4-flash/2026-11-06-v41-reap-single-gpu-design.md``. + +``MoE.forward`` is replaced by a copy that also records, for every routed slot, + +- REAP statistics (``keep.quality.reap.ReapStats`` arrays, ``reap__*``); +- imatrix sums for the two expert input spaces (``hidden`` = the post + ``ffn_norm`` row that w1/w3 read, ``down`` = the GLU activation w2 reads), + under the keys ``keep.quality.dsv4_teacher_runner`` finalizes; +- the first ``ROUTE_RECORD_TOKENS`` routes of each session (``routes__*``), + for ``benchmarks/replay_dsv4_expert_residency.py``. + +Two sweeps over the weights: + +- ``--sweep calibration`` writes REAP + imatrix + routes per calibration session. +- ``--sweep eval`` runs the teacher and every ``--survivors`` student side by + side. The teacher writes REAP + routes and top-K logits per eval session. Each + student is scored against the teacher's top-K (KL, NLL, top-1 agreement) and + written to ``/student-KEEP/``. + +Output is one npz per session under ``//sessions/``, so KEEP's +``finalize-reap``, the imatrix finalizers and ``score.py`` read it unchanged. + +The replacement forward computes exactly what the reference does (same ops, +same order). ``tests/test_runpod_v41_reap.py`` checks it bit for bit against a +CPU copy of the reference Gate / Expert / MoE, and ``tests/test_runpod_v41_stream.py`` +checks both sweeps against ``Transformer.forward``. One approximation: the +reference folds the gate weight in *before* w2 (``w2(g * h)``), so the expert +output norm ``||f||`` is taken as ``||w2(g * h)|| / g``. With FP8 activation +quantization in w2 that differs from ``||w2(h)||`` by rounding only. + +Usage, on the pod, one process per GPU (``--shard I/N`` splits the sessions):: + + CUDA_VISIBLE_DEVICES=0 python3 collect.py --hf-path /workspace/v41-hf \\ + --pack /workspace/pack/v41-pack.jsonl --out /workspace/out --sweep calibration --shard 0/1 + +then finalize-reap and reap-select, then the same command with ``--sweep eval`` +and ``--survivors 192=/workspace/out/survivors-192.json`` (repeatable). +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import hashlib +import json +import sys +import time +from collections import OrderedDict +from pathlib import Path +from typing import Any + +import numpy as np + +MODEL_ID = "deepseek-ai/DeepSeek-V4.1-Flash-1017" +PINNED_REVISION = "1bae0e1088993b9b82de08f59de5be5a34e42170" +#: sha256 of inference/model.py at PINNED_REVISION. The patched forward below +#: mirrors that file's MoE.forward and Expert.forward; refuse anything else. +REFERENCE_MODEL_PY_SHA256 = "e5decd76fe70fd56ef2f10124d245c6f64fb6a0003cc9068a1d3182a168b8d2d" +ROUTE_RECORD_TOKENS = 2048 # matches keep.quality.expert_residency +REAP_PREFIX = "reap__" +ROUTE_PREFIX = "routes__" +SPACES = ("hidden", "down") +CAPTURE_KEYS = ( + "topk_logit_ids", "topk_logit_values", "logsumexp", "tail_mass", "target_logit", # teacher pass + "kl_on_teacher_topk", "student_nll", "top1_agree", # masked-router student pass +) + + +# --------------------------------------------------------------------------- +# Accumulation (torch, float64 on device; one instance per process) +# --------------------------------------------------------------------------- + + +class SessionStats: + """Per-session sums over this rank's local experts; all-reduced at the end.""" + + def __init__(self, *, torch: Any, n_layers: int, n_experts: int, dims: dict[str, int], device: Any): + self.torch = torch + self.n_layers, self.n_experts = n_layers, n_experts + f64 = dict(dtype=torch.float64, device=device) + z = lambda *shape: torch.zeros(shape, **f64) + self.reap = { + "route_count": z(n_layers, n_experts), + "gate_sum": z(n_layers, n_experts), + "gate_norm_sum": z(n_layers, n_experts), + "norm_sum": z(n_layers, n_experts), + "norm_max": z(n_layers, n_experts), + "abs_max": z(n_layers, n_experts), + } + self.routed_slots = z(n_layers) + self.imatrix = { + space: { + "importance_sum": z(n_layers, n_experts, dim), + "affinity_weighted_importance": z(n_layers, n_experts, dim), + "route_count": z(n_layers, n_experts), + "affinity_score_sum": z(n_layers, n_experts), + "total_route_count": z(n_layers), + } + for space, dim in dims.items() + } + self.route_ids: list[list[np.ndarray]] = [[] for _ in range(n_layers)] + self.route_gates: list[list[np.ndarray]] = [[] for _ in range(n_layers)] + self.route_tokens = [0] * n_layers + + def add_routes(self, layer: int, indices: Any, weights: Any, *, record: bool) -> None: + """Once per layer per forward, on every rank (the gate is replicated).""" + + rows, top_k = indices.shape + if record: # rank 0 only: identical on every rank + self.routed_slots[layer] += rows * top_k + for m in self.imatrix.values(): + m["total_route_count"][layer] += rows * top_k + room = ROUTE_RECORD_TOKENS - self.route_tokens[layer] + if room > 0: + self.route_ids[layer].append(indices[:room].to("cpu").numpy().astype(np.int16)) + self.route_gates[layer].append(weights[:room].float().to("cpu").numpy().astype(np.float16)) + self.route_tokens[layer] += min(room, rows) + + def add_expert(self, layer: int, expert: int, x: Any, gate: Any, out: Any, h: Any) -> None: + """One local expert's routed rows: input ``x``, gate ``gate`` [r], output, GLU ``h``.""" + + torch = self.torch + g = gate.double() + out64 = out.double() + weighted_norm = out64.norm(dim=-1) # ||g f|| (g > 0) + f_norm = weighted_norm / g + r = self.reap + r["route_count"][layer, expert] += g.numel() + r["gate_sum"][layer, expert] += g.sum() + r["gate_norm_sum"][layer, expert] += weighted_norm.sum() + r["norm_sum"][layer, expert] += f_norm.sum() + r["norm_max"][layer, expert] = torch.maximum(r["norm_max"][layer, expert], f_norm.max()) + abs_max = (out64.abs().amax(dim=-1) / g).max() + r["abs_max"][layer, expert] = torch.maximum(r["abs_max"][layer, expert], abs_max) + for space, rows in (("hidden", x), ("down", h)): + if space not in self.imatrix: + continue + sq = rows.double().square() + m = self.imatrix[space] + m["importance_sum"][layer, expert] += sq.sum(dim=0) + m["affinity_weighted_importance"][layer, expert] += (sq * g[:, None]).sum(dim=0) + m["route_count"][layer, expert] += g.numel() + m["affinity_score_sum"][layer, expert] += g.sum() + + def reduce(self, dist: Any) -> None: + if dist is None or not dist.is_initialized(): + return + for name, value in self.reap.items(): + op = dist.ReduceOp.MAX if name in ("norm_max", "abs_max") else dist.ReduceOp.SUM + dist.all_reduce(value, op=op) + dist.all_reduce(self.routed_slots) + for m in self.imatrix.values(): + for value in m.values(): + dist.all_reduce(value) + + def arrays(self, layers: list[int]) -> dict[str, np.ndarray]: + np_ = lambda t: t.to("cpu").numpy() + out: dict[str, np.ndarray] = { + "layers": np.asarray(layers, dtype=np.int32), + "num_experts": np.asarray(self.n_experts, dtype=np.int32), + f"{REAP_PREFIX}layers": np.asarray(layers, dtype=np.int64), + f"{REAP_PREFIX}num_experts": np.asarray(self.n_experts, dtype=np.int64), + f"{REAP_PREFIX}routed_slots": np.rint(np_(self.routed_slots)).astype(np.int64), + } + for name, value in self.reap.items(): + arr = np_(value) + out[f"{REAP_PREFIX}{name}"] = np.rint(arr).astype(np.int64) if name == "route_count" else arr + for space, m in self.imatrix.items(): + # float32 on disk, as KEEP's validator requires (sums stay float64 until here) + out[f"importance_sum__{space}"] = np_(m["importance_sum"]).astype(np.float32) + out[f"affinity_weighted_importance__{space}"] = np_(m["affinity_weighted_importance"]).astype(np.float32) + out[f"route_count__{space}"] = np.rint(np_(m["route_count"])).astype(np.int64) + out[f"affinity_score_sum__{space}"] = np_(m["affinity_score_sum"]) + out[f"total_route_count__{space}"] = np.rint(np_(m["total_route_count"])).astype(np.int64) + if all(self.route_ids[layer] for layer in range(self.n_layers)): + n = min(self.route_tokens) + out[f"{ROUTE_PREFIX}layers"] = np.asarray(layers, dtype=np.int64) + out[f"{ROUTE_PREFIX}expert_ids"] = np.stack( + [np.concatenate(self.route_ids[layer])[:n] for layer in range(self.n_layers)], axis=1 + ) + out[f"{ROUTE_PREFIX}gate_weights"] = np.stack( + [np.concatenate(self.route_gates[layer])[:n] for layer in range(self.n_layers)], axis=1 + ) + return out + + +# --------------------------------------------------------------------------- +# Patched reference modules +# --------------------------------------------------------------------------- + + +def expert_forward_with_activation(F: Any, torch: Any, expert: Any, x: Any, weights: Any) -> tuple[Any, Any]: + """``Expert.forward(x, weights)`` from the reference, also returning the GLU activation.""" + + dtype = x.dtype + gate = expert.w1(x).float() + up = expert.w3(x).float() + if expert.swiglu_limit > 0: + up = torch.clamp(up, min=-expert.swiglu_limit, max=expert.swiglu_limit) + gate = torch.clamp(gate, max=expert.swiglu_limit) + h = F.silu(gate) * up + return expert.w2((weights * h).to(dtype)), h + + +def make_moe_forward(ref: Any, torch: Any, F: Any, dist: Any, state: dict[str, Any]): + """A drop-in ``MoE.forward`` that records into ``state['stats']`` when set.""" + + def forward(self, x, image_mask=None): + shape = x.size() + x = x.view(-1, self.dim) + weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten()) + y = torch.zeros_like(x, dtype=torch.float32) + counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() + stats = state.get("stats") + layer = self.layer_id if stats is not None and self.layer_id < stats.n_layers else None + if layer is not None: + stats.add_routes(layer, indices, weights, record=state["rank"] == 0) + for i in range(self.experts_start_idx, self.experts_end_idx): + if counts[i] == 0: + continue + expert = self.experts[i] + idx, top = torch.where(indices == i) + w = weights[idx, top, None] + if layer is None: + y[idx] += expert(x[idx], w) + continue + out, h = expert_forward_with_activation(F, torch, expert, x[idx], w) + stats.add_expert(layer, i, x[idx], w[:, 0], out, h) + y[idx] += out + if ref.world_size > 1: + dist.all_reduce(y) + y += self.shared_experts(x) + return y.type_as(x).view(shape) + + return forward + + +def make_head_forward(head: Any, torch: Any, F: Any, dist: Any, state: dict[str, Any], *, top_k: int, rows_per_slice: int): + """``ParallelHead.forward`` that also captures top-K logits at ``state['positions']``.""" + + original = head.forward + + def forward(x, full_logits=False): + positions = state.get("positions") + if positions is not None: + picked = x[0, positions] # [P, dim], already normalized + for start in range(0, picked.shape[0], rows_per_slice): + part = F.linear(picked[start : start + rows_per_slice].float(), head.weight) + if dist is not None and dist.is_initialized() and dist.get_world_size() > 1: + gathered = [torch.empty_like(part) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, part) + part = torch.cat(gathered, dim=-1) + lse = torch.logsumexp(part, dim=-1) + teacher = state.get("teacher") + if teacher is not None: # masked-router student pass: score against the teacher, keep scalars + rows = slice(start, start + part.shape[0]) + t_logp = teacher["logprob"][rows] + s_logp = part.gather(1, teacher["ids"][rows]) - lse[:, None] + kl = (t_logp.exp() * (t_logp - s_logp)).sum(dim=-1) + nll = lse - part.gather(1, teacher["targets"][rows, None])[:, 0] + agree = part.argmax(dim=-1) == teacher["ids"][rows, 0] + if state["rank"] == 0: + cap = state["captured"] + cap["kl_on_teacher_topk"].append(kl.cpu().numpy().astype(np.float32)) + cap["student_nll"].append(nll.cpu().numpy().astype(np.float32)) + cap["top1_agree"].append(agree.cpu().numpy()) + continue + values, ids = part.topk(top_k, dim=-1) + tail = (1.0 - torch.exp(values - lse[:, None]).sum(dim=-1)).clamp_min(0.0) + target = part.gather(1, state["targets"][start : start + part.shape[0], None])[:, 0] + if state["rank"] == 0: + cap = state["captured"] + cap["target_logit"].append(target.cpu().numpy().astype(np.float32)) + cap["topk_logit_ids"].append(ids.to(torch.int32).cpu().numpy()) + cap["topk_logit_values"].append(values.cpu().numpy().astype(np.float16)) + cap["logsumexp"].append(lse.cpu().numpy().astype(np.float32)) + cap["tail_mass"].append(tail.cpu().numpy().astype(np.float16)) + return original(x, full_logits) + + return forward + + +# --------------------------------------------------------------------------- +# Pack and output +# --------------------------------------------------------------------------- + + +def load_sessions(pack: Path) -> OrderedDict[str, list[dict[str, Any]]]: + """Pack lines grouped by ``session`` in file order. + + Each line: ``{"id", "session", "split": "calibration"|"eval", "domain", + "token_ids": [...], "target_spans": [[start, end], ...]}``. Target spans + index ``token_ids``; logits are captured at ``span_index - 1``. + """ + + sessions: OrderedDict[str, list[dict[str, Any]]] = OrderedDict() + with pack.open() as handle: + for line in handle: + if line.strip(): + record = json.loads(line) + sessions.setdefault(record["session"], []).append(record) + for name, records in sessions.items(): + splits = {r["split"] for r in records} + if len(splits) != 1 or not splits <= {"calibration", "eval"}: + raise ValueError(f"session {name} mixes or misnames splits: {sorted(splits)}") + if splits == {"eval"} and len(records) != 1: + raise ValueError(f"eval session {name} must hold exactly one sample (one logits stream)") + return sessions + + +def target_positions(record: dict[str, Any]) -> np.ndarray: + n = len(record["token_ids"]) + targets = sorted({t for start, end in record["target_spans"] for t in range(max(int(start), 1), min(int(end), n))}) + return np.asarray(targets, dtype=np.int64) - 1 + + +def sha256_ids(ids: list[int]) -> str: + return hashlib.sha256(np.asarray(ids, dtype=np.int32).tobytes()).hexdigest() + + +def identity(*, session: str, split: str, records: list[dict[str, Any]], config_sha: str, chunk: int) -> dict[str, np.ndarray]: + ids_hash = hashlib.sha256("".join(sha256_ids(r["token_ids"]) for r in records).encode()).hexdigest() + mode = "calibration" if split == "calibration" else "logits" + return { + "record_type": np.asarray(f"dsv4_teacher_{mode}_v1"), + "schema_version": np.asarray(1, dtype=np.int32), + "prompt_id": np.asarray(session), + "campaign_split": np.asarray(split), + "token_ids_sha256": np.asarray(ids_hash), + "token_count": np.asarray(sum(len(r["token_ids"]) for r in records), dtype=np.int64), + "prefill_chunk_tokens": np.asarray(chunk, dtype=np.int32), + "generation_config_sha256": np.asarray(config_sha), + "producer": np.asarray("keep/scripts/runpod/v41_reap/collect.py (DeepSeek reference torch)"), + "model_id": np.asarray(MODEL_ID), + "model_revision": np.asarray(PINNED_REVISION), + "domains": np.asarray(sorted({r["domain"] for r in records})), + "sample_ids": np.asarray([r["id"] for r in records]), + } + + +def save_npz(path: Path, arrays: dict[str, np.ndarray]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp.npz") + np.savez(tmp, **arrays) + tmp.replace(path) + + +# --------------------------------------------------------------------------- +# Sweeps (stream.LayerStreamer: one GPU, one layer at a time) +# --------------------------------------------------------------------------- + +TEACHER = "teacher" + + +def chunk_sessions(sessions: list[tuple[str, list[dict[str, Any]]]], budget: int, streams: int) -> list[list[tuple[str, list[dict[str, Any]]]]]: + """Whole sessions, in order, up to ``budget`` resident tokens (tokens x streams) per chunk.""" + + chunks: list[list[tuple[str, list[dict[str, Any]]]]] = [] + size = 0 + for name, records in sessions: + n = sum(len(r["token_ids"]) for r in records) * streams + if not chunks or (chunks[-1] and size + n > budget): + chunks.append([]) + size = 0 + chunks[-1].append((name, records)) + size += n + return chunks + + +def survivor_masks(table: dict[str, Any], n_experts: int) -> dict[int, Any]: + """reap-select table -> {layer: bool tensor, True = pruned}. Masking the router to -inf + equals slicing it (tests/test_reap.py), so one set of weights scores every candidate.""" + + import torch + + masks = {} + for layer, kept in table["survivors"].items(): + mask = torch.ones(n_experts, dtype=torch.bool) + mask[torch.as_tensor(kept)] = False + masks[int(layer)] = mask + return masks + + +def install_hooks(streamer: Any, *, top_k: int, rows_per_slice: int) -> dict[str, Any]: + """Patch the reference MoE and the streamer's head once; returns the shared ``state``.""" + + import torch + import torch.nn.functional as F + + state: dict[str, Any] = {"rank": 0, "stats": None, "positions": None, "teacher": None} + streamer.ref.MoE.forward = make_moe_forward(streamer.ref, torch, F, None, state) + streamer.head.forward = make_head_forward(streamer.head, torch, F, None, state, top_k=top_k, rows_per_slice=rows_per_slice) + return state + + +def run_chunk( + streamer: Any, + state: dict[str, Any], + chunk: list[tuple[str, list[dict[str, Any]]]], + *, + split: str, + students: dict[str, dict[int, Any]], + dims: dict[str, int], + device: Any, + teacher_mask: dict[int, Any] | None = None, +) -> dict[tuple[str, str], dict[str, np.ndarray]]: + """One pass over the weights for ``chunk``; returns arrays per (stream, session). + + Calibration: the teacher only, REAP + imatrix + routes per session. + Eval: the teacher (REAP + routes + top-K logits) and every student (scored against + the teacher) side by side. Teacher records go through the MoE per session, so the + stats stay per session; each student's records go through it together. + + ``teacher_mask`` (calibration only) masks the measured stream's router to a survivors table, + so the stats describe that pruned model's own routing and activations. + """ + + import torch + from stream import Item + + args = streamer.args + n_layers, n_experts = int(args.n_layers), int(args.n_routed_experts) + stats = { + name: SessionStats(torch=torch, n_layers=n_layers, n_experts=n_experts, dims=dims if split == "calibration" else {}, device=device) + for name, _ in chunk + } + items: list[Item] = [] + streams_of: dict[str, list[tuple[str, Item]]] = {} + for name, records in chunk: + for record in records: + tokens = torch.as_tensor([record["token_ids"]], dtype=torch.long, device=device) + for stream in (TEACHER, *students) if split == "eval" else (TEACHER,): + item = Item(tokens=tokens, group=(TEACHER, name) if stream == TEACHER else (stream, None)) + items.append(item) + streams_of.setdefault(name, []).append((stream, item)) + + def before_moe(layer: int, block: Any, key: tuple[str, str | None]) -> None: + stream, name = key + bias = block.ffn.gate.bias + if not hasattr(block, "full_bias"): + block.full_bias = bias.data.clone() + bias.data.copy_(block.full_bias) + mask = students[stream].get(layer) if stream != TEACHER else (teacher_mask or {}).get(layer) + if mask is not None: + bias.data[mask.to(bias.device)] = float("-inf") + state["stats"] = stats[name] if stream == TEACHER and layer < n_layers else None + + with torch.inference_mode(): + streamer.forward(items, before_moe) + state["stats"] = None + + layers = list(range(n_layers)) + results: dict[tuple[str, str], dict[str, np.ndarray]] = {} + for name, records in chunk: + results[(TEACHER, name)] = stats[name].arrays(layers) + if split != "eval": + continue + (record,) = records + positions = target_positions(record) + targets = np.asarray(record["token_ids"], dtype=np.int64)[positions + 1] + state["positions"] = torch.as_tensor(positions, device=device) + state["targets"] = torch.as_tensor(targets, device=device) + base = {"positions": positions.astype(np.int32), "target_token_ids": targets.astype(np.int32)} + state["teacher"] = None + for stream, item in streams_of[name]: # the teacher first, then the students it scores + state["captured"] = {k: [] for k in CAPTURE_KEYS} + with torch.inference_mode(): + streamer.head(streamer.final_hidden(item)) + cap = {k: np.concatenate(v) for k, v in state["captured"].items() if v} + results.setdefault((stream, name), {}).update(base | cap) + if stream == TEACHER: + state["teacher"] = { # the fp16 values the npz keeps, as the old two-pass flow read them + "ids": torch.as_tensor(cap["topk_logit_ids"].astype(np.int64), device=device), + "logprob": torch.as_tensor(cap["topk_logit_values"].astype(np.float32) - cap["logsumexp"][:, None], device=device), + "targets": state["targets"], + } + state["positions"] = state["targets"] = state["teacher"] = None + return results + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--hf-path", required=True, type=Path, help="the pinned HF checkpoint (safetensors shards + tokenizer)") + parser.add_argument("--ref-dir", type=Path, help="reference inference/ directory (default: HF_PATH/inference)") + parser.add_argument("--pack", required=True, type=Path) + parser.add_argument("--out", required=True, type=Path) + parser.add_argument("--sweep", required=True, choices=("calibration", "eval")) + parser.add_argument( + "--survivors", action="append", default=[], metavar="KEEP=PATH", + help="eval sweep (repeatable): also run a student with the router masked to this reap-select " + "table, scored against the teacher, into /student-KEEP/", + ) + parser.add_argument( + "--mask", type=Path, + help="calibration sweep: mask the router to this survivors table, so REAP, routes and imatrix " + "are measured on that pruned model's own routing (for re-ranking after a prune)", + ) + parser.add_argument("--chunk-tokens", type=int, default=524288, help="resident tokens per weight pass (x streams)") + parser.add_argument("--shard", default="0/1", help="I/N: this process takes every Nth session from I") + parser.add_argument("--max-seq-len", type=int, default=16384) + parser.add_argument("--top-k", type=int, default=2048, help="teacher logits kept per eval position") + parser.add_argument("--rows-per-slice", type=int, default=256) + parser.add_argument("--deadline-utc", help="stop cleanly before the first chunk that starts after this time") + parser.add_argument("--skip-hash-check", action="store_true", help="tests only") + args = parser.parse_args(argv) + if args.survivors and args.sweep != "eval": + parser.error("--survivors needs --sweep eval") + if args.mask and args.sweep != "calibration": + parser.error("--mask needs --sweep calibration") + shard, n_shards = map(int, args.shard.split("/")) + + import torch + from transformers import AutoTokenizer + + sys.path.insert(0, str(Path(__file__).resolve().parent)) + import stream + + ref_dir = args.ref_dir or args.hf_path / "inference" + digest = hashlib.sha256((ref_dir / "model.py").read_bytes()).hexdigest() + if not args.skip_hash_check and digest != REFERENCE_MODEL_PY_SHA256: + raise SystemExit(f"{ref_dir / 'model.py'} sha256 {digest} is not the pinned reference ({REFERENCE_MODEL_PY_SHA256})") + ref = stream.import_reference(ref_dir) + + torch.cuda.memory._set_allocator_settings("expandable_segments:True") + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.manual_seed(33377335) + model_args = ref.ModelArgs(**json.loads((ref_dir / "config.json").read_text())) + model_args.max_batch_size = 1 + model_args.max_seq_len = args.max_seq_len + + ckpt = stream.HFCheckpoint(args.hf_path) + layout = ref.EngramLayout.from_args(model_args) + hash_state = None + if layout is not None: + with torch.device("cuda"): + hash_state = ref.NgramHashState(model_args, layout, AutoTokenizer.from_pretrained(args.hf_path)) + engram_rows = {layer: stream.EngramRows.from_checkpoint(ckpt, layer, ref.fp8_block_size) for layer in model_args.engram_layer_ids} + streamer = stream.LayerStreamer(ref, model_args, ckpt.load, engram_rows, hash_state, "cuda") + state = install_hooks(streamer, top_k=args.top_k, rows_per_slice=args.rows_per_slice) + + n_experts = int(model_args.n_routed_experts) + students: dict[str, dict[int, Any]] = {} + survivors_sha: dict[str, str] = {} + for spec in args.survivors: + keep, _, table = spec.partition("=") + students[keep] = survivor_masks(json.loads(Path(table).read_text()), n_experts) + survivors_sha[keep] = hashlib.sha256(Path(table).read_bytes()).hexdigest() + + teacher_mask = survivor_masks(json.loads(args.mask.read_text()), n_experts) if args.mask else None + run_config = { + "model_id": MODEL_ID, + "revision": PINNED_REVISION, + "model_py_sha256": digest, + "pack_sha256": hashlib.sha256(args.pack.read_bytes()).hexdigest(), + "max_seq_len": args.max_seq_len, + "top_k": args.top_k, + "driver": "layer-streamed, world_size 1", + } + if args.mask: + run_config["mask_sha256"] = hashlib.sha256(args.mask.read_bytes()).hexdigest() + text = json.dumps(run_config, indent=2, sort_keys=True) + "\n" + config_path = args.out / "run-config.json" + args.out.mkdir(parents=True, exist_ok=True) + if not config_path.exists() or config_path.read_text() != text: + config_path.write_text(text) + config_sha = hashlib.sha256(json.dumps(run_config, sort_keys=True).encode()).hexdigest() + teacher_config_sha = hashlib.sha256(text.encode()).hexdigest() + + def target(stream_name: str, name: str) -> Path: + root = args.out if stream_name == TEACHER else args.out / f"student-{stream_name}" + return root / args.sweep / "sessions" / f"{name}.npz" + + streams = (TEACHER, *students) + todo = [ + (name, records) + for i, (name, records) in enumerate(item for item in load_sessions(args.pack).items() if item[1][0]["split"] == args.sweep) + if i % n_shards == shard and not all(target(s, name).exists() for s in streams) + ] + for name, records in todo: + if any(len(r["token_ids"]) > args.max_seq_len for r in records): + raise ValueError(f"session {name} has a record over --max-seq-len {args.max_seq_len}") + deadline = None + if args.deadline_utc: + deadline = dt.datetime.fromisoformat(args.deadline_utc.replace("Z", "+00:00")).timestamp() + dims = {"hidden": int(model_args.dim), "down": int(model_args.moe_inter_dim)} + started = time.time() + progress = args.out / f"progress-{args.sweep}-{shard}of{n_shards}.jsonl" + for chunk in chunk_sessions(todo, args.chunk_tokens, len(streams)): + if deadline is not None and time.time() > deadline: + print("deadline reached; stopping cleanly", flush=True) + break + t0 = time.time() + results = run_chunk(streamer, state, chunk, split=args.sweep, students=students, dims=dims, device="cuda", + teacher_mask=teacher_mask) + for name, records in chunk: + ident = identity(session=name, split=args.sweep, records=records, config_sha=config_sha, chunk=args.max_seq_len) + for stream_name in streams: + arrays = ident | results[(stream_name, name)] + if stream_name != TEACHER: + arrays["record_type"] = np.asarray("dsv4_masked_student_eval_v1") + arrays["survivors_sha256"] = np.asarray(survivors_sha[stream_name]) + arrays["teacher_run_config_sha256"] = np.asarray(teacher_config_sha) + save_npz(target(stream_name, name), arrays) + tokens = sum(len(r["token_ids"]) for _, records in chunk for r in records) + entry = { + "sessions": [name for name, _ in chunk], + "tokens": tokens, + "streams": len(streams), + "seconds": round(time.time() - t0, 3), + "tok_per_s": round(tokens * len(streams) / max(time.time() - t0, 1e-9), 1), + "max_gpu_gb": round(torch.cuda.max_memory_allocated() / 2**30, 1), + "elapsed_h": round((time.time() - started) / 3600, 3), + } + with progress.open("a") as handle: + handle.write(json.dumps(entry) + "\n") + print(json.dumps(entry), flush=True) + del results + torch.cuda.empty_cache() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/extract_tensors.py b/scripts/runpod/v41_reap/extract_tensors.py new file mode 100644 index 0000000..b0c6414 --- /dev/null +++ b/scripts/runpod/v41_reap/extract_tensors.py @@ -0,0 +1,76 @@ +"""Raw tensors out of an HF safetensors checkpoint, as one safetensors file on stdout. + +Either one layer's chosen routed experts, renumbered 0..n-1 in the order given, or every router +tensor of the decoder layers. Names are the reference names (``stream.reference_name``), bytes +are copied untouched. Standard library only, so it runs on any pod: + + ssh pod python3 extract_tensors.py /workspace/v41-hf --layer 3 --experts 1,2,5 > layer-03.safetensors + ssh pod python3 extract_tensors.py /workspace/v41-hf --gates > router.safetensors +""" + +import argparse +import json +import os +import re +import struct +import sys +from pathlib import Path + + +def reference_name(name: str) -> str: # stream.reference_name + name = name.removeprefix("model.").replace("self_attn", "attn") + if not name.startswith("vision."): + name = name.replace("mlp", "ffn") + return name.replace("weight_scale_inv", "scale").replace("e_score_correction_bias", "bias") + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("root", type=Path) + p.add_argument("--layer", type=int) + p.add_argument("--experts", help="comma-separated expert ids of --layer") + p.add_argument("--gates", action="store_true") + args = p.parse_args() + index = json.loads((args.root / "model.safetensors.index.json").read_text())["weight_map"] + order = {int(e): i for i, e in enumerate(args.experts.split(","))} if args.experts else {} + picked: dict[str, tuple[str, str]] = {} # new name -> (shard, stored name) + for stored, shard in index.items(): + ref = reference_name(stored) + if args.gates: + if re.fullmatch(r"layers\.\d+\.ffn\.gate\..+", ref): + picked[ref] = (shard, stored) + else: + m = re.fullmatch(rf"layers\.{args.layer}\.ffn\.experts\.(\d+)\.(.+)", ref) + if m and int(m[1]) in order: + picked[f"layers.{args.layer}.ffn.experts.{order[int(m[1])]}.{m[2]}"] = (shard, stored) + assert picked and (args.gates or len(picked) == 6 * len(order)), len(picked) + + shards: dict[str, tuple[int, dict, int]] = {} + header, slabs, offset = {}, [], 0 + for name in sorted(picked): + shard, stored = picked[name] + if shard not in shards: + fd = os.open(args.root / shard, os.O_RDONLY) + (size,) = struct.unpack(" dict[str, np.ndarray]: + """Summed ``importance_sum__{hidden,down}`` [L, 384, in] over the calibration sessions.""" + + total: dict[str, np.ndarray] = {} + for path in sorted((stats / "calibration" / "sessions").glob("cal-*.npz")): + with np.load(path) as z: + for space in ("hidden", "down"): + a = z[f"importance_sum__{space}"].astype(np.float64) + total[space] = total.get(space, 0) + a + assert total, "no calibration sessions found" + return total + + +def decode_fp4(weight, scale) -> np.ndarray: + import torch_kernel + + w = torch_kernel.decode_e2m1(weight) # [out, in] fp32 + return (w.unflatten(-1, (-1, 32)) * scale.float()[..., None]).flatten(-2).numpy() + + +def decode_codes(z, proj: str, e: int) -> np.ndarray: + """One expert's projection from a keep_quant layer file, as ``local_eval.expert_override`` decodes it.""" + + from keep.vq.e8 import e8_1bit_grid, e8p_full_grid + + bits = int(z["code_bits"]) if "code_bits" in z else 16 # the first E8P runs wrote no code_bits + grid = (e8p_full_grid() if bits == 16 else e8_1bit_grid()).astype(np.float32) + vec = grid[z[f"{proj}_codes"][e].astype(np.int64)] # [out, in/8, 8] + scales = z[f"{proj}_scales"][e].astype(np.float32) + vec = vec * np.repeat(scales, vec.shape[1] // scales.shape[-1], axis=-1)[..., None] + return vec.reshape(vec.shape[0], -1) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--hf-path", required=True, type=Path) + p.add_argument("--checkpoint-survivors", required=True, type=Path) + p.add_argument("--stats", type=Path, help="local REAP-stats copy (for --importance pod)") + p.add_argument("--importance", required=True, help="'pod' or a local imatrix npz") + p.add_argument("--out", required=True, type=Path) + p.add_argument("--layers", default="all", help="'all' or comma-separated layer ids") + p.add_argument("--iterations", type=int, default=8) + p.add_argument("--backend", default="metal", help="E8P search backend: metal or numpy") + p.add_argument("--code-bits", type=int, default=16, choices=(8, 16), + help="16: E8P, 2 bits/weight + scales; 8: E8 1-bit codebook, 1 bit/weight + scales") + p.add_argument("--experts", type=int, default=0, help="debug: fit only the first N experts per layer") + p.add_argument("--residual-of", type=Path, help="fit W minus this keep_quant run's decoded codes (nested tier)") + args = p.parse_args(argv) + + import stream + from keep.convert.dsv4_vq_pilot import fit_projection, projection_error, reconstruct_quantized + + survivors = json.loads(args.checkpoint_survivors.read_text())["survivors"] + ckpt = stream.HFCheckpoint(args.hf_path) + n_layers = len(survivors) + layers = range(n_layers) if args.layers == "all" else [int(x) for x in args.layers.split(",")] + if args.importance == "pod": + full = pod_importance(args.stats) + imp = {space: np.stack([full[space][l][sorted(survivors[str(l)])] for l in range(n_layers)]) for space in full} + else: + with np.load(args.importance) as z: + imp = {space: z[f"importance__{space}"].astype(np.float64) for space in ("hidden", "down")} + args.out.mkdir(parents=True, exist_ok=True) + (args.out / "config.json").write_text(json.dumps(vars(args) | {"hf_path": str(args.hf_path), "out": str(args.out), + "checkpoint_survivors": str(args.checkpoint_survivors), + "stats": str(args.stats)}, indent=2, default=str)) + + for layer in layers: + dest = args.out / f"layer-{layer:02d}.npz" + if dest.exists(): + continue + t0 = time.time() + sd = ckpt.load(f"layers.{layer}.ffn.experts.") + n_exp = len(survivors[str(layer)]) if not args.experts else args.experts + arrays: dict[str, list[np.ndarray]] = {} + errs = {k: [] for k in PROJ_SPACE} + total = {k: [] for k in PROJ_SPACE} + base = np.load(args.residual_of / f"layer-{layer:02d}.npz") if args.residual_of else None + for e in range(n_exp): + for proj, space in PROJ_SPACE.items(): + w = decode_fp4(sd[f"{e}.{proj}.weight"], sd[f"{e}.{proj}.scale"]) + group = 512 if w.shape[1] % 512 == 0 else 256 + d = imp[space][layer, e].astype(np.float32) + b = decode_codes(base, proj, e) if base is not None else 0.0 + q, metrics, _ = fit_projection(w - b, d, group_size=group, + code_bits=args.code_bits, iterations=args.iterations, backend=args.backend) + if base is not None: + total[proj].append(projection_error(w, b + reconstruct_quantized(q), d).weighted_relative_mse) + # 1-bit E8 codes are one byte; E8P codes are two + arrays.setdefault(f"{proj}_codes", []).append(np.asarray(q.codes, dtype=np.uint8 if args.code_bits == 8 else np.uint16)) + arrays.setdefault(f"{proj}_scales", []).append(np.asarray(q.scales, dtype=np.float16)) + errs[proj].append(metrics.weighted_relative_mse if hasattr(metrics, "weighted_relative_mse") else float("nan")) + tmp = dest.with_suffix(".tmp.npz") + np.savez(tmp, code_bits=np.asarray(args.code_bits), **{k: np.stack(v) for k, v in arrays.items()}) + tmp.rename(dest) + row = {"layer": layer, "experts": n_exp, "seconds": round(time.time() - t0, 1), + **{f"{k}_wrel_mse": float(np.mean(v)) for k, v in errs.items()}, + **{f"{k}_total_wrel_mse": float(np.mean(v)) for k, v in total.items() if v}} + with (args.out / "metrics.jsonl").open("a") as f: + f.write(json.dumps(row) + "\n") + print(json.dumps(row), flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/keep_quant_torch.py b/scripts/runpod/v41_reap/keep_quant_torch.py new file mode 100644 index 0000000..4c5de85 --- /dev/null +++ b/scripts/runpod/v41_reap/keep_quant_torch.py @@ -0,0 +1,253 @@ +"""``keep_quant.py``'s 2-bit E8P fit on torch, for the experts a REAP prune removed. + +``keep_quant.py`` fits the survivors of a pruned checkpoint with KEEP's Metal search, on a Mac. +This fits the other experts from the full checkpoint on a CUDA pod, so the cache's pool can be +all 384. The fit is the same alternation as ``quantize_weight_importance_aware`` with its +default policy (weighted nearest E8P codeword under the per-input-channel importance, then the +exact weighted least-squares group scale; 16-bit codes; group 512 for w1/w3, 256 for w2), in +torch. Rows are independent, so ``check`` can compare a few rows against a ``keep_quant`` file. + + # Mac: the importance of the removed experts, one small file to upload + keep_quant_torch.py importance --stats STATS --survivors survivors-192.json --out imp-pruned.npz + # Mac: the torch fit against stored codes of the survivors (fitted by the Metal search) + keep_quant_torch.py check --hf-path PRUNED_CKPT --survivors survivors-192.json --stats STATS --codes QUANT_DIR + # pod: fit, one process per GPU (--layers splits the work); finished layers are skipped + keep_quant_torch.py fit --hf-path FULL_CKPT --importance imp-pruned.npz --out OUT [--layers 0,1,...] + # pod: every router row of the full checkpoint + keep_quant_torch.py router --hf-path FULL_CKPT --out router-384.npz + +``fit`` writes ``layer-XX.npz`` in ``keep_quant``'s format (experts in the order of +``imp-pruned.npz``'s ``experts`` row, which ``experts.json`` repeats) and ``metrics.jsonl``. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch + +KIT = Path(__file__).resolve().parent +sys.path.insert(0, str(KIT)) +sys.path.insert(0, str(KIT.parents[2] / "src")) + +from keep.vq.e8 import E8P_SHUFFLE_MAP, e8p_abs_grid, e8p_full_grid + +PROJ_SPACE = {"w1": "hidden", "w3": "hidden", "w2": "down"} + + +def _parity(x: torch.Tensor) -> torch.Tensor: + x = x ^ (x >> 4) + x = x ^ (x >> 2) + return (x ^ (x >> 1)) & 1 + + +class Search: + """Weighted nearest E8P codeword; ``e8._encode_e8p_rtn_per_row_diagonal_hessian`` in torch.""" + + def __init__(self, device, vectors_per_chunk: int = 1 << 17): + base = e8p_abs_grid().astype(np.float32) # signed base rows [256, 8] + self.abs_rows = torch.from_numpy(np.abs(base)).to(device) + self.bit = torch.from_numpy(E8P_SHUFFLE_MAP.astype(np.int32)).to(device) # dim -> sign bit + self.base_signs = ((torch.from_numpy(base < 0).to(device).int() << self.bit).sum(-1)).int() + self.grid = torch.from_numpy(e8p_full_grid().astype(np.float32)).to(device) # [65536, 8] + self.vectors_per_chunk = vectors_per_chunk + + def assign(self, vectors: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """vectors [rows, P, 8] (already divided by their scales), w [P, 8] -> codes [rows, P] int32.""" + abs_norm_w = w @ (self.abs_rows * self.abs_rows).T # [P, 256] + step = max(1, self.vectors_per_chunk // vectors.shape[1]) + return torch.cat([self._assign(vectors[i:i + step], w, abs_norm_w) for i in range(0, vectors.shape[0], step)]) + + def _assign(self, chunk, w, abs_norm_w): + best_dist = torch.full(chunk.shape[:2], torch.inf, device=chunk.device) + best_code = torch.zeros(chunk.shape[:2], dtype=torch.int32, device=chunk.device) + for parity, shift in ((0, 0.25), (1, -0.25)): + target = chunk - shift + weighted = target.abs() * w # [R, P, 8] + dist = (target.abs() * weighted).sum(-1, keepdim=True) + abs_norm_w - 2.0 * (weighted @ self.abs_rows.T) + signs = ((target < 0).int() << self.bit).sum(-1, dtype=torch.int32)[..., None] ^ self.base_signs # [R, P, 256] + odd = _parity(signs).bool() + # An odd sign count is fixed by flipping the dimension that costs least. + dist = torch.where(odd, dist + (4 * weighted[..., None, :] * self.abs_rows).amin(-1), dist) + row = dist.argmin(-1, keepdim=True) + row_dist, signs, odd = dist.gather(-1, row)[..., 0], signs.gather(-1, row)[..., 0], odd.gather(-1, row)[..., 0] + row = row[..., 0] + penalty = 4 * weighted * self.abs_rows[row] # [R, P, 8] + flipped = signs[..., None] ^ (1 << self.bit) + tied = penalty == penalty.amin(-1, keepdim=True) + flip = torch.where(tied, flipped, 65535).argmin(-1) # ties: the lowest resulting code + signs = torch.where(odd, signs ^ (1 << self.bit[flip]), signs) + code = (row.int() << 8) | (signs ^ parity) + take = (row_dist < best_dist) | ((row_dist == best_dist) & (code < best_code)) + best_dist, best_code = torch.where(take, row_dist, best_dist), torch.where(take, code, best_code) + return best_code + + def fit(self, weight: torch.Tensor, importance: torch.Tensor, group: int, iterations: int = 8): + """weight [out, in] fp32, importance [in] -> (codes uint16 [out, in/8], scales fp16 [out, in/group], + weighted relative squared error of the stored code).""" + out_dim, in_dim = weight.shape + grouped = weight.reshape(out_dim, in_dim // group, group // 8, 8) + h = importance.reshape(in_dim // group, group // 8, 8) + w = importance.reshape(-1, 8) + scales = grouped.abs().amax((2, 3)) / self.grid.abs().max() + scales = torch.where(scales > 0, scales, torch.ones_like(scales)) + + def assign(): + return self.assign((grouped / scales[:, :, None, None]).reshape(out_dim, -1, 8), w).reshape(grouped.shape[:3]) + + for _ in range(iterations): + decoded = self.grid[assign().long()] + numerator = (h * grouped * decoded).double().sum((2, 3)) + denominator = (h * decoded * decoded).double().sum((2, 3)) + ok = (denominator > 0) & (numerator > 0) + scales = torch.where(ok, numerator / denominator.clamp_min(1e-300), scales.double()).float() + codes = assign() + stored = scales.half() + error = self.grid[codes.long()] * stored.float()[:, :, None, None] - grouped + wrel = float((h * error * error).double().sum() / (h * grouped * grouped).double().sum()) + return codes.reshape(out_dim, in_dim // 8).cpu().numpy().astype(np.uint16), stored.cpu().numpy(), wrel + + +def decode_fp4(weight: torch.Tensor, scale: torch.Tensor, device) -> torch.Tensor: + import torch_kernel + + # The packed dtypes leave the CPU as bytes and floats; not every backend has them. + w = torch_kernel.decode_e2m1(weight.contiguous().view(torch.uint8).to(device)) # [out, in] fp32 + return (w.unflatten(-1, (-1, 32)) * scale.float().to(device)[..., None]).flatten(-2) + + +def load_expert(ckpt, layer: int, expert: int, device) -> dict[str, torch.Tensor]: + sd = ckpt.load(f"layers.{layer}.ffn.experts.{expert}.") + return {proj: decode_fp4(sd[f"{proj}.weight"], sd[f"{proj}.scale"], device) for proj in PROJ_SPACE} + + +def importance(args) -> None: + from keep_quant import pod_importance + + # Without --survivors: every expert, in order. + survivors = json.loads(args.survivors.read_text())["survivors"] if args.survivors else {} + full = pod_importance(args.stats) # {space: [L, 384, in]} + n_layers, n_experts = full["hidden"].shape[:2] + removed = np.array([sorted(set(range(n_experts)) - set(survivors.get(str(layer), ()))) for layer in range(n_layers)]) + out, replaced = {"experts": removed}, 0 + for space, total in full.items(): + picked = np.stack([total[layer][removed[layer]] for layer in range(n_layers)]) + # The search needs some weight in every group of eight inputs. An expert that calibration + # never routed has none: it gets its layer's mean over all experts. + dead = ~(picked.reshape(*picked.shape[:2], -1, 8) > 0).any(-1).all(-1) # [L, E] + picked[dead] = np.broadcast_to(total.mean(1)[:, None, :], picked.shape)[dead] + replaced += int(dead.sum()) + assert (picked.reshape(*picked.shape[:2], -1, 8) > 0).any(-1).all() and np.isfinite(picked).all(), space + out[space] = picked.astype(np.float32) + np.savez(args.out, **out) + print(f"{args.out}: {removed.shape[1]} experts in each of {n_layers} layers; " + f"{replaced} (expert, space) importances replaced by the layer mean") + + +def check(args) -> None: + """A few rows of some surviving experts, fitted here, against the stored Metal-search codes.""" + import stream + from keep_quant import pod_importance + + device = torch.device(args.device) + survivors = json.loads(args.survivors.read_text())["survivors"] + full, ckpt, search = pod_importance(args.stats), stream.HFCheckpoint(args.hf_path), Search(device) + grid = e8p_full_grid() + for layer, expert in ((0, 3), (20, 0), (39, 191)): + z = np.load(args.codes / f"layer-{layer:02d}.npz") + weights = load_expert(ckpt, layer, expert, device) + original = sorted(survivors[str(layer)])[expert] + for proj, space in PROJ_SPACE.items(): + w = weights[proj][: args.rows] + d = torch.from_numpy(full[space][layer][original].astype(np.float32)).to(device) + started = time.time() + codes, _, wrel = search.fit(w, d, 512 if w.shape[1] % 512 == 0 else 256) + seconds = time.time() - started + ref_codes, ref_scales = z[f"{proj}_codes"][expert, : args.rows], z[f"{proj}_scales"][expert, : args.rows] + group = w.shape[1] // ref_scales.shape[1] + ref = (grid[ref_codes].reshape(args.rows, -1, group // 8, 8) * ref_scales.astype(np.float32)[:, :, None, None]) + ref = torch.from_numpy(ref.reshape(args.rows, -1)).to(device) + ref_wrel = float((d * (ref - w) ** 2).double().sum() / (d * w * w).double().sum()) + same = float((codes == ref_codes).mean()) + print(f"layer {layer} expert {expert} {proj}: codes equal {same:.4f}, weighted error {wrel:.6f} " + f"(stored {ref_wrel:.6f}), {seconds:.2f} s for {args.rows} rows") + assert same > 0.99 and abs(wrel - ref_wrel) < 1e-3 * ref_wrel, (same, wrel, ref_wrel) + + +def fit(args) -> None: + import stream + + device = torch.device(args.device) + with np.load(args.importance) as z: + imp = {"hidden": z["hidden"], "down": z["down"]} + experts = z["experts"] + ckpt, search = stream.HFCheckpoint(args.hf_path), Search(device) + args.out.mkdir(parents=True, exist_ok=True) + (args.out / "experts.json").write_text(json.dumps({"experts": experts.tolist(), "source": str(args.hf_path)})) + layers = range(experts.shape[0]) if args.layers == "all" else [int(x) for x in args.layers.split(",")] + for layer in layers: + dest = args.out / f"layer-{layer:02d}.npz" + if dest.exists(): + continue + started = time.time() + arrays: dict[str, list[np.ndarray]] = {} + errors: dict[str, list[float]] = {proj: [] for proj in PROJ_SPACE} + for slot, expert in enumerate(experts[layer][: args.experts or None].tolist()): + weights = load_expert(ckpt, layer, expert, device) + for proj, space in PROJ_SPACE.items(): + w = weights[proj] + codes, scales, wrel = search.fit(w, torch.from_numpy(imp[space][layer, slot]).to(device), + 512 if w.shape[1] % 512 == 0 else 256, args.iterations) + arrays.setdefault(f"{proj}_codes", []).append(codes) + arrays.setdefault(f"{proj}_scales", []).append(scales) + errors[proj].append(wrel) + tmp = dest.with_suffix(".tmp.npz") + np.savez(tmp, code_bits=np.asarray(16), **{k: np.stack(v) for k, v in arrays.items()}) + tmp.rename(dest) + row = {"layer": layer, "experts": len(errors["w1"]), "seconds": round(time.time() - started, 1), + **{f"{proj}_wrel_mse": float(np.mean(v)) for proj, v in errors.items()}} + with (args.out / "metrics.jsonl").open("a") as f: + f.write(json.dumps(row) + "\n") + print(json.dumps(row), flush=True) + + +def router(args) -> None: + import stream + + ckpt, rows = stream.HFCheckpoint(args.hf_path), {} + for layer in range(args.n_layers): + for name, tensor in ckpt.load(f"layers.{layer}.ffn.gate.").items(): + rows.setdefault(name, []).append(tensor.float().numpy()) # bf16 widens to fp32 without loss + np.savez(args.out, **{name: np.stack(v) for name, v in rows.items()}) + print({name: np.stack(v).shape for name, v in rows.items()}) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="command", required=True) + for name in ("importance", "check", "fit", "router"): + s = sub.add_parser(name) + s.add_argument("--hf-path", type=Path) + s.add_argument("--survivors", type=Path) + s.add_argument("--stats", type=Path) + s.add_argument("--codes", type=Path) + s.add_argument("--importance", type=Path) + s.add_argument("--out", type=Path) + s.add_argument("--layers", default="all") + s.add_argument("--iterations", type=int, default=8) + s.add_argument("--experts", type=int, default=0, help="debug: only the first N experts of a layer") + s.add_argument("--rows", type=int, default=64) + s.add_argument("--n-layers", type=int, default=40) + s.add_argument("--device", default="cuda" if torch.cuda.is_available() else "mps") + args = p.parse_args() + with torch.no_grad(): + {"importance": importance, "check": check, "fit": fit, "router": router}[args.command](args) + + +if __name__ == "__main__": + main() diff --git a/scripts/runpod/v41_reap/local_eval.py b/scripts/runpod/v41_reap/local_eval.py new file mode 100644 index 0000000..de71c18 --- /dev/null +++ b/scripts/runpod/v41_reap/local_eval.py @@ -0,0 +1,649 @@ +"""Score router-masked variants of a local REAP checkpoint against the pod's teacher logits. + +Every variant keeps a subset of the checkpoint's experts in each layer (any count per layer, +so non-uniform budgets work) by setting the dropped experts' gate bias to -inf, which is what +``collect.py --survivors`` did on the pod and equals physically slicing them out. All variants +go through one layer-streamed pass on CPU, through DeepSeek's reference code with the +``torch_kernel`` stand-ins. + +Sessions are cut to their first ``--max-tokens`` tokens. The model is causal, so the logits at +the kept positions are the ones the pod computed on the whole record; the teacher's top-K rows +for those positions are compared directly. + + python local_eval.py --hf-path ~/Models/v41-reap50 \\ + --checkpoint-survivors ~/Models/v41-reap-run/stats/survivors-192.json \\ + --pack ~/Models/v41-reap-run/pack/v41-pack.jsonl --stats ~/Models/v41-reap-run/stats \\ + --sessions-file ~/Models/v41-reap-run/mac-eval-sessions.json --max-tokens 1024 \\ + --variant reap50=all --variant reap75=~/Models/v41-reap-run/stats/survivors-96.json \\ + --out ~/Models/v41-reap-run/evals/baseline + +A variant spec is ``all`` (the checkpoint as is) or a survivors JSON in the base model's expert +ids (``{"survivors": {"": [ids...]}}``), each layer a subset of the checkpoint's. + +``--expert-override`` swaps every routed expert's weights for bf16 ones before the pass: +``fp4-bf16`` decodes the checkpoint's FP4 weights exactly (the control), a directory of +``keep_quant.py`` layer files decodes KEEP's codes (E8P or 1-bit E8), and +``HOT_DIR@SET.json+REST_DIR`` takes the experts in a survivors set from one directory and the +rest from another (a per-expert precision tier). Overridden experts run as plain bf16 +GEMMs without FP8 activation quantization, as KEEP's runtime does. + +``--capture-imatrix OUT.npz`` measures KEEP's importance diagonal on this checkpoint instead of +scoring: one variant, any sessions (``name`` or ``name#record``, calibration records allowed). +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import resource +import sys +import time +from pathlib import Path + +import numpy as np + +KIT = Path(__file__).resolve().parent +sys.path.insert(0, str(KIT)) +sys.path.insert(0, str(KIT.parents[2] / "src")) + +from local_parity import rss_guard # noqa: E402 + + +def local_masks(spec: str, ckpt_survivors: dict[str, list[int]]) -> dict[int, np.ndarray] | None: + """Per layer, a bool array over the checkpoint's local experts: True = dropped.""" + + if spec == "all": + return None + wanted = json.loads(Path(spec).expanduser().read_text())["survivors"] + masks = {} + for layer, have in ckpt_survivors.items(): + have = sorted(have) + keep = set(wanted[layer]) + assert keep <= set(have), f"layer {layer}: variant keeps experts the checkpoint dropped" + masks[int(layer)] = np.asarray([e not in keep for e in have]) + return masks + + +def gate_masked(gate, x, allowed, linear): + """``Gate.forward`` with a per-token candidate set: ``allowed`` [n, n_experts] bool. + + The router picks its top-k among the allowed experts and normalizes over those, which is what + a router masked to a pruned set does, here with a set that changes from token to token. A row + with fewer than top-k allowed experts gives its surplus slots zero weight.""" + + import torch + import torch.nn.functional as F + + scores = linear(x.float(), gate.weight.float()) / gate.gate_temp + if gate.score_func == "softmax": + scores = scores.softmax(dim=-1) + elif gate.score_func == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + indices = (scores + gate.bias).masked_fill(~allowed, float("-inf")).topk(gate.topk, dim=-1)[1] + weights = scores.gather(1, indices) * allowed.gather(1, indices) + if gate.norm_topk_prob and gate.topk > 1: + weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) + return weights * gate.route_scale, indices + + +def decode_fp4(weight, scale): + """bf16 [out, in] of one FP4 projection (E2M1 codes, one scale per 32 inputs).""" + import torch_kernel + + return (torch_kernel.decode_e2m1(weight).unflatten(-1, (-1, 32)) * scale.float()[..., None]).flatten(-2).bfloat16() + + +def hot_set(order, n: int, have) -> set[int]: + """The first ``n`` experts of ``order`` (hottest first) that are among ``have``.""" + have = set(have) + return set([int(e) for e in order if int(e) in have][:n]) + + +def expert_override(spec: str, n_experts: int, ckpt_survivors: dict): + """``layer -> {(expert, proj): bf16 [out, in]}`` from FP4 (``fp4-bf16``) or E8P codes.""" + + import torch + + if spec == "fp4-bf16": + def from_state(layer: int, sd: dict) -> dict: + out = {} + for e in range(n_experts): + for proj in ("w1", "w2", "w3"): + out[(e, proj)] = decode_fp4(sd[f"ffn.experts.{e}.{proj}.weight"], sd[f"ffn.experts.{e}.{proj}.scale"]) + return out + return from_state + + from keep.vq.e8 import e8_1bit_grid, e8p_full_grid + + grids = {16: torch.as_tensor(e8p_full_grid().astype(np.float32)), 8: torch.as_tensor(e8_1bit_grid().astype(np.float32))} + + def from_dir(root: Path): + def from_codes(layer: int, sd: dict, only: set[int] | None = None) -> dict: + out = {} + with np.load(root / f"layer-{layer:02d}.npz") as z: + grid = grids[int(z["code_bits"]) if "code_bits" in z.files else 16] + for proj in ("w1", "w2", "w3"): + codes = torch.as_tensor(z[f"{proj}_codes"].astype(np.int64)) # [E, out, in/8] + scales = torch.as_tensor(z[f"{proj}_scales"].astype(np.float32)) # [E, out, in/g] + assert codes.shape[0] == n_experts, (layer, proj, codes.shape) + for e in range(n_experts): + if only is not None and e not in only: + continue + vec = grid[codes[e]] # [out, in/8, 8] + words = vec.shape[1] // scales.shape[-1] + vec = vec * scales[e].repeat_interleave(words, dim=-1)[..., None] + out[(e, proj)] = vec.flatten(-2).bfloat16() + return out + return from_codes + + # "HOT_DIR@HOT_SET.json+REST_DIR": the hot set's experts from one source, the rest from another + if "@" in spec: + hot_spec, _, rest = spec.partition("@") + hot_set_path, _, rest_spec = rest.partition("+") + hot = from_dir(Path(hot_spec).expanduser()) + other = from_dir(Path(rest_spec).expanduser()) + hot_ids = json.loads(Path(hot_set_path).expanduser().read_text())["survivors"] + + def mixed(layer: int, sd: dict) -> dict: + have = sorted(ckpt_survivors[str(layer)]) + is_hot = {i for i, e in enumerate(have) if e in set(hot_ids[str(layer)])} + out = hot(layer, sd, is_hot) + out.update(other(layer, sd, set(range(n_experts)) - is_hot)) # decode each expert once + return out + return mixed + + return from_dir(Path(spec).expanduser()) + + +def layer_codes(spec: str, layer: int, have: list[int]) -> dict: + """One layer's code arrays, a row per expert of ``have`` (the checkpoint's experts, ascending). + + ``spec`` is a ``keep_quant.py`` directory holding exactly those experts, or + ``HOT_DIR@SET.json+REST_DIR``: the set's experts in one directory and the rest in the other, + each in ascending expert order.""" + if "@" not in spec: + return dict(np.load(Path(spec).expanduser() / f"layer-{layer:02d}.npz")) + hot_dir, _, rest = spec.partition("@") + set_path, _, rest_dir = rest.partition("+") + hot = set(json.loads(Path(set_path).expanduser().read_text())["survivors"][str(layer)]) + stored = [e for e in have if e in hot] + [e for e in have if e not in hot] # the two files back to back + row = {e: j for j, e in enumerate(stored)} + parts = [np.load(Path(d).expanduser() / f"layer-{layer:02d}.npz") for d in (hot_dir, rest_dir)] + assert sum(z["w1_codes"].shape[0] for z in parts) == len(have) and parts[0]["w1_codes"].shape[0] == len(hot & set(have)) + out = {k: np.concatenate([z[k] for z in parts])[[row[e] for e in have]] + for k in parts[0].files if k.endswith(("_codes", "_scales"))} + bits = {int(z["code_bits"]) if "code_bits" in z.files else 16 for z in parts} + assert len(bits) == 1, bits + return {**out, "code_bits": bits.pop()} + + +def residual_cache(ind: np.ndarray, wts: np.ndarray, k: int, tau: float, bounds: list[int], + warm: np.ndarray | None, counts: np.ndarray, policy: str = "lru", + trace: list | None = None, resident: np.ndarray | None = None) -> np.ndarray: + """Which (token, slot) pairs may run base + residual, under an LRU of ``k`` resident residuals. + + ``ind``/``wts``: [T, top_k] routed experts and gate weights, the records of ``bounds`` back to + back. A slot runs high precision only if its expert's residual is already resident. After each + token, every slot with gate >= ``tau`` admits (or refreshes) its residual: an async SSD fetch + that serves the next token on, so decode never stalls. Each record starts from ``warm`` (or + empty). ``counts[e, 4]`` accumulates admissions, i.e. residual fetches; ``trace`` collects + them as (record, token in record, expert). + + ``policy="oracle"`` is the upper bound for any cache of ``k``: each record's hot set is its own + top ``k`` experts by routed gate weight, known in advance and never changed (no admissions). + + ``resident`` ([T, n_experts] bool, zeroed by the caller) receives the resident set each token + sees, for callers that re-route a miss among resident experts.""" + + from collections import OrderedDict + + hi = np.zeros(ind.shape, dtype=bool) + start = 0 + for rec, n in enumerate(bounds): + if policy == "oracle": + mass = np.bincount(ind[start:start + n].ravel(), weights=wts[start:start + n].ravel(), minlength=counts.shape[0]) + hot = np.argsort(-mass, kind="stable")[:k] + hi[start:start + n] = np.isin(ind[start:start + n], hot[mass[hot] > 0]) + if resident is not None: + resident[start:start + n, hot] = True + start += n + continue + lru = OrderedDict((int(e), None) for e in (warm if warm is not None else [])) + for t in range(start, start + n): + row = ind[t].tolist() + hi[t] = [e in lru for e in row] + if resident is not None: + resident[t, list(lru)] = True + for s, e in enumerate(row): + if wts[t, s] >= tau: + if e in lru: + lru.move_to_end(e) + else: + counts[e, 4] += 1 + if trace is not None: + trace.append((rec, t - start, e)) + lru[e] = None + if len(lru) > k: + lru.popitem(last=False) + start += n + return hi + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--hf-path", required=True, type=Path) + p.add_argument("--checkpoint-survivors", type=Path, + help="the survivors JSON the checkpoint was exported from; omit for an unpruned checkpoint") + p.add_argument("--device", default="cpu", help="cpu, or cuda on a pod") + p.add_argument("--router", type=Path, help="safetensors with every router row of the unpruned model: the pool is " + "all its experts, taken from --dynamic codes alone (HOT_DIR@SET.json+REST_DIR), on a pruned checkpoint") + p.add_argument("--pack", required=True, type=Path) + p.add_argument("--stats", required=True, type=Path) + p.add_argument("--sessions-file", required=True, type=Path) + p.add_argument("--max-tokens", type=int, default=1024) + p.add_argument("--variant", action="append", required=True, metavar="NAME=SPEC") + p.add_argument("--out", required=True, type=Path) + p.add_argument("--max-rss-gb", type=float, default=64.0) + p.add_argument("--threads", type=int, default=os.cpu_count()) + p.add_argument("--expert-override", help="'fp4-bf16' or a keep_quant.py output directory") + p.add_argument("--capture-imatrix", type=Path, help="measure the importance diagonal instead of scoring") + p.add_argument("--dynamic", help="per-token tiers. BASE_DIR+RESIDUAL_DIR: base alone or base + nested residual. " + "BASE_DIR|HI_DIR: base alone or a separate high-precision copy (a hot expert then holds both). " + "With '|', a side may be HOT_DIR@SET.json+REST_DIR (the set's experts from one directory, the rest " + "from the other), and BASE may be empty when every variant has --miss") + p.add_argument("--hi-proj", action="append", default=[], metavar="NAME=w2|w1w3", + help="nested only, per variant: the high tier adds the residual for these projections only") + p.add_argument("--cache-policy", action="append", default=[], metavar="NAME=lru|oracle", + help="with --cache-k: lru (default) or oracle (each record's own top-K, the upper bound)") + p.add_argument("--trace", action="append", default=[], metavar="NAME", + help="with --cache-k: write this variant's residual admissions (layer, record, token, expert)") + p.add_argument("--tau", action="append", default=[], metavar="NAME=VALUE", + help="with --dynamic, per variant: a slot runs base + residual when its gate weight >= VALUE") + p.add_argument("--tau-drop", action="append", default=[], metavar="NAME=VALUE", + help="with --dynamic, per variant: a slot with gate weight < VALUE is skipped (rung 0 per token)") + p.add_argument("--miss", action="append", default=[], metavar="NAME=drop|reroute", + help="with --cache-k and BASE|HI codes: no base copy. A slot whose expert is not resident is " + "skipped (drop), or the router picks its top-k among resident experts (reroute)") + p.add_argument("--cache-k", action="append", default=[], metavar="NAME=K", + help="with --dynamic: hi tier only from an LRU of K resident residuals per layer, admitted at gate >= tau") + p.add_argument("--hot", action="append", default=[], metavar="NAME=H", + help="with --miss: the H hottest experts (--cache-warm order) that the checkpoint holds run from " + "its FP4 weights while resident, the other residents from the codes") + p.add_argument("--hot-tier", default="fp4", metavar="fp4|DIR", + help="with --hot: what a hot expert runs. fp4, or its codes plus the residual in DIR, a keep_quant.py " + "--residual-of fit over the checkpoint's experts") + p.add_argument("--cache-warm", type=Path, help="imatrix npz (reap__gate_norm_sum): warm-start each cache with the top-K") + p.add_argument("--sidecar", type=Path, help="a sidecar_fit.py output directory, added to every variant's MoE output") + p.add_argument("--sidecar-rank", type=int, help="use only the top-k sidecar components (stored largest first)") + args = p.parse_args(argv) + rss_guard(args.max_rss_gb) + + import torch + from transformers import AutoTokenizer + + import collect + import stream + import torch_kernel + + if args.checkpoint_survivors: + prune = json.loads((args.hf_path / "keep-prune.json").read_text()) + sha = hashlib.sha256(args.checkpoint_survivors.read_bytes()).hexdigest() + assert sha == prune["survivors_sha256"], "--checkpoint-survivors is not what this checkpoint was exported from" + ckpt_survivors = json.loads(args.checkpoint_survivors.read_text())["survivors"] + else: # an unpruned checkpoint: every expert, under its own number + shape = json.loads((args.hf_path / "inference" / "config.json").read_text()) + ckpt_survivors = {str(l): list(range(shape["n_routed_experts"])) for l in range(shape["n_layers"])} + wide = None + if args.router: # the checkpoint's own experts are never run: every variant re-routes over the codes + from safetensors.torch import load_file + + assert args.dynamic and {v.partition("=")[0] for v in args.miss} == {v.partition("=")[0] for v in args.variant}, \ + "--router needs --dynamic codes and --miss for every variant" + wide, loaded = load_file(args.router), ckpt_survivors + ckpt_survivors = {l: list(range(wide[f"layers.{l}.ffn.gate.bias"].numel())) for l in loaded} + variants = {} + for spec in args.variant: + name, _, path = spec.partition("=") + masks = local_masks(path, ckpt_survivors) + variants[name] = None if masks is None else {l: torch.as_tensor(m) for l, m in masks.items()} + + torch_kernel.exact_quant = True + torch.set_num_threads(args.threads) + ref_dir = args.hf_path / "inference" + dev = args.device + # On a GPU the reference's own kernels run, as in collect.py; elsewhere the torch stand-ins. + ref = stream.import_reference(ref_dir, kernel=torch_kernel if dev == "cpu" else None) + torch.set_default_dtype(torch.bfloat16) + if dev != "cpu": + torch.cuda.memory._set_allocator_settings("expandable_segments:True") + torch.set_default_device(dev) + torch.manual_seed(33377335) + model_args = ref.ModelArgs(**json.loads((ref_dir / "config.json").read_text())) + model_args.max_batch_size = 1 + + sessions = collect.load_sessions(args.pack) + names = json.loads(args.sessions_file.read_text()) + records = {} + for name in names: + session, _, idx = name.partition("#") + recs = sessions[session] + if not idx: + assert len(recs) == 1, f"{session} has {len(recs)} records; name one as {session}#N" + record = recs[int(idx or 0)] + record = dict(record, token_ids=record["token_ids"][: args.max_tokens]) + records[name] = record + model_args.max_seq_len = max(len(r["token_ids"]) for r in records.values()) + + ckpt = stream.HFCheckpoint(args.hf_path) + layout = ref.EngramLayout.from_args(model_args) + with torch.device(dev): + hash_state = ref.NgramHashState(model_args, layout, AutoTokenizer.from_pretrained(args.hf_path)) + engram_rows = {l: stream.EngramRows.from_checkpoint(ckpt, l, ref.fp8_block_size) for l in model_args.engram_layer_ids} + streamer = stream.LayerStreamer(ref, model_args, ckpt.load, engram_rows, hash_state, dev) + if wide: + build_narrow = streamer.build_block + + def build_wide(layer, sd): + block = build_narrow(layer, sd) + gate, ids = block.ffn.gate, loaded[str(layer)] + for f in ("weight", "bias", "bias_vl"): + if getattr(gate, f, None) is not None: + rows = wide[f"layers.{layer}.ffn.gate.{f}"].to(getattr(gate, f)) + assert torch.equal(rows[ids], getattr(gate, f).data), f"layer {layer}: router {f} is not this checkpoint's" + setattr(gate, f, torch.nn.Parameter(rows, requires_grad=False)) + return block + + streamer.build_block = build_wide + state = collect.install_hooks(streamer, top_k=2048, rows_per_slice=256) + n_local = len(next(iter(ckpt_survivors.values()))) + + dyn_counts: dict[str, np.ndarray] = {} + if args.dynamic: # per-token precision: each (token, expert) slot picks its tier from its gate weight + import torch.nn.functional as F + + from keep.vq.e8 import e8_1bit_grid, e8p_full_grid + + copy_mode = "|" in args.dynamic # the high tier is its own copy, not base + residual + args_cache_k_names = {s.partition("=")[0] for s in args.cache_k} + tier_specs = [s for s in args.dynamic.split("|" if copy_mode else "+", 1) if s] # base, then high + hi_proj = {name: {"w2": ("w2",), "w1w3": ("w1", "w3")}[v] for name, _, v in (s.partition("=") for s in args.hi_proj)} + assert not (copy_mode and hi_proj), "--hi-proj needs nested codes" + policies = {name: v for name, _, v in (s.partition("=") for s in args.cache_policy)} + traces: dict[str, list] = {name: [] for name in args.trace} + taus = {name: float(t) for name, _, t in (s.partition("=") for s in args.tau)} + tau_drop = {name: float(t) for name, _, t in (s.partition("=") for s in args.tau_drop)} + miss = {name: mode for name, _, mode in (s.partition("=") for s in args.miss)} + assert all(m in ("drop", "reroute") and copy_mode and v in args_cache_k_names for v, m in miss.items()), \ + "--miss takes drop or reroute, and needs BASE|HI codes and --cache-k for the same variant" + assert set(taus) == set(variants), "--tau needs one NAME=VALUE per --variant" + assert len(tier_specs) == 2 or set(miss) == set(variants), "codes without a base need --miss for every variant" + grids = {bits: torch.as_tensor(grid().astype(np.float32), device=dev) for bits, grid in ((16, e8p_full_grid), (8, e8_1bit_grid))} + cache: dict = {"layer": None} + n_layers_ = int(model_args.n_layers) + cache_k = {name: int(k) for name, _, k in (s.partition("=") for s in args.cache_k)} + warm = None + if args.cache_warm: # warm start: each layer's top-K experts by routed mass (the static hot set) + warm = np.argsort(-np.load(args.cache_warm)["reap__gate_norm_sum"], axis=1, kind="stable") + hot_n = {name: int(h) for name, _, h in (s.partition("=") for s in args.hot)} + assert set(hot_n) <= set(miss) and (warm is not None or not hot_n), "--hot needs --miss and --cache-warm" + own = loaded if wide else ckpt_survivors # pool numbers of the checkpoint's own experts, in its order + hot_slots = dict.fromkeys(variants, 0) + for v in variants: # [layer, expert, (lo slots, hi slots, lo mass, hi mass, admissions, dropped slots, dropped mass)] + dyn_counts[v] = np.zeros((n_layers_, n_local, 7)) + + def simulate_cache(layer, ind, wts, k, tau, counts, resident=None): + v, found = state["variant"], [] + hi = residual_cache(ind, wts, k, tau, state["bounds"], None if warm is None else warm[layer, :k], counts, + policy=policies.get(v, "lru"), trace=found if v in traces else None, resident=resident) + if v in traces: + traces[v].extend((layer, *row) for row in found) + return hi + + def decode(z, proj: str, e: int): + """float32 weights of expert ``e`` of one keep_quant layer file.""" + vec = grids[int(z.get("code_bits", 16))][torch.as_tensor(z[f"{proj}_codes"][e].astype(np.int64), device=dev)] + s = torch.as_tensor(z[f"{proj}_scales"][e].astype(np.float32), device=dev) # vec: [out, in/8, 8] + return (vec * s.repeat_interleave(vec.shape[1] // s.shape[-1], dim=-1)[..., None]).flatten(-2) + + def tier_weights(layer: int, e: int) -> tuple[dict, dict]: + """(base, base + residual) bf16 weights of local expert ``e``, decoded once per layer.""" + if cache["layer"] != layer: + cache.clear() + cache["layer"] = layer + cache["z"] = [layer_codes(s, layer, sorted(ckpt_survivors[str(layer)])) for s in tier_specs] + if e not in cache: + tiers = [{proj: decode(z, proj, e) for proj in ("w1", "w2", "w3")} for z in cache["z"]] + lo = {p: w.bfloat16() for p, w in tiers[0].items()} + hi = lo if len(tiers) == 1 else {p: (tiers[1][p] if copy_mode else lo_w + tiers[1][p]).bfloat16() + for p, lo_w in tiers[0].items()} + cache[e] = (lo, hi) + return cache[e] + + def run_expert(x, w, tier, limit): + gate, up = F.linear(x, tier["w1"]).float(), F.linear(x, tier["w3"]).float() + if limit > 0: + up, gate = up.clamp(-limit, limit), gate.clamp(max=limit) + return F.linear((w * (F.silu(gate) * up)).to(x.dtype), tier["w2"]) + + def dynamic_forward(self, x, image_mask=None): + shape, layer = x.size(), self.layer_id + x = x.view(-1, self.dim) + weights, indices = self.gate(x, None) + tau, counts = taus[state["variant"]], dyn_counts[state["variant"]][layer] + y = torch.zeros_like(x, dtype=torch.float32) + limit = self.shared_experts.swiglu_limit + k = cache_k.get(state["variant"]) + on_miss = miss.get(state["variant"]) + resident = np.zeros((x.size(0), self.gate.weight.size(0)), dtype=bool) if on_miss == "reroute" else None + hi_all = None if k is None else torch.as_tensor( + simulate_cache(layer, indices.cpu().numpy(), weights.float().cpu().numpy(), k, tau, counts, resident), + device=x.device) + if resident is not None: # demand (above) drove the cache; the tokens run on what is resident + weights, indices = gate_masked(self.gate, x, torch.as_tensor(resident, device=x.device), ref.linear) + hi_all = weights > 0 + for i in torch.unique(indices).tolist(): + idx, top = torch.where(indices == i) + w = weights[idx, top] + hi_mask = w >= tau if hi_all is None else hi_all[idx, top] + drop = w < tau_drop.get(state["variant"], 0.0) # rung 0 per token: the slot is skipped + if on_miss: + drop = drop | ~hi_mask # no base copy: a slot that is not resident is not run + counts[i, 5] += int(drop.sum()) + counts[i, 6] += float(w[drop].sum()) + hi_mask = hi_mask & ~drop + lo_w, hi_w = tier_weights(layer, i) + if hot_n.get(state["variant"]): # residents of the hot set run at --hot-tier precision + if "hot" not in cache: + cache["hot"] = {n: hot_set(warm[layer], n, own[str(layer)]) for n in set(hot_n.values())} + if i in cache["hot"][hot_n[state["variant"]]]: + if ("hot", i) not in cache: + j = own[str(layer)].index(i) # its place in the checkpoint, and in a fit over it + if args.hot_tier == "fp4": + mine = self.experts[j] + cache["hot", i] = {p: decode_fp4(getattr(mine, p).weight, getattr(mine, p).scale) for p in lo_w} + else: + if "residual" not in cache: + cache["residual"] = np.load(Path(args.hot_tier).expanduser() / f"layer-{layer:02d}.npz") + assert cache["residual"]["w1_codes"].shape[0] == len(own[str(layer)]) + cache["hot", i] = {p: (decode(cache["z"][0], p, i) + decode(cache["residual"], p, j)).bfloat16() + for p in lo_w} + hi_w = cache["hot", i] + hot_slots[state["variant"]] += int(hi_mask.sum()) + only = hi_proj.get(state["variant"]) + if only is not None: # a finer rung: the residual of some projections only + hi_w = {p: hi_w[p] if p in only else lo_w[p] for p in lo_w} + for mask, tier, col in ((~hi_mask & ~drop, lo_w, 0), (hi_mask, hi_w, 1)): + if mask.any(): + y[idx[mask]] += run_expert(x[idx[mask]], w[mask, None], tier, limit).float() + counts[i, col] += int(mask.sum()) + counts[i, col + 2] += float(w[mask].sum()) + y += self.shared_experts(x).float() + return y.type_as(x).view(shape) + + ref.MoE.forward = dynamic_forward + + if args.sidecar: # sidecar_fit.py output: y += ([x, 1] @ U) @ Vt on every MoE call + sidecars = {} + for layer in range(int(model_args.n_layers)): + if not (args.sidecar / f"layer-{layer:02d}.npz").exists(): + continue # a fit still in progress: the layers it has are valid without the rest + with np.load(args.sidecar / f"layer-{layer:02d}.npz") as z: + U, Vt = z["U"][:, : args.sidecar_rank or None], z["Vt"][: args.sidecar_rank or None] + sidecars[layer] = (torch.as_tensor(U), torch.as_tensor(Vt)) + moe_forward = ref.MoE.forward + + def with_sidecar(self, x, image_mask=None): + y = moe_forward(self, x, image_mask) + if self.layer_id not in sidecars: + return y + U, Vt = sidecars[self.layer_id] + xf = x.reshape(-1, x.size(-1)).float() + corr = (xf @ U[:-1] + U[-1]) @ Vt + return (y.reshape(-1, y.size(-1)).float() + corr).type_as(y).view(y.shape) + + ref.MoE.forward = with_sidecar + + if args.expert_override: + decode = expert_override(args.expert_override, n_local, ckpt_survivors) + decoded: dict[int, dict] = {} + load_layer, build_block = streamer.load_layer, streamer.build_block + + def load_with_override(prefix, skip=()): # runs in the prefetch thread + sd = load_layer(prefix, skip) + if prefix.startswith("layers."): + layer = int(prefix.split(".")[1]) + decoded[layer] = decode(layer, sd) + return sd + + def build_with_override(layer, sd): + block = build_block(layer, sd) + for (e, proj), w in decoded.pop(layer).items(): + lin = getattr(block.ffn.experts[e], proj) + lin.weight = torch.nn.Parameter(w, requires_grad=False) + if hasattr(lin, "scale"): + del lin.scale + return block + + streamer.load_layer, streamer.build_block = load_with_override, build_with_override + + capture = None + if args.capture_imatrix: + assert len(variants) == 1, "--capture-imatrix takes one variant" + capture = collect.SessionStats(torch=torch, n_layers=int(model_args.n_layers), n_experts=n_local, + dims={"hidden": int(model_args.dim), "down": int(model_args.moe_inter_dim)}, device=dev) + + state["bounds"] = [len(r["token_ids"]) for r in records.values()] # one MoE call holds a variant's records in order + items = { + (v, n): stream.Item(tokens=torch.as_tensor([r["token_ids"]], dtype=torch.long, device=dev), group=(v, None)) + for v in variants + for n, r in records.items() + } + + def before_moe(layer: int, block, key) -> None: + bias = block.ffn.gate.bias + if not hasattr(block, "full_bias"): + block.full_bias = bias.data.clone() + bias.data.copy_(block.full_bias) + masks = variants[key[0]] + if masks is not None: + bias.data[masks[layer].to(bias.device)] = float("-inf") + state["stats"] = capture + state["variant"] = key[0] + + t0 = time.time() + with torch.inference_mode(): + streamer.forward(list(items.values()), before_moe) + forward_s = time.time() - t0 + state["stats"] = None + + if capture is not None: + arrays = capture.arrays(list(range(int(model_args.n_layers)))) + args.capture_imatrix.parent.mkdir(parents=True, exist_ok=True) + np.savez(args.capture_imatrix, + importance__hidden=arrays["importance_sum__hidden"], importance__down=arrays["importance_sum__down"], + route_count=arrays["route_count__hidden"], reap__gate_norm_sum=arrays["reap__gate_norm_sum"], + reap__route_count=arrays["reap__route_count"], + sessions=np.asarray(names), tokens=np.asarray(sum(len(r["token_ids"]) for r in records.values()))) + print(json.dumps({"forward_seconds": round(forward_s, 1), "tokens": int(sum(len(r["token_ids"]) for r in records.values())), + "peak_rss_gb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**30, 2)}), flush=True) + return 0 + + args.out.mkdir(parents=True, exist_ok=True) + summary: dict = {"forward_seconds": round(forward_s, 1), "max_tokens": args.max_tokens, "variants": {}} + for name, record in records.items(): + teacher = np.load(args.stats / "eval" / "sessions" / f"{name}.npz") + positions = collect.target_positions(record) + rows = slice(0, positions.size) + assert np.array_equal(positions, teacher["positions"][rows]), name + targets = np.asarray(record["token_ids"], dtype=np.int64)[positions + 1] + teacher_nll = (teacher["logsumexp"][rows] - teacher["target_logit"][rows]).astype(np.float64) + state["positions"] = torch.as_tensor(positions, device=dev) + state["targets"] = torch.as_tensor(targets, device=dev) + state["teacher"] = { + "ids": torch.as_tensor(teacher["topk_logit_ids"][rows].astype(np.int64), device=dev), + "logprob": torch.as_tensor(teacher["topk_logit_values"][rows].astype(np.float32) - teacher["logsumexp"][rows, None], + device=dev), + "targets": state["targets"], + } + for v in variants: + state["captured"] = {k: [] for k in collect.CAPTURE_KEYS} + with torch.inference_mode(): + streamer.head(streamer.final_hidden(items[(v, name)])) + cap = {k: np.concatenate(x) for k, x in state["captured"].items() if x} + out = { + "positions": positions.astype(np.int32), + "kl": cap["kl_on_teacher_topk"], + "nll_delta": (cap["student_nll"] - teacher_nll).astype(np.float32), + "top1": cap["top1_agree"], + } + (args.out / v).mkdir(exist_ok=True) + np.savez(args.out / v / f"{name}.npz", **out) + s = summary["variants"].setdefault(v, {}) + s[name] = {"domain": record["domain"], "positions": int(positions.size), + **{k: float(np.mean(out[k])) for k in ("kl", "nll_delta", "top1")}} + state["teacher"] = state["positions"] = state["targets"] = None + + for v, per in summary["variants"].items(): + doms: dict[str, list] = {} + for row in per.values(): + doms.setdefault(row["domain"], []).append(row) + by_dom = {d: {k: float(np.average([r[k] for r in rs], weights=[r["positions"] for r in rs])) for k in ("kl", "nll_delta", "top1")} + for d, rs in sorted(doms.items())} + per["_domains"] = by_dom + per["_domain_mean"] = {k: float(np.mean([x[k] for x in by_dom.values()])) for k in ("kl", "nll_delta", "top1")} + print(v, json.dumps(per["_domain_mean"]), flush=True) + for v, c in dyn_counts.items(): # per-token precision: how often (and with how much gate mass) the hi tier ran + np.save(args.out / v / "tier-counts.npy", c) + tot = c.sum(axis=(0, 1)) + if v in traces: + t = np.asarray(traces[v], dtype=np.int32).reshape(-1, 4) + np.savez(args.out / v / "admissions.npz", layer=t[:, 0], record=t[:, 1], token=t[:, 2], expert=t[:, 3], + records=np.asarray(list(records)), tokens=np.asarray(state["bounds"])) + kept = n_local if variants[v] is None else float(np.mean([int((~m).sum()) for m in variants[v].values()])) + # hi_units: memory (and bytes fetched) the high tier adds per expert, in units of one 1-bit expert + hi_units = 2.0 if copy_mode else len(hi_proj.get(v, ("w1", "w2", "w3"))) / 3 + summary["variants"][v]["_tiers"] = {"tau": taus[v], "cache_k": cache_k.get(v), + "policy": policies.get(v, "lru") if v in cache_k else None, + "copy_mode": copy_mode, "hi_units": hi_units, + "n_bases": 0 if v in miss else kept, "pool": kept, "miss": miss.get(v), + "hot": hot_n.get(v, 0), "hot_tier": "fp4" if args.hot_tier == "fp4" else "nested", + "hot_slot_share": hot_slots[v] / max(tot[:2].sum() + tot[5], 1), + "tau_drop": tau_drop.get(v, 0.0), + "hi_slot_share": float(tot[1] / (tot[:2].sum() + tot[5])), + "drop_slot_share": float(tot[5] / (tot[:2].sum() + tot[5])), + "hi_mass_share": float(tot[3] / (tot[2:4].sum() + tot[6])), + "drop_mass_share": float(tot[6] / (tot[2:4].sum() + tot[6])), + "admissions_per_token_layer": float(tot[4] / ((tot[:2].sum() + tot[5]) / int(model_args.n_activated_experts)) / int(model_args.n_layers))} + print(v, json.dumps(summary["variants"][v]["_tiers"]), flush=True) + summary["peak_rss_gb"] = round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**30, 2) + (args.out / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps({k: summary[k] for k in ("forward_seconds", "peak_rss_gb")}), flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/local_parity.py b/scripts/runpod/v41_reap/local_parity.py new file mode 100644 index 0000000..e5343b5 --- /dev/null +++ b/scripts/runpod/v41_reap/local_parity.py @@ -0,0 +1,140 @@ +"""Run an exported REAP checkpoint on CPU and score it against the pod's teacher, per position. + +The Runpod eval scored each candidate by masking pruned experts out of the full model's +router (``collect.py --survivors``). This script runs the exported checkpoint itself, which +has those experts physically removed and the router sliced, through the same layer-streamed +reference code with the CPU kernel stand-ins in ``torch_kernel.py``. It reads the teacher's +top-K logits from the pod's eval npz files, so the full model is never needed here, then +compares KL, NLL and top-1 agreement with the pod's masked-router numbers for the same +positions. + + python local_parity.py --hf-path ~/Models/v41-reap75 \\ + --pack ~/Models/v41-reap-run/pack/v41-pack.jsonl \\ + --stats ~/Models/v41-reap-run/stats --keep 96 \\ + --sessions eval-00054 eval-00100 --out ~/Models/v41-reap-run/parity-96.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import resource +import sys +import threading +import time +from pathlib import Path + +import numpy as np + +KIT = Path(__file__).resolve().parent +sys.path.insert(0, str(KIT)) + + +def rss_guard(limit_gb: float) -> None: + """Exit hard if this process's peak resident size passes ``limit_gb`` (Jack's memory budget).""" + + def watch() -> None: + while True: + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**30 # bytes on macOS + if peak > limit_gb: + print(f"peak RSS {peak:.1f} GB passed the {limit_gb} GB budget; exiting", flush=True) + os._exit(3) + time.sleep(2) + + threading.Thread(target=watch, daemon=True).start() + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--hf-path", required=True, type=Path, help="the exported REAP checkpoint") + p.add_argument("--pack", required=True, type=Path) + p.add_argument("--stats", required=True, type=Path, help="local copy of the REAP-stats dataset (eval/ and student-KEEP/)") + p.add_argument("--keep", required=True, help="the pod candidate to compare with, e.g. 96") + p.add_argument("--sessions", nargs="+", required=True) + p.add_argument("--out", required=True, type=Path) + p.add_argument("--max-rss-gb", type=float, default=24.0) + p.add_argument("--threads", type=int, default=os.cpu_count()) + args = p.parse_args(argv) + rss_guard(args.max_rss_gb) + + import torch + from transformers import AutoTokenizer + + import collect + import stream + import torch_kernel + + torch_kernel.exact_quant = True + torch.set_num_threads(args.threads) + ref_dir = args.hf_path / "inference" + ref = stream.import_reference(ref_dir, kernel=torch_kernel) + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(33377335) + model_args = ref.ModelArgs(**json.loads((ref_dir / "config.json").read_text())) + model_args.max_batch_size = 1 + + sessions = collect.load_sessions(args.pack) + records = {name: sessions[name][0] for name in args.sessions} + model_args.max_seq_len = max(len(r["token_ids"]) for r in records.values()) + + ckpt = stream.HFCheckpoint(args.hf_path) + layout = ref.EngramLayout.from_args(model_args) + hash_state = ref.NgramHashState(model_args, layout, AutoTokenizer.from_pretrained(args.hf_path)) + engram_rows = {l: stream.EngramRows.from_checkpoint(ckpt, l, ref.fp8_block_size) for l in model_args.engram_layer_ids} + streamer = stream.LayerStreamer(ref, model_args, ckpt.load, engram_rows, hash_state, "cpu") + state = collect.install_hooks(streamer, top_k=2048, rows_per_slice=256) + + items = {name: stream.Item(tokens=torch.as_tensor([r["token_ids"]], dtype=torch.long), group=("student", None)) # one MoE pass for all + for name, r in records.items()} + t0 = time.time() + with torch.inference_mode(): + streamer.forward(list(items.values())) + forward_s = time.time() - t0 + + report: dict[str, dict[str, float]] = {} + for name, record in records.items(): + teacher = np.load(args.stats / "eval" / "sessions" / f"{name}.npz") + pod = np.load(args.stats / f"student-{args.keep}" / "eval" / "sessions" / f"{name}.npz") + positions = collect.target_positions(record) + assert np.array_equal(positions, teacher["positions"]) and np.array_equal(positions, pod["positions"]), name + targets = np.asarray(record["token_ids"], dtype=np.int64)[positions + 1] + state["positions"] = torch.as_tensor(positions) + state["targets"] = torch.as_tensor(targets) + state["teacher"] = { + "ids": torch.as_tensor(teacher["topk_logit_ids"].astype(np.int64)), + "logprob": torch.as_tensor(teacher["topk_logit_values"].astype(np.float32) - teacher["logsumexp"][:, None]), + "targets": state["targets"], + } + state["captured"] = {k: [] for k in collect.CAPTURE_KEYS} + with torch.inference_mode(): + streamer.head(streamer.final_hidden(items[name])) + cap = {k: np.concatenate(v) for k, v in state["captured"].items() if v} + teacher_nll = teacher["logsumexp"] - teacher["target_logit"] + row = {"domain": record["domain"], "positions": int(positions.size)} + for key, local, remote in ( + ("kl", cap["kl_on_teacher_topk"], pod["kl_on_teacher_topk"]), + ("nll_delta", cap["student_nll"] - teacher_nll, pod["student_nll"] - teacher_nll), + ("top1_agree", cap["top1_agree"].astype(np.float64), pod["top1_agree"].astype(np.float64)), + ): + row[f"{key}_local"] = float(np.mean(local)) + row[f"{key}_pod"] = float(np.mean(remote)) + row[f"{key}_mean_abs_diff"] = float(np.mean(np.abs(local - remote))) + report[name] = row + print(json.dumps({name: row}), flush=True) + + summary = { + "hf_path": str(args.hf_path), + "keep": args.keep, + "forward_seconds": round(forward_s, 1), + "peak_rss_gb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**30, 2), + "sessions": report, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps({k: v for k, v in summary.items() if k != "sessions"}), flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/pod.sh b/scripts/runpod/v41_reap/pod.sh new file mode 100755 index 0000000..2a2b432 --- /dev/null +++ b/scripts/runpod/v41_reap/pod.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +# DeepSeek-V4.1-Flash REAP run on RunPod. Runbook: +# docs/deepseek-v4-flash/2026-11-05-v41-runpod-reap-runbook.md +# +# pod.sh fetch CPU pod: download the pinned checkpoint, verify every LFS sha256 +# pod.sh pack CPU pod: build the calibration + eval pack from public datasets +# pod.sh gpu GPU pod (1-8 GPUs, 80 GB+ each): calibration sweep, reap-select, +# eval sweep (teacher + every candidate side by side), scores +# pod.sh prune KEEP CPU pod: write the pruned HF checkpoint for the Mac to pull +# pod.sh smoke GPU pod: the longest calibration record through all 40 layers; memory check +# pod.sh all one GPU pod does everything: fetch || pack, smoke, gpu, prune $PRUNE +# +# Everything lives under /workspace (a network volume, or the pod's own volume +# disk in single-pod mode). +set -euo pipefail + +MODEL_ID=deepseek-ai/DeepSeek-V4.1-Flash-1017 +REVISION=1bae0e1088993b9b82de08f59de5be5a34e42170 +W=${WORKSPACE:-/workspace} +KIT=$W/keep/scripts/runpod/v41_reap +HF=${HF:-$W/v41-hf} +OUT=$W/out +CANDIDATES=${CANDIDATES:-"256 192 128 96"} +PRUNE=${PRUNE:-"192"} # REAP 50%. REAP 75% (96) after the Mac pulls 192: 1.2 TB holds one +MAX_GPU_HOURS=${MAX_GPU_HOURS:-8} # hard stop for the GPU pod (1xH100 SXM secure $3.99/h) +# Tokens (x streams) resident per chunk. The hidden state carried between layers costs ~63 KB/token +# on the GPU: 524288 left ~65 GB in use and the 16K-record indexer's 16 GiB alloc OOMed (2026-11-07). +CHUNK_TOKENS=${CHUNK_TOKENS:-262144} +export PYTHONPATH=$W/keep/src +# ssh sessions miss the pod's env (RUNPOD_POD_ID, the pod-scoped RUNPOD_API_KEY): load it, never print it +if [[ -z "${RUNPOD_POD_ID:-}" && -r /etc/rp_environment ]]; then source /etc/rp_environment; fi + +watchdog() { # whatever happens next, the pod stops at MAX_GPU_HOURS; a failed stage stops it at once + [[ -n "${WATCHDOG_STARTED:-}" ]] && return + export WATCHDOG_STARTED=1 + ( sleep $((MAX_GPU_HOURS * 3600)); stop_pod "watchdog: ${MAX_GPU_HOURS}h limit" ) & + trap 'status=$?; [[ $status -ne 0 ]] && stop_pod "stage failed (status $status)"' EXIT +} + +log() { printf '[%s] %s\n' "$(date -u +%FT%TZ)" "$*" | tee -a "$W/pod.log"; } + +stop_pod() { + log "stopping pod: $*" + if [[ -n "${RUNPOD_POD_ID:-}" ]] && command -v runpodctl >/dev/null; then + runpodctl pod stop "$RUNPOD_POD_ID" || runpodctl stop pod "$RUNPOD_POD_ID" || log "runpodctl stop FAILED: stop the pod by hand" + else + log "no runpodctl/RUNPOD_POD_ID: stop the pod by hand" + fi +} + +fetch_small() { # tokenizer, configs, encoding/ and inference/: all pack and collect need besides weights + pip install -q "huggingface_hub>=0.34" "datasets>=3" tokenizers + hf download "$MODEL_ID" --revision "$REVISION" --local-dir "$HF" \ + --include "*.json" --include "encoding/*" --include "inference/*" + test -f "$HF/model.safetensors.index.json" || { log "fetch_small: no index json in $HF"; return 1; } +} + +fetch() { + fetch_small + hf download "$MODEL_ID" --revision "$REVISION" --local-dir "$HF" + python3 - "$HF" "$MODEL_ID" "$REVISION" <<'EOF' +import hashlib, sys +from pathlib import Path +from huggingface_hub import HfApi +root, repo, rev = Path(sys.argv[1]), sys.argv[2], sys.argv[3] +bad = 0 +for f in HfApi().model_info(repo, revision=rev, files_metadata=True).siblings: + path = root / f.rfilename + want = f.lfs.sha256 if f.lfs else None + if not path.exists() or (f.size is not None and path.stat().st_size != f.size): + print("MISSING/SIZE", f.rfilename); bad += 1; continue + if want: + h = hashlib.sha256() + with path.open("rb") as fh: + for chunk in iter(lambda: fh.read(64 << 20), b""): + h.update(chunk) + if h.hexdigest() != want: + print("SHA256", f.rfilename); bad += 1 +sys.exit(f"{bad} files failed verification" if bad else 0) +EOF + sha256sum "$HF/inference/model.py" + log "fetch verified" +} + +pack() { + python3 -I "$KIT/build_pack.py" --model-dir "$HF" --out "$W/pack" + log "pack built: $(python3 -c "import json;m=json.load(open('$W/pack/manifest.json'));print(m['records'],'records',m['sessions'],'sessions')")" +} + +sweep() { # one collect.py per GPU, sessions sharded across them; extra args go to collect.py + local name=$1; shift + local n i pids=() + n=$(nvidia-smi -L | wc -l) + log "sweep $name on $n GPU(s)" + for ((i = 0; i < n; i++)); do + CUDA_VISIBLE_DEVICES=$i python3 "$KIT/collect.py" --hf-path "$HF" --pack "$W/pack/v41-pack.jsonl" \ + --out "$OUT" --sweep "$name" --shard "$i/$n" --deadline-utc "$deadline" --chunk-tokens "$CHUNK_TOKENS" "$@" \ + > "$W/sweep-$name-$i.log" 2>&1 & + pids+=($!) + done + for i in "${!pids[@]}"; do + wait "${pids[$i]}" || { log "sweep $name shard $i failed: see $W/sweep-$name-$i.log"; return 1; } + done +} + +gpu_setup() { + [[ -f "$W/pip-freeze-gpu.txt" ]] && return + pip install -q -r "$HF/inference/requirements.txt" + pip freeze > "$W/pip-freeze-gpu.txt" + nvidia-smi --query-gpu=name,memory.total --format=csv | tee -a "$W/pod.log" +} + +smoke() { # the longest calibration record through all 40 layers, before any sweep + watchdog + gpu_setup + python3 - "$W/pack/v41-pack.jsonl" "$W/smoke/pack.jsonl" <<'PY' +import json, pathlib, sys +rows = [json.loads(line) for line in open(sys.argv[1]) if line.strip()] +row = max((r for r in rows if r["split"] == "calibration"), key=lambda r: len(r["token_ids"])) +row["session"] = "smoke" +pathlib.Path(sys.argv[2]).parent.mkdir(parents=True, exist_ok=True) +pathlib.Path(sys.argv[2]).write_text(json.dumps(row) + "\n") +print("smoke record", row["id"], len(row["token_ids"]), "tokens") +PY + rm -rf "$W/smoke/calibration" "$W/smoke"/progress-* + python3 "$KIT/collect.py" --hf-path "$HF" --pack "$W/smoke/pack.jsonl" --out "$W/smoke" --sweep calibration \ + 2>&1 | tee "$W/smoke.log" + python3 - "$W/smoke" <<'PY' +import glob, json, sys +import numpy as np +(entry,) = [json.loads(line) for f in glob.glob(sys.argv[1] + "/progress-*.jsonl") for line in open(f)] +with np.load(sys.argv[1] + "/calibration/sessions/smoke.npz") as z: + slots = z["reap__routed_slots"] + assert (slots == slots[0]).all() and slots[0] > 0, slots + assert np.isfinite(z["reap__gate_norm_sum"]).all() and np.isfinite(z["importance_sum__hidden"]).all() +print("smoke ok:", json.dumps(entry)) +sys.exit(0 if entry["max_gpu_gb"] < 76 else f"smoke: {entry['max_gpu_gb']} GB peak is too close to 80 GB") +PY + log "smoke passed: $(tail -1 "$W/smoke.log")" +} + +gpu() { + local deadline + deadline=$(date -u -d "+$((MAX_GPU_HOURS * 60 - 20)) minutes" +%FT%TZ) + watchdog + gpu_setup + + sweep calibration + + log "finalize-reap + candidates: $CANDIDATES" + python3 "$W/keep/benchmarks/produce_dsv4_teacher_cache.py" finalize-reap --out-dir "$OUT/calibration" + local survivors=() + for keep in $CANDIDATES; do + python3 "$W/keep/benchmarks/produce_dsv4_teacher_cache.py" reap-select \ + --stats "$OUT/calibration/reap-stats.npz" --keep "$keep" --output "$OUT/survivors-$keep.json" \ + > "$OUT/reap-select-$keep.json" + survivors+=(--survivors "$keep=$OUT/survivors-$keep.json") + done + + sweep eval "${survivors[@]}" + + local students=() + for keep in $CANDIDATES; do students+=(--student "$keep=$OUT/student-$keep"); done + python3 "$KIT/score.py" --teacher "$OUT" "${students[@]}" | tee "$OUT/candidate-scores.json" + log "gpu stage done; scores in $OUT/candidate-scores.json" +} + +prune() { + local keep=${1:?usage: pod.sh prune KEEP} + pip install -q numpy + python3 "$KIT/prune.py" --src "$HF" --survivors "$OUT/survivors-$keep.json" --out "$W/v41-reap$keep" + (cd "$W/v41-reap$keep" && sha256sum -c SHA256SUMS --quiet) && log "pruned checkpoint $W/v41-reap$keep verified" +} + +all() { + watchdog + fetch_small + log "fetch (weights) and pack in parallel" + fetch > "$W/fetch.log" 2>&1 & + local fetch_pid=$! + pack > "$W/pack.log" 2>&1 + wait "$fetch_pid" || { log "fetch failed: see $W/fetch.log"; return 1; } + smoke + gpu + for keep in $PRUNE; do prune "$keep"; done + log "all done: $(echo "$PRUNE" | sed "s|[0-9]*|$W/v41-reap&|g"). Pod left running for the Mac to pull; watchdog still armed." +} + +stage=${1:?usage: pod.sh fetch|pack|smoke|gpu|prune KEEP|all} +shift +"$stage" "$@" diff --git a/scripts/runpod/v41_reap/pool384_pod.sh b/scripts/runpod/v41_reap/pool384_pod.sh new file mode 100755 index 0000000..e0ee904 --- /dev/null +++ b/scripts/runpod/v41_reap/pool384_pod.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# Pool of 384 on one GPU pod: 2-bit E8P codes for the 192 experts per layer that REAP-50 removed, +# and every router row of the full checkpoint. +# +# The volume already holds what this needs, so the GPU starts work after one copy: +# /workspace/v41-hf the pinned full checkpoint +# /workspace/keep this repo (keep_quant_torch.py uploaded into the kit) +# /workspace/pool384/imp-pruned.npz importance of the removed experts (built on the Mac) +# Output, on the volume: /workspace/pool384/e8p-podimp-pruned/layer-XX.npz, router-384.npz, DONE. +set -euo pipefail +W=/workspace +P=$W/pool384 +KIT=$W/keep/scripts/runpod/v41_reap +LOCAL=${LOCAL:-/root/v41-experts} +export PYTHONPATH=$W/keep/src PIP_BREAK_SYSTEM_PACKAGES=1 +log() { printf '[%s] %s\n' "$(date -u +%FT%TZ)" "$*" | tee -a "$P/pod.log"; } +trap 'status=$?; [[ $status -ne 0 ]] && { log "FAILED (status $status)"; touch "$P/FAILED"; }' EXIT + +# Never read the checkpoint from the volume (one stream is 120 MB/s). Copy the shards that hold +# routed experts or router rows to the container's disk with 16 streams. +log "copy expert and router shards to $LOCAL" +mkdir -p "$LOCAL" +python3 - "$W/v41-hf" > /tmp/shards.txt <<'PY' +import json, sys +names = json.load(open(sys.argv[1] + "/model.safetensors.index.json"))["weight_map"] +print("\n".join(sorted({shard for name, shard in names.items() if ".ffn.experts." in name or ".ffn.gate." in name}))) +PY +(cd "$W/v41-hf" && xargs -P 16 -I{} cp {} "$LOCAL/" < /tmp/shards.txt && cp ./*.json "$LOCAL/") +want=$(cd "$W/v41-hf" && xargs stat -c %s < /tmp/shards.txt | awk '{s+=$1} END {print s}') +have=$(cd "$LOCAL" && xargs stat -c %s < /tmp/shards.txt | awk '{s+=$1} END {print s}') +[[ $want == "$have" ]] +log "copy done: $(wc -l < /tmp/shards.txt) shards, $((have >> 30)) GiB" + +pip install -q --no-index --find-links "$W/wheels" -r "$KIT/requirements-ref.txt" \ + || pip install -q -r "$KIT/requirements-ref.txt" +nvidia-smi --query-gpu=name,memory.total --format=csv | tee -a "$P/pod.log" + +log "router rows" +python3 "$KIT/keep_quant_torch.py" router --hf-path "$LOCAL" --out "$P/router-384.npz" | tee -a "$P/pod.log" + +log "smoke: two experts of layer 0" +rm -rf "$P/smoke" +python3 "$KIT/keep_quant_torch.py" fit --hf-path "$LOCAL" --importance "$P/imp-pruned.npz" --out "$P/smoke" \ + --layers 0 --experts 2 | tee -a "$P/pod.log" + +gpus=$(nvidia-smi -L | wc -l) +log "fit on $gpus GPU(s)" +for g in $(seq 0 $((gpus - 1))); do + CUDA_VISIBLE_DEVICES=$g python3 "$KIT/keep_quant_torch.py" fit --hf-path "$LOCAL" --importance "$P/imp-pruned.npz" \ + --out "$P/e8p-podimp-pruned" --layers "$(seq -s, "$g" "$gpus" 39)" > "$P/fit-$g.log" 2>&1 & +done +wait +[[ $(ls "$P/e8p-podimp-pruned"/layer-??.npz | wc -l) == 40 ]] +log "all done" +touch "$P/DONE" diff --git a/scripts/runpod/v41_reap/precision_alloc.py b/scripts/runpod/v41_reap/precision_alloc.py new file mode 100644 index 0000000..d1fea57 --- /dev/null +++ b/scripts/runpod/v41_reap/precision_alloc.py @@ -0,0 +1,101 @@ +"""Static per-expert precision allocation inside REAP-50 (step 2c): rungs 0 (dropped), 1 and 2 bits. + +Within each layer, under a budget of B units (1-bit expert = 1 unit, 2-bit = 2), greedily buy the +upgrade with the largest estimated error reduction per unit: + + cost(e, r) = m_e * eps_r m_e = expert e's share of the layer's routed output mass + (sum of g * ||f|| on the pruned model, ``--capture-imatrix``) + eps_r = relative output error at rung r (eps_0 = 1, dropped) + +eps_1 and eps_2 are the layer means of KEEP's weighted relative MSE (gate/up averaged, plus down) +from ``keep_quant.py``'s ``metrics.jsonl``. Every layer gets the same budget, because moving +budget between layers by mass share lost to uniform in step 1. + +Writes ``keep.json`` (experts at rung >= 1, original ids, for ``local_eval.py --variant``) and +``hot.json`` (rung 2, for ``--expert-override HOT@hot.json+REST``), plus ``alloc.json``. +""" + +from __future__ import annotations + +import argparse +import heapq +import json +from pathlib import Path + +import numpy as np + + +def layer_eps(metrics: Path) -> dict[int, float]: + rows = [json.loads(line) for line in metrics.read_text().splitlines() if line.strip()] + return {r["layer"]: (r["w1_wrel_mse"] + r["w3_wrel_mse"]) / 2 + r["w2_wrel_mse"] for r in rows} + + +def allocate(mass: np.ndarray, eps: tuple[float, float, float], budget: int) -> np.ndarray: + """Rung per expert (0/1/2) maximizing the estimated error reduction under ``budget`` units.""" + + rung = np.zeros(len(mass), dtype=int) + heap = [(-m * (eps[0] - eps[1]), e) for e, m in enumerate(mass)] + heapq.heapify(heap) + left = budget + while heap and left > 0: + _, e = heapq.heappop(heap) + rung[e] += 1 + left -= 1 + if rung[e] == 1: + heapq.heappush(heap, (-mass[e] * (eps[1] - eps[2]), e)) + return rung + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--checkpoint-survivors", required=True, type=Path) + p.add_argument("--imatrix", required=True, type=Path, help="pruned-model capture npz (reap__gate_norm_sum)") + p.add_argument("--metrics-1bit", required=True, type=Path) + p.add_argument("--metrics-2bit", required=True, type=Path) + p.add_argument("--budget", required=True, type=int, help="units per layer (1-bit expert = 1, 2-bit = 2)") + p.add_argument("--out", required=True, type=Path) + args = p.parse_args(argv) + + survivors = json.loads(args.checkpoint_survivors.read_text())["survivors"] + gate_mass = np.load(args.imatrix)["reap__gate_norm_sum"].astype(np.float64) # [L, E_local] + e1, e2 = layer_eps(args.metrics_1bit), layer_eps(args.metrics_2bit) + keep, hot, rows = {}, {}, [] + for layer in range(len(survivors)): + ids = sorted(survivors[str(layer)]) + mass = gate_mass[layer] / gate_mass[layer].sum() + eps = (1.0, e1[layer], e2[layer]) + rung = allocate(mass, eps, args.budget) + keep[str(layer)] = [ids[i] for i in np.flatnonzero(rung >= 1)] + hot[str(layer)] = [ids[i] for i in np.flatnonzero(rung == 2)] + rows.append({"layer": layer, "dropped": int((rung == 0).sum()), "rung1": int((rung == 1).sum()), + "rung2": int((rung == 2).sum()), "eps1": round(eps[1], 4), "eps2": round(eps[2], 4), + "est_cost": float((mass * np.asarray(eps)[rung]).sum())}) + args.out.mkdir(parents=True, exist_ok=True) + (args.out / "keep.json").write_text(json.dumps({"survivors": keep})) + (args.out / "hot.json").write_text(json.dumps({"survivors": hot})) + (args.out / "alloc.json").write_text(json.dumps({"budget": args.budget, "layers": rows}, indent=1)) + print(json.dumps({"budget": args.budget, "kept_mean": float(np.mean([len(v) for v in keep.values()])), + "hot_mean": float(np.mean([len(v) for v in hot.values()])), + "est_cost_mean": float(np.mean([r["est_cost"] for r in rows]))})) + return 0 + + +def _check() -> None: + eps = (1.0, 0.3, 0.1) + # flat mass: a 1-bit floor for everyone before any upgrade (0.175 per unit vs 0.05) + assert allocate(np.full(4, 0.25), eps, 4).tolist() == [1, 1, 1, 1] + m = np.array([0.5, 0.3, 0.15, 0.05]) + assert allocate(m, eps, 3).tolist() == [1, 1, 1, 0] + # upgrading the 0.5 expert (0.10) beats keeping the 0.05 one (0.035) + assert allocate(m, eps, 4).tolist() == [2, 1, 1, 0] + assert allocate(m, eps, 5).tolist() == [2, 2, 1, 0] + + +if __name__ == "__main__": + import sys + + if sys.argv[1:] == ["--check"]: + _check() + print("ok") + sys.exit(0) + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/prune.py b/scripts/runpod/v41_reap/prune.py new file mode 100644 index 0000000..ab55b14 --- /dev/null +++ b/scripts/runpod/v41_reap/prune.py @@ -0,0 +1,218 @@ +"""Write a REAP-pruned copy of a DeepSeek-V4.1 HF checkpoint, byte for byte. + +Streams each source shard and writes a shard of the same name holding: + +- every untouched tensor, copied unchanged; +- the kept routed experts, renumbered ``0..keep-1`` in ascending source order; +- the router ``weight`` / ``bias`` / ``bias_vl`` rows of the kept experts. + +Tensors are copied as raw safetensors bytes, so FP8 weights, ue8m0 scales and +packed FP4 experts are never decoded and nothing needs torch. The plan comes +from ``keep.quality.reap.prune_checkpoint_plan``; the drafter (``mtp.*``) is +left untouched unless ``--mtp-survivors`` is given. ``config.json`` and +``inference/config.json`` get the new ``n_routed_experts``; the index and a +sha256 manifest are rewritten. + +Usage:: + + python scripts/runpod/v41_reap/prune.py --src /workspace/v41-hf \\ + --survivors /workspace/out/survivors-192.json --out /workspace/v41-reap192 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import shutil +import struct +import sys +from pathlib import Path +from typing import Any, BinaryIO + +REPO = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO / "src")) + +from keep.quality.reap import prune_checkpoint_plan + +INDEX = "model.safetensors.index.json" +COPY_CHUNK = 64 << 20 + + +def read_header(handle: BinaryIO) -> tuple[dict[str, Any], int]: + (size,) = struct.unpack(" None: + src.seek(start) + while length: + chunk = src.read(min(COPY_CHUNK, length)) + if not chunk: + raise OSError("source shard ended early") + dst.write(chunk) + length -= len(chunk) + + +def prune_shard(src_path: Path, dst_path: Path, plan: dict[str, Any]) -> list[str]: + """Write one pruned shard; return the tensor names it holds.""" + + rename, slice_rows, drop = plan["rename"], plan["slice_rows"], set(plan["drop"]) + with src_path.open("rb") as src: + header, base = read_header(src) + metadata = header.pop("__metadata__", None) + jobs: list[tuple[str, dict[str, Any], list[tuple[int, int]]]] = [] + for name in sorted(header, key=lambda n: header[n]["data_offsets"][0]): + info = header[name] + start, end = info["data_offsets"] + if name in drop: + continue + if name in slice_rows: + rows = info["shape"][0] + row_bytes, rem = divmod(end - start, rows) + if rem: + raise ValueError(f"{name}: {end - start} bytes is not {rows} equal rows") + kept = slice_rows[name] + ranges = [(base + start + r * row_bytes, row_bytes) for r in kept] + shape = [len(kept), *info["shape"][1:]] + jobs.append((name, {"dtype": info["dtype"], "shape": shape}, ranges)) + else: + jobs.append((rename.get(name, name), {"dtype": info["dtype"], "shape": info["shape"]}, [(base + start, end - start)])) + out_header: dict[str, Any] = {} if metadata is None else {"__metadata__": metadata} + offset = 0 + for name, info, ranges in jobs: + length = sum(n for _, n in ranges) + out_header[name] = {**info, "data_offsets": [offset, offset + length]} + offset += length + blob = json.dumps(out_header, separators=(",", ":")).encode() + blob += b" " * (-len(blob) % 8) + dst_path.parent.mkdir(parents=True, exist_ok=True) + tmp = dst_path.with_name(dst_path.name + ".tmp") + with tmp.open("wb") as dst: + dst.write(struct.pack(" int: + """Set every ``n_routed_experts`` key (not the drafter's ``dspark_*``); return how many.""" + + hits = 0 + if isinstance(node, dict): + for key, value in node.items(): + if key == "n_routed_experts": + node[key] = keep + hits += 1 + else: + hits += set_expert_count(value, keep) + elif isinstance(node, list): + hits += sum(set_expert_count(item, keep) for item in node) + return hits + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(COPY_CHUNK), b""): + digest.update(chunk) + return digest.hexdigest() + + +def load_survivors(path: Path) -> dict[int, list[int]]: + payload = json.loads(path.read_text()) + if payload.get("record_type") != "keep_reap_survivors_v1": + raise ValueError(f"{path} is not a keep_reap_survivors_v1 table") + return {int(block): [int(e) for e in kept] for block, kept in payload["survivors"].items()} + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--src", required=True, type=Path, help="HF checkpoint directory") + parser.add_argument("--survivors", required=True, type=Path, help="reap-select output for the backbone layers") + parser.add_argument("--mtp-survivors", type=Path, help="optional survivor table for the mtp.* drafter") + parser.add_argument("--out", required=True, type=Path) + parser.add_argument("--skip-shard", action="append", default=[], + help="a shard the plan leaves untouched (the Engram tables): indexed but not written; use the source's file") + args = parser.parse_args(argv) + + index = json.loads((args.src / INDEX).read_text()) + weight_map: dict[str, str] = index["weight_map"] + survivors = load_survivors(args.survivors) + keep = len(next(iter(survivors.values()))) + plan = prune_checkpoint_plan(weight_map, survivors, prefix="layers") + if args.mtp_survivors: + mtp = prune_checkpoint_plan(plan["untouched"], load_survivors(args.mtp_survivors), prefix="mtp") + plan = { + "rename": {**plan["rename"], **mtp["rename"]}, + "slice_rows": {**plan["slice_rows"], **mtp["slice_rows"]}, + "drop": plan["drop"] + mtp["drop"], + } + # ponytail: the drafter's expert count lives under dspark_n_routed_experts; set it by hand if you prune it. + + args.out.mkdir(parents=True, exist_ok=True) + new_map: dict[str, str] = {} + total = 0 + for shard in sorted(set(weight_map.values())): + if shard in args.skip_shard: + with (args.src / shard).open("rb") as handle: + names = [n for n in read_header(handle)[0] if n != "__metadata__"] + touched = set(names) & (set(plan["rename"]) | set(plan["slice_rows"]) | set(plan["drop"])) + if touched: + raise ValueError(f"{shard} cannot be skipped: the plan changes {len(touched)} of its tensors") + total += (args.src / shard).stat().st_size + else: + names = prune_shard(args.src / shard, args.out / shard, plan) + total += (args.out / shard).stat().st_size + new_map.update({name: shard for name in names}) + print(f"{shard}: {len(names)} tensors", flush=True) + + for path in sorted(args.src.rglob("*")): + rel = path.relative_to(args.src) + if path.is_dir() or path.suffix == ".safetensors" or rel.name == INDEX or rel.parts[0].startswith("."): + continue + target = args.out / rel + target.parent.mkdir(parents=True, exist_ok=True) + if rel.name == "config.json": + config = json.loads(path.read_text()) + if set_expert_count(config, keep) == 0: + raise ValueError(f"{rel}: no n_routed_experts key to update") + target.write_text(json.dumps(config, indent=4) + "\n") + else: + shutil.copy2(path, target) + + expected = len(weight_map) - len(plan["drop"]) + if len(new_map) != expected: + raise AssertionError(f"wrote {len(new_map)} tensors, expected {expected}") + out_index = {"metadata": {**index.get("metadata", {}), "total_size": total}, "weight_map": dict(sorted(new_map.items()))} + (args.out / INDEX).write_text(json.dumps(out_index, indent=2) + "\n") + (args.out / "keep-prune.json").write_text( + json.dumps( + { + "record_type": "keep_reap_pruned_checkpoint_v1", + "source_index_sha256": sha256_file(args.src / INDEX), + "survivors_sha256": sha256_file(args.survivors), + "keep": keep, + "tensors_dropped": len(plan["drop"]), + "router_tensors_sliced": len(plan["slice_rows"]), + "experts_renamed": len(plan["rename"]), + }, + indent=2, + ) + + "\n" + ) + manifest = { + str(p.relative_to(args.out)): sha256_file(p) + for p in sorted(args.out.rglob("*")) + if p.is_file() and p.name != "SHA256SUMS" + } + (args.out / "SHA256SUMS").write_text("".join(f"{digest} {name}\n" for name, digest in manifest.items())) + print(json.dumps({"tensors": len(new_map), "bytes": total, "keep": keep})) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/requirements-ref.txt b/scripts/runpod/v41_reap/requirements-ref.txt new file mode 100644 index 0000000..0ac24b8 --- /dev/null +++ b/scripts/runpod/v41_reap/requirements-ref.txt @@ -0,0 +1,9 @@ +torch>=2.10.0 +transformers +tokenizers +safetensors>=0.7.0 +numpy +sympy +Pillow +tilelang==0.1.12 +tqdm diff --git a/scripts/runpod/v41_reap/rerank.py b/scripts/runpod/v41_reap/rerank.py new file mode 100644 index 0000000..71100c8 --- /dev/null +++ b/scripts/runpod/v41_reap/rerank.py @@ -0,0 +1,73 @@ +"""Re-rank experts by total routed mass (sum of g * ||f||), optionally inside an earlier survivors set. + +"Prune twice": measure routing on the model that is already pruned (``collect.py --sweep +calibration --mask TABLE``), then keep each layer's top K of that table by the new mass. On the +Mac this cut KL by 0.18 going from 192 to 96 experts; neither the mass criterion nor the +pruned-model routing did so alone (recovery log, step 2c). + + python rerank.py --stats OUT/calibration/sessions --within fm-288.json --keep 192 --out pt-192.json + +``--stats`` is a directory of ``collect.py`` session files or one npz holding +``reap__gate_norm_sum`` [layers, experts]. Without ``--within`` every expert is a candidate. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + + +def total_mass(stats: Path) -> np.ndarray: + paths = sorted(stats.glob("*.npz")) if stats.is_dir() else [stats] + assert paths, f"no npz under {stats}" + total = 0 + for path in paths: + with np.load(path) as z: + total = total + z["reap__gate_norm_sum"].astype(np.float64) + return total + + +def rerank(mass: np.ndarray, keep: int, within: dict[str, list[int]] | None) -> dict[str, list[int]]: + out = {} + for layer in range(mass.shape[0]): + ids = np.asarray(sorted(within[str(layer)]) if within else range(mass.shape[1])) + assert keep <= len(ids), f"layer {layer}: keep {keep} of {len(ids)}" + order = np.argsort(-mass[layer, ids], kind="stable")[:keep] + out[str(layer)] = sorted(int(e) for e in ids[order]) + return out + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--stats", required=True, type=Path) + p.add_argument("--within", type=Path, help="survivors JSON: candidates are this table's experts") + p.add_argument("--keep", required=True, type=int) + p.add_argument("--out", required=True, type=Path) + args = p.parse_args(argv) + within = json.loads(args.within.read_text())["survivors"] if args.within else None + mass = total_mass(args.stats) + table = rerank(mass, args.keep, within) + kept = np.mean([mass[int(l), ids].sum() / mass[int(l)].sum() for l, ids in table.items()]) + args.out.write_text(json.dumps({"record_type": "keep_reap_survivors_v1", "criterion": "total routed mass", + "stats": str(args.stats), "within": str(args.within) if args.within else None, + "survivors": table})) + print(json.dumps({"keep": args.keep, "layers": len(table), "mass_kept_mean": round(float(kept), 4)})) + return 0 + + +def _check() -> None: + mass = np.array([[5.0, 1.0, 4.0, 3.0], [0.0, 2.0, 9.0, 1.0]]) + assert rerank(mass, 2, None) == {"0": [0, 2], "1": [1, 2]} + assert rerank(mass, 1, {"0": [1, 3], "1": [0, 3]}) == {"0": [3], "1": [3]} # only inside the table + + +if __name__ == "__main__": + if sys.argv[1:] == ["--check"]: + _check() + print("ok") + sys.exit(0) + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/routed_report.py b/scripts/runpod/v41_reap/routed_report.py new file mode 100644 index 0000000..d2d04c6 --- /dev/null +++ b/scripts/runpod/v41_reap/routed_report.py @@ -0,0 +1,129 @@ +"""Costs and quality of routed-precision variants (``local_eval.py --dynamic``), per variant: + +- KL / top-1 / NLL change: domain means from ``summary.json``. +- resident expert memory: every base (tier 1) plus K residuals per layer when cached, otherwise + every residual too (nothing streamed), in units per layer (1 unit = one 1-bit expert) and GB. +- SSD traffic: residual admissions per token, in MB/token and GB/s at ``--tok-s``. +- unified-memory bytes read per token by routed experts (tier 1 = base, tier 2 = base + residual). + ``tok_s_bound`` divides ``--bandwidth-gbs`` by those bytes alone; routed experts were about 11% + of decode traffic on V4-Flash, so it is an upper bound on what expert tiers can buy, not a speed + estimate. + + python routed_report.py --eval ~/Models/v41-reap-run/evals/step4-routed [--eval ...] + python routed_report.py --compare EVAL_DIR:c96t2 EVAL_DIR:s96 +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + +DIM, INTER = 5120, 2304 +# 1-bit E8 code + fp16 group scales: gate/up group 512 (1.031 bpw), down group 256 (1.0625 bpw) +UNIT_BYTES = (2 * DIM * INTER * (1 + 16 / 512) + DIM * INTER * (1 + 16 / 256)) / 8 # ~4.6 MB +FP4_BYTES = 3 * DIM * INTER * (4 + 8 / 32) / 8 # an FP4 expert with one scale byte per 32 weights, 18.8 MB + + +def rows(eval_dir: Path, tok_s: float, bandwidth_gbs: float) -> list[dict]: + summary = json.loads((eval_dir / "summary.json").read_text()) + out = [] + for name, per in summary["variants"].items(): + tiers = per.get("_tiers") + row = {"variant": f"{eval_dir.name}:{name}", **{k: round(v, 4) for k, v in per["_domain_mean"].items()}} + if tiers is None: + out.append(row) + continue + c = np.load(eval_dir / name / "tier-counts.npy") # [L, E, 7] + n_layers = c.shape[0] + used = c[:, :, :2].sum() + c[:, :, 5].sum() # every routed slot + tokens = used / 6 / n_layers + k = tiers.get("cache_k") + n_bases, hi_units = tiers.get("n_bases", c.shape[1]), tiers.get("hi_units", 1.0) + units = n_bases + (k if k is not None else n_bases) * hi_units + hi_read = hi_units if tiers.get("copy_mode") else 1 + hi_units # a separate copy is read alone + read = (c[:, :, 0].sum() + c[:, :, 1].sum() * hi_read) * UNIT_BYTES / tokens + ssd = c[:, :, 4].sum() * hi_units * UNIT_BYTES / tokens if k is not None else 0.0 + if tiers.get("copy_mode") and k is not None and not tiers.get("miss"): + # A separate-copy cache can instead drop a hot expert's low copy (1 extra unit, not 2) and read + # it back on demotion: every promote/demote cycle then moves 3 units where a residual moves 1. + row |= {"units_per_layer_swap": round(n_bases + k, 1), + "ssd_gbs_swap": round(c[:, :, 4].sum() * 3 * UNIT_BYTES / tokens * tok_s / 1e9, 2)} + row |= { + "tau": tiers["tau"], "tau_drop": tiers.get("tau_drop", 0.0), "cache_k": k, "policy": tiers.get("policy"), + "miss": tiers.get("miss"), + "n_bases": n_bases, "hi_units": round(hi_units, 2), + "hi_slots": round(tiers["hi_slot_share"], 3), "drop_slots": round(tiers.get("drop_slot_share", 0.0), 3), + "units_per_layer": round(units, 1), "resident_gb": round(units * n_layers * UNIT_BYTES / 1e9, 1), + "ssd_mb_per_tok": round(ssd / 1e6, 1), "ssd_gbs": round(ssd * tok_s / 1e9, 2), + "read_gb_per_tok": round(read / 1e9, 2), "tok_s_bound": round(bandwidth_gbs * 1e9 / read, 1), + } + if tiers.get("hot"): # a cache whose hottest experts are FP4 copies, or carry one more bit as a residual + hot_bytes = FP4_BYTES if tiers.get("hot_tier", "fp4") == "fp4" else (hi_units + 1) * UNIT_BYTES + mixed = (k - tiers["hot"]) * hi_units * UNIT_BYTES + tiers["hot"] * hot_bytes + row |= {"hot": tiers["hot"], "hot_slots": round(tiers["hot_slot_share"], 3), + "resident_gb": round(mixed * n_layers / 1e9, 1)} + out.append(row) + return out + + +def session_kl(spec: str) -> dict[str, np.ndarray]: + """``EVAL_DIR:VARIANT`` -> {session: per-position KL}.""" + + d, _, name = spec.rpartition(":") + return {p.stem: np.load(p)["kl"].astype(np.float64) for p in sorted((Path(d).expanduser() / name).glob("*.npz")) + if p.stem != "admissions"} + + +def compare(a: str, b: str, reps: int = 2000, block: int = 64, seed: int = 0) -> dict: + """Paired block bootstrap of the domain-mean KL difference A - B. + + Positions are resampled in blocks of ``block`` within each session (the same blocks for A and + B), sessions are pooled by domain as in ``local_eval.py``, and domains are averaged.""" + + ka, kb = session_kl(a), session_kl(b) + da, _, na = a.rpartition(":") + doms = {s: r["domain"] for s, r in json.loads((Path(da).expanduser() / "summary.json").read_text())["variants"][na].items() + if not s.startswith("_")} + assert set(ka) == set(kb) == set(doms), "the two variants cover different sessions" + rng = np.random.default_rng(seed) + + def stat(pick) -> float: + num: dict[str, float] = {} + den: dict[str, int] = {} + for s in ka: + i = pick(len(ka[s])) + num[doms[s]] = num.get(doms[s], 0.0) + float((ka[s][i] - kb[s][i]).sum()) + den[doms[s]] = den.get(doms[s], 0) + len(i) + return float(np.mean([num[d] / den[d] for d in num])) + + def blocks(n: int) -> np.ndarray: + starts = rng.integers(0, max(n - block, 0) + 1, size=-(-n // block)) + return (starts[:, None] + np.arange(min(block, n))[None, :]).ravel()[:n] + + point = stat(lambda n: np.arange(n)) + draws = np.sort([stat(blocks) for _ in range(reps)]) + return {"a": a, "b": b, "kl_diff": round(point, 4), "ci95": [round(float(draws[int(0.025 * reps)]), 4), + round(float(draws[int(0.975 * reps)]), 4)]} + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--eval", action="append", default=[], type=Path) + p.add_argument("--compare", nargs=2, action="append", default=[], metavar=("A", "B"), + help="EVAL_DIR:VARIANT pairs: paired block-bootstrap 95%% interval of domain-mean KL(A) - KL(B)") + p.add_argument("--tok-s", type=float, default=30.0) + p.add_argument("--bandwidth-gbs", type=float, default=546.0) + args = p.parse_args(argv) + for d in args.eval: + for row in rows(d.expanduser(), args.tok_s, args.bandwidth_gbs): + print(json.dumps(row)) + for a, b in args.compare: + print(json.dumps(compare(a, b))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/runpod/v41_reap/score.py b/scripts/runpod/v41_reap/score.py new file mode 100644 index 0000000..7aab6fa --- /dev/null +++ b/scripts/runpod/v41_reap/score.py @@ -0,0 +1,51 @@ +"""Score masked-router student runs against the teacher, per domain. + + python scripts/runpod/v41_reap/score.py --teacher /workspace/out/teacher \\ + --student 256=/workspace/out/student-256 --student 192=/workspace/out/student-192 + +For each candidate: mean KL(teacher || student) over the teacher's top-K, the +NLL increase on the real next token, and top-1 agreement, per domain and +overall (position-weighted). Writes JSON to stdout. +""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path + +import numpy as np + + +def score(teacher: Path, student: Path) -> dict[str, dict[str, float]]: + sums: dict[str, np.ndarray] = defaultdict(lambda: np.zeros(4)) + for path in sorted((student / "eval" / "sessions").glob("*.npz")): + with np.load(path) as s, np.load(teacher / "eval" / "sessions" / path.name) as t: + if not np.array_equal(s["positions"], t["positions"]): + raise ValueError(f"{path.name}: positions differ from the teacher's") + teacher_nll = t["logsumexp"] - t["target_logit"] + row = np.array([s["kl_on_teacher_topk"].sum(), (s["student_nll"] - teacher_nll).sum(), s["top1_agree"].sum(), s["positions"].size]) + for domain in [*map(str, s["domains"]), "all"]: + sums[domain] += row + return { + domain: {"kl": kl / n, "nll_delta": dn / n, "top1_agree": agree / n, "positions": int(n)} + for domain, (kl, dn, agree, n) in sorted(sums.items()) + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--teacher", required=True, type=Path) + parser.add_argument("--student", action="append", required=True, help="KEEP=DIR, repeatable") + args = parser.parse_args(argv) + report = {} + for item in args.student: + keep, _, directory = item.partition("=") + report[keep] = score(args.teacher, Path(directory)) + print(json.dumps(report, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/runpod/v41_reap/sidecar_fit.py b/scripts/runpod/v41_reap/sidecar_fit.py new file mode 100644 index 0000000..2c8b615 --- /dev/null +++ b/scripts/runpod/v41_reap/sidecar_fit.py @@ -0,0 +1,332 @@ +"""Step 3: layer-local residual sidecars for a pruned, KEEP-quantized V4.1 student, in one pass. + +The student is a survivors set (base expert ids) whose experts run from KEEP codes as bf16 GEMMs. +The teacher is the same layer with every expert the checkpoint holds, at its own FP4 weights. In +one layer-streamed pass over the student's own hidden states, each layer's MoE computes both: + + r = y_teacher(x) - y_student(x) routed experts only; the shared expert cancels + +and fits a reduced-rank linear map x~ = [x, 1] -> r by ridge regression on the fit tokens: + + W = (X~'X~ + lam * mean(diag) * I)^-1 X~'R, then the best rank-k W V V' (V: top eigvecs of W'X~'X~W) + +The rank-k sidecar is added to the student's MoE output before the next layer runs, so every +later layer is fit on inputs that already carry the earlier sidecars (sequential, like GPTQ). +Held-out records score each layer: the share of teacher output energy the residual holds, +before and after the sidecar. + +On the pod the checkpoint is the full 384-expert source; on the Mac, REAP-50 stands in as a +proxy teacher for smoke tests. Student codes come from ``pack-codes`` (only the student's +experts, keyed by base id), so the pod never needs the full code directory: + + python sidecar_fit.py pack-codes --codes ~/Models/v41-reap-run/quant/e8p-podimp \\ + --checkpoint-survivors ~/Models/v41-reap-run/stats/survivors-192.json \\ + --student ~/Models/v41-reap-run/alloc/sel/prunedmass-96.json --out ~/Models/v41-reap-run/sidecar/codes + + python sidecar_fit.py fit --hf-path HF --pack PACK --records RECORDS.json \\ + --student SURVIVORS.json --codes CODES_DIR --rank 256 --out OUT + +``RECORDS.json`` is a list of ``session#index`` names; every ``--holdout-every``-th one is held out. +Writes ``OUT/layer-LL.npz`` (U [dim+1, k], Vt [k, dim], fp32; optional X~'X~ and X~'R with +``--save-stats``) and ``OUT/metrics.jsonl``. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import numpy as np + +KIT = Path(__file__).resolve().parent +sys.path.insert(0, str(KIT)) +sys.path.insert(0, str(KIT.parents[2] / "src")) + + +def local_ids(hf_path: Path, ckpt_survivors: Path | None, n_experts: int) -> dict[int, list[int]]: + """Per layer, the base expert id of each local expert slot in this checkpoint.""" + + if not (hf_path / "keep-prune.json").exists(): + return {layer: list(range(n_experts)) for layer in range(64)} + assert ckpt_survivors is not None, "a pruned checkpoint needs --checkpoint-survivors" + return {int(k): sorted(v) for k, v in json.loads(ckpt_survivors.read_text())["survivors"].items()} + + +def pack_codes(args: argparse.Namespace) -> int: + have = {int(k): sorted(v) for k, v in json.loads(args.checkpoint_survivors.read_text())["survivors"].items()} + want = {int(k): sorted(v) for k, v in json.loads(args.student.read_text())["survivors"].items()} + args.out.mkdir(parents=True, exist_ok=True) + for layer in sorted(want): + local = [have[layer].index(e) for e in want[layer]] + with np.load(args.codes / f"layer-{layer:02d}.npz") as z: + out = {"ids": np.asarray(want[layer], dtype=np.int32), + "code_bits": z["code_bits"] if "code_bits" in z.files else np.asarray(16)} + for proj in ("w1", "w2", "w3"): + out[f"{proj}_codes"] = z[f"{proj}_codes"][local] + out[f"{proj}_scales"] = z[f"{proj}_scales"][local] + np.savez(args.out / f"layer-{layer:02d}.npz", **out) + print(json.dumps({"layers": len(want), "experts": int(np.mean([len(v) for v in want.values()]))})) + return 0 + + +class LayerCodes: + """One ``pack-codes`` layer file, held as codes; ``[base id]`` decodes that expert to bf16. + + Decoding on use keeps only one expert's bf16 weights alive at a time (all 96 of a layer, twice + over with the prefetch, cost ~13 GB of host RAM and tripped the Mac's 80 GB guard).""" + + def __init__(self, path: Path, torch, grids): + self.torch = torch + with np.load(path) as z: + self.grid = grids[int(z["code_bits"])] + self.row = {e: i for i, e in enumerate(z["ids"].tolist())} + self.arrays = {k: z[k] for k in z.files if k.endswith(("_codes", "_scales"))} + + def __getitem__(self, e: int) -> dict[str, object]: + torch, i, out = self.torch, self.row[e], {} + with torch.device("cpu"): + for proj in ("w1", "w2", "w3"): + vec = self.grid[torch.as_tensor(self.arrays[f"{proj}_codes"][i].astype(np.int64))] + scales = torch.as_tensor(self.arrays[f"{proj}_scales"][i].astype(np.float32)) + vec = vec * scales.repeat_interleave(vec.shape[1] // scales.shape[-1], dim=-1)[..., None] + out[proj] = vec.flatten(-2).bfloat16() + return out + + +def reduced_rank(A, B, lam: float, rank: int): + """Ridge W = (A + reg)^-1 B, cut to the rank-k map that keeps the most fitted output energy.""" + + import torch + + n = A.shape[0] + reg = lam * A.diagonal()[:-1].mean() + W = torch.linalg.solve(A + reg * torch.eye(n, dtype=A.dtype, device=A.device), B) # [n, d] + M = W.T @ A @ W + _, V = torch.linalg.eigh((M + M.T) / 2) + V = V[:, -rank:].flip(-1) # largest first, so U[:, :j] @ Vt[:j] is the best rank-j map + return W @ V, V.T # U [n, k], Vt [k, d] + + +def fit(args: argparse.Namespace) -> int: + import torch + import torch.nn.functional as F + from transformers import AutoTokenizer + + import collect + import stream + from keep.vq.e8 import e8_1bit_grid, e8p_full_grid + + cuda = args.device == "cuda" + if not cuda: + import torch_kernel + + torch_kernel.exact_quant = True + ref = stream.import_reference(args.hf_path / "inference", kernel=torch_kernel) + else: + ref = stream.import_reference(args.hf_path / "inference") + torch.cuda.memory._set_allocator_settings("expandable_segments:True") + torch.set_default_device("cuda") + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(33377335) + model_args = ref.ModelArgs(**json.loads((args.hf_path / "inference" / "config.json").read_text())) + model_args.max_batch_size = 1 + dim, n_layers = int(model_args.dim), int(model_args.n_layers) + + sessions = collect.load_sessions(args.pack) + names = json.loads(args.records.read_text()) + records = [] + for name in names: + session, _, idx = name.partition("#") + r = sessions[session][int(idx or 0)] + records.append(dict(r, token_ids=r["token_ids"][: args.max_tokens])) + model_args.max_seq_len = max(len(r["token_ids"]) for r in records) + holdout = torch.cat([torch.full((len(r["token_ids"]),), i % args.holdout_every == args.holdout_every - 1) + for i, r in enumerate(records)]).to(args.device) + + ckpt = stream.HFCheckpoint(args.hf_path) + layout = ref.EngramLayout.from_args(model_args) + with torch.device(args.device): + hash_state = ref.NgramHashState(model_args, layout, AutoTokenizer.from_pretrained(args.hf_path)) + engram_rows = {l: stream.EngramRows.from_checkpoint(ckpt, l, ref.fp8_block_size) for l in model_args.engram_layer_ids} + streamer = stream.LayerStreamer(ref, model_args, ckpt.load, engram_rows, hash_state, args.device) + + ids = local_ids(args.hf_path, args.checkpoint_survivors, int(model_args.n_routed_experts)) + student = {int(k): v for k, v in json.loads(args.student.read_text())["survivors"].items()} + grids = {16: torch.as_tensor(e8p_full_grid().astype(np.float32), device="cpu"), + 8: torch.as_tensor(e8_1bit_grid().astype(np.float32), device="cpu")} + decoded: dict[int, dict] = {} + load_layer = streamer.load_layer + + def load_with_codes(prefix, skip=()): # prefetch thread: the next layer's student weights too + sd = load_layer(prefix, skip) + if prefix.startswith("layers."): + layer = int(prefix.split(".")[1]) + decoded[layer] = LayerCodes(args.codes / f"layer-{layer:02d}.npz", torch, grids) + return sd + + streamer.load_layer = load_with_codes + args.out.mkdir(parents=True, exist_ok=True) + # Resume: a layer with a metrics row keeps its saved sidecar; only the student is run through it. + done = {json.loads(line)["layer"] for line in (args.out / "metrics.jsonl").read_text().splitlines()} \ + if (args.out / "metrics.jsonl").exists() else set() + metrics = (args.out / "metrics.jsonl").open("a") + + def expert(x, w, w1, w2, w3, limit): + gate, up = F.linear(x, w1).float(), F.linear(x, w3).float() + if limit > 0: + up, gate = up.clamp(-limit, limit), gate.clamp(max=limit) + return F.linear((w * (F.silu(gate) * up)).to(x.dtype), w2) + + def moe_forward(self, x, image_mask=None): + t0 = time.time() + shape, layer = x.size(), self.layer_id + x = x.view(-1, dim) + local = ids[layer] + # This layer's Engram lookups are done. Unmapping the table drops the pages they touched + # (tens of GB at 150k tokens) from RSS; left mapped, layers 1 and 14 together passed 80 GB. + engram_rows.pop(layer, None) + n = x.size(0) + + def blocks(): # [x, 1] in row blocks, so no full fp32 copy of x is ever resident + for s in range(0, n, args.block_rows): + e = min(n, s + args.block_rows) + yield s, e, torch.cat([x[s:e].float(), torch.ones(e - s, 1, device=x.device)], dim=1) + + full_bias = self.gate.bias.data.clone() + if layer not in done: + weights, indices = self.gate(x, None) # teacher: every expert in the checkpoint + y_t = torch.zeros_like(x, dtype=torch.float32) + for i in torch.unique(indices).tolist(): + idx, top = torch.where(indices == i) + y_t[idx] += self.experts[i](x[idx], weights[idx, top, None]).float() + keep = set(student[layer]) + self.gate.bias.data[torch.as_tensor([e not in keep for e in local], device=x.device)] = float("-inf") + weights, indices = self.gate(x, None) # student: masked router, KEEP-decoded experts + self.gate.bias.data.copy_(full_bias) + weights_of = decoded.pop(layer) + y_s = torch.zeros_like(x, dtype=torch.float32) + limit = self.experts[0].swiglu_limit if self.experts[0] is not None else 0.0 + for i in torch.unique(indices).tolist(): + idx, top = torch.where(indices == i) + w = {p: t.to(x.device) for p, t in weights_of[local[i]].items()} + y_s[idx] += expert(x[idx], weights[idx, top, None], w["w1"], w["w2"], w["w3"], limit).float() + + if layer in done: + if args.apply: + with np.load(args.out / f"layer-{layer:02d}.npz") as z: + U, Vt = torch.as_tensor(z["U"], device=x.device), torch.as_tensor(z["Vt"], device=x.device) + for s, e, xt in blocks(): + y_s[s:e] += (xt @ U) @ Vt + print(json.dumps({"layer": layer, "resumed": True, "seconds": round(time.time() - t0, 1)}), flush=True) + y_s += self.shared_experts(x).float() + return y_s.type_as(x).view(shape) + + energy = {s: float(y_t[m].pow(2).sum()) for s, m in (("fit", ~holdout), ("holdout", holdout))} + r = y_t.sub_(y_s) # in place: y_t is not needed past its energy + del y_t + A = torch.zeros(dim + 1, dim + 1, dtype=torch.float64, device=x.device) + B = torch.zeros(dim + 1, dim, dtype=torch.float64, device=x.device) + + for s, e, xt in blocks(): + m = ~holdout[s:e] + A += (xt[m].T @ xt[m]).double() + B += (xt[m].T @ r[s:e][m]).double() + U, Vt = reduced_rank(A, B, args.ridge, args.rank) + U, Vt = U.float(), Vt.float() + before = {"fit": 0.0, "holdout": 0.0} + after = {"fit": 0.0, "holdout": 0.0} + for s, e, xt in blocks(): + corr = (xt @ U) @ Vt + for split, m in (("fit", ~holdout[s:e]), ("holdout", holdout[s:e])): + before[split] += float(r[s:e][m].pow(2).sum()) + after[split] += float((r[s:e][m] - corr[m]).pow(2).sum()) + if args.apply: + y_s[s:e] += corr + del r + row = {"layer": layer, "tokens_fit": int((~holdout).sum()), "tokens_holdout": int(holdout.sum()), + "seconds": round(time.time() - t0, 1)} + for split in ("fit", "holdout"): + row[f"{split}_res_before"] = before[split] / energy[split] + row[f"{split}_res_after"] = after[split] / energy[split] + out = {"U": U.cpu().numpy(), "Vt": Vt.cpu().numpy(), "rank": np.asarray(args.rank), "ridge": np.asarray(args.ridge)} + if args.save_stats: + out |= {"XtX": A.float().cpu().numpy(), "XtR": B.float().cpu().numpy()} + np.savez(args.out / f"layer-{layer:02d}.npz", **out) + metrics.write(json.dumps(row) + "\n") + metrics.flush() + print(json.dumps(row), flush=True) + y_s += self.shared_experts(x).float() + return y_s.type_as(x).view(shape) + + ref.MoE.forward = moe_forward + items = [stream.Item(tokens=torch.as_tensor([r["token_ids"]], dtype=torch.long, device=args.device), group="student") + for r in records] + t0 = time.time() + with torch.inference_mode(): + streamer.forward(items) + summary = {"seconds": round(time.time() - t0, 1), "tokens": int(holdout.numel()), "layers": n_layers} + if cuda: + summary["max_gpu_gb"] = round(torch.cuda.max_memory_allocated() / 2**30, 1) + print(json.dumps(summary), flush=True) + return 0 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="cmd", required=True) + pc = sub.add_parser("pack-codes") + pc.add_argument("--codes", required=True, type=Path) + pc.add_argument("--checkpoint-survivors", required=True, type=Path) + pc.add_argument("--student", required=True, type=Path) + pc.add_argument("--out", required=True, type=Path) + f = sub.add_parser("fit") + f.add_argument("--hf-path", required=True, type=Path) + f.add_argument("--checkpoint-survivors", type=Path, help="only for a pruned checkpoint (the Mac proxy)") + f.add_argument("--pack", required=True, type=Path) + f.add_argument("--records", required=True, type=Path) + f.add_argument("--max-tokens", type=int, default=1024) + f.add_argument("--holdout-every", type=int, default=10) + f.add_argument("--student", required=True, type=Path) + f.add_argument("--codes", required=True, type=Path) + f.add_argument("--rank", type=int, default=256) + f.add_argument("--ridge", type=float, default=1e-3) + f.add_argument("--no-apply", dest="apply", action="store_false", help="fit only; do not feed sidecars forward") + f.add_argument("--save-stats", action="store_true") + f.add_argument("--block-rows", type=int, default=32768, help="rows per [x, 1] block in the fit") + f.add_argument("--device", default="cuda") + f.add_argument("--out", required=True, type=Path) + f.add_argument("--max-rss-gb", type=float, default=80.0) + args = p.parse_args(argv) + if args.cmd == "pack-codes": + return pack_codes(args) + from local_parity import rss_guard + + rss_guard(args.max_rss_gb) + return fit(args) + + +def _check() -> None: + import torch + + torch.manual_seed(0) + x = torch.randn(4000, 16, dtype=torch.float64) + xt = torch.cat([x, torch.ones(4000, 1, dtype=torch.float64)], 1) + true = torch.randn(17, 3, dtype=torch.float64) @ torch.randn(3, 16, dtype=torch.float64) # rank 3 + r = xt @ true + U, Vt = reduced_rank(xt.T @ xt, xt.T @ r, 1e-9, 3) + assert torch.allclose(U @ Vt, true, atol=1e-6), (U @ Vt - true).abs().max() + U3, Vt3 = U, Vt + U, Vt = reduced_rank(xt.T @ xt, xt.T @ r, 1e-9, 1) + assert ((xt @ U @ Vt - r).pow(2).sum() / r.pow(2).sum()) < 1 # rank 1 keeps part of it + assert torch.allclose(U3[:, :1] @ Vt3[:1], U @ Vt, atol=1e-6) # stored largest first: truncation works + + +if __name__ == "__main__": + if sys.argv[1:] == ["--check"]: + _check() + print("ok") + sys.exit(0) + sys.exit(main()) diff --git a/scripts/runpod/v41_reap/sidecar_pod.sh b/scripts/runpod/v41_reap/sidecar_pod.sh new file mode 100755 index 0000000..4a33e0d --- /dev/null +++ b/scripts/runpod/v41_reap/sidecar_pod.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Step 3 on one H100 pod: fetch the source checkpoint, then fit the residual sidecars. +# The Mac uploads /workspace/in/ (pack, records, student survivors, student codes) meanwhile +# and touches /workspace/in/READY when it is complete. Writes /workspace/sidecar/ and DONE. +set -euo pipefail +W=/workspace +KIT=$W/keep/scripts/runpod/v41_reap +log() { printf '[%s] %s\n' "$(date -u +%FT%TZ)" "$*" | tee -a "$W/pod.log"; } + +log "fetch + pip in parallel" +bash "$KIT/pod.sh" fetch > "$W/fetch.log" 2>&1 & +fetch_pid=$! +pip install -q -r "$KIT/requirements-ref.txt" > "$W/pip.log" 2>&1 # the pinned inference/requirements.txt +wait "$fetch_pid" || { log "fetch failed: see $W/fetch.log"; exit 1; } +log "fetch verified; waiting for the Mac upload" +while [[ ! -f "$W/in/READY" ]]; do sleep 30; done +log "fit start" +python3 "$KIT/sidecar_fit.py" fit --hf-path "$W/v41-hf" --pack "$W/in/v41-pack.jsonl" \ + --records "$W/in/records-pod.json" --student "$W/in/prunedmass-96.json" --codes "$W/in/codes" \ + --rank 256 --save-stats --device cuda --max-rss-gb 400 --out "$W/sidecar" > "$W/fit.log" 2>&1 \ + || { log "fit failed: see $W/fit.log"; exit 1; } +log "fit done: $(tail -1 "$W/fit.log")" +touch "$W/DONE" diff --git a/scripts/runpod/v41_reap/stream.py b/scripts/runpod/v41_reap/stream.py new file mode 100644 index 0000000..0011efe --- /dev/null +++ b/scripts/runpod/v41_reap/stream.py @@ -0,0 +1,333 @@ +"""DeepSeek-V4.1-Flash forward on one GPU, one layer at a time. + +Design: ``docs/deepseek-v4-flash/2026-11-06-v41-reap-single-gpu-design.md``. + +The reference ``Block`` modules run unchanged. Only the order of work changes: +for a chunk of records, each layer's weights are loaded once (straight from the +HF shards, no ``convert.py``), every record runs through that layer's attention, +then the MoE runs over whole groups of records at once. Between layers each +record keeps three things: + +- its hyper-connection residual; +- its ``pre_mix``; +- its slice of ``shared_attn``, the module-level slot the reference uses to hand + compressed KV, index keys, top-k indices and candidate blocks down the stack. + +Engram rows are gathered from the shard files on the host and never put on the +GPU. ``tests/test_runpod_v41_stream.py`` checks the whole thing against +``Transformer.forward`` on a toy model. +""" + +from __future__ import annotations + +import importlib +import json +import struct +import sys +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +INDEX = "model.safetensors.index.json" +SHARED_FIELDS = ("compress_kv", "index_k", "topk_idxs", "candidates") +ST_DTYPES = { + "F8_E4M3": torch.float8_e4m3fn, + "F8_E8M0": torch.float8_e8m0fnu, + "F32": torch.float32, + "BF16": torch.bfloat16, +} + + +def import_reference(ref_dir: Path, kernel: Any | None = None) -> Any: + """Import the reference ``model`` module, optionally with a stand-in ``kernel``.""" + + if kernel is not None: + sys.modules["kernel"] = kernel + path = str(ref_dir) + if path not in sys.path: + sys.path.insert(0, path) + return importlib.import_module("model") + + +# --------------------------------------------------------------------------- +# Weights +# --------------------------------------------------------------------------- + + +def reference_name(name: str) -> str: + """HF tensor name -> reference parameter name (``convert.py`` at --model-parallel 1).""" + + name = name.removeprefix("model.") + name = name.replace("self_attn", "attn") + if not name.startswith("vision."): + name = name.replace("mlp", "ffn") + return name.replace("weight_scale_inv", "scale").replace("e_score_correction_bias", "bias") + + +def read_header(path: Path) -> tuple[dict[str, Any], int]: + with path.open("rb") as handle: + (size,) = struct.unpack(" dict[str, torch.Tensor]: + """Every tensor under ``prefix``, keyed relative to it, with convert.py's fixups.""" + + from safetensors import safe_open + + wanted: dict[str, list[tuple[str, str]]] = {} + for ref_name, (hf_name, shard) in self.shard_of.items(): + if ref_name.startswith(prefix) and not any(s in ref_name for s in skip): + wanted.setdefault(shard, []).append((ref_name, hf_name)) + out: dict[str, torch.Tensor] = {} + for shard, names in wanted.items(): + with safe_open(str(self.root / shard), framework="pt", device="cpu") as handle: + for ref_name, hf_name in names: + out[ref_name[len(prefix):]] = handle.get_tensor(hf_name) + return fix_tensors(out) + + +def fix_tensors(sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """``convert.py``'s per-tensor changes: wo_a dequantized to bf16, FP4 experts viewed as fp4.""" + + for name in [n for n in sd if n.endswith("wo_a.weight")]: + scale_name = name.replace("weight", "scale") + if scale_name in sd: + weight, scale = sd[name], sd.pop(scale_name) + ob, ib = weight.size(0) // scale.size(0), weight.size(1) // scale.size(1) + assert (ob, ib) in ((32, 32), (128, 128)), (name, weight.shape, scale.shape) + full = weight.unflatten(0, (-1, ob)).unflatten(-1, (-1, ib)).float() * scale[:, None, :, None].float() + sd[name] = full.flatten(2, 3).flatten(0, 1).bfloat16() + for name, tensor in sd.items(): + if "experts" in name and tensor.dtype == torch.int8: + sd[name] = tensor.view(torch.float4_e2m1fn_x2) + return sd + + +# --------------------------------------------------------------------------- +# Engram rows on the host +# --------------------------------------------------------------------------- + + +class EngramRows: + """One layer's n-gram table as host byte arrays (np.memmap of the shard, or plain arrays).""" + + def __init__(self, weight: np.ndarray, scale: np.ndarray, scale_dtype: torch.dtype, block_size: int): + self.weight, self.scale = weight, scale # uint8 [rows, dim], uint8 [rows, scale bytes] + self.scale_dtype, self.block_size = scale_dtype, block_size + + @classmethod + def from_checkpoint(cls, ckpt: HFCheckpoint, layer: int, block_size: int) -> EngramRows: + arrays = {} + for leaf in ("weight", "scale"): + hf_name, shard = ckpt.shard_of[f"layers.{layer}.engram.embed.{leaf}"] + header, base = read_header(ckpt.root / shard) + info = header[hf_name] + start, end = info["data_offsets"] + rows = info["shape"][0] + arrays[leaf] = (np.memmap(ckpt.root / shard, dtype=np.uint8, mode="r", offset=base + start, + shape=(rows, (end - start) // rows)), info["dtype"]) + assert arrays["weight"][1] == "F8_E4M3", arrays["weight"][1] + return cls(arrays["weight"][0], arrays["scale"][0], ST_DTYPES[arrays["scale"][1]], block_size) + + def lookup(self, indices: torch.Tensor) -> torch.Tensor: + """``ParallelEngramEmbedding.forward`` at world_size 1, reading rows from the host.""" + + flat = indices.reshape(-1).cpu().numpy() + unique, inverse = np.unique(flat, return_inverse=True) # sorted: near-sequential reads + w = torch.from_numpy(np.ascontiguousarray(self.weight[unique]))[torch.from_numpy(inverse)] + s = torch.from_numpy(np.ascontiguousarray(self.scale[unique]))[torch.from_numpy(inverse)] + w = w.to(indices.device).view(torch.float8_e4m3fn).view(*indices.shape, -1) + s = s.to(indices.device).view(self.scale_dtype).view(*indices.shape, -1) + values = w.float().unflatten(-1, (-1, self.block_size)) * s.float().unsqueeze(-1) + return values.flatten(-2).to(torch.bfloat16) + + +class HostTableEmbedding(torch.nn.Module): + """Stand-in for ``ParallelEngramEmbedding``: same forward, no table on the device.""" + + def __init__(self, num_embeddings: int, dim: int): + super().__init__() + self.num_embeddings, self.dim = num_embeddings, dim + self.rows: EngramRows | None = None + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + assert self.rows is not None, "engram rows not attached" + return self.rows.lookup(indices) + + +# --------------------------------------------------------------------------- +# shared_attn per record +# --------------------------------------------------------------------------- + + +def pack_bits(mask: torch.Tensor) -> tuple[torch.Tensor, int]: + n = mask.shape[-1] + padded = torch.nn.functional.pad(mask.to(torch.uint8), (0, -n % 8)).unflatten(-1, (-1, 8)) + shifts = torch.arange(8, dtype=torch.uint8, device=mask.device) + return (padded << shifts).sum(dim=-1, dtype=torch.uint8), n + + +def unpack_bits(packed: tuple[torch.Tensor, int]) -> torch.Tensor: + bits, n = packed + shifts = torch.arange(8, dtype=torch.uint8, device=bits.device) + return ((bits.unsqueeze(-1) >> shifts) & 1).bool().flatten(-2)[..., :n] + + +# --------------------------------------------------------------------------- +# The streamer +# --------------------------------------------------------------------------- + + +@dataclass +class Item: + """One record in one stream (the teacher, or a masked-router student).""" + + tokens: torch.Tensor # [1, s] long + group: Any # items sharing a group go through the MoE together + h: torch.Tensor | None = None + pre_mix: torch.Tensor | None = None + hashes: torch.Tensor | None = None + shared: dict[str, Any] = field(default_factory=dict) + restored: dict[str, Any] = field(default_factory=dict) + scratch: dict[str, Any] = field(default_factory=dict) + + +class LayerStreamer: + def __init__( + self, + ref: Any, + args: Any, + load_layer: Callable[[str], dict[str, torch.Tensor]], + engram_rows: dict[int, EngramRows], + hash_state: Any, + device: str, + ): + self.ref, self.args, self.device = ref, args, device + self.load_layer, self.engram_rows, self.hash_state = load_layer, engram_rows, hash_state + ref.world_size, ref.rank = 1, 0 + ref.default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16 + self.layout = ref.EngramLayout.from_args(args) + self.hc = args.hc_mult + with torch.device(device): + self.embed = ref.ParallelEmbedding(args.vocab_size, args.dim) + self.norm = ref.RMSNorm(args.dim, args.norm_eps) + self.head = ref.ParallelHead(args.vocab_size, args.dim, args.norm_eps, args.hc_eps) + for prefix, module in (("embed.", self.embed), ("norm.", self.norm), ("head.", self.head)): + module.load_state_dict(load_layer(prefix), strict=True) + self._pool = ThreadPoolExecutor(max_workers=1) + + def build_block(self, layer: int, state: dict[str, torch.Tensor]) -> Any: + ref = self.ref + real = ref.ParallelEngramEmbedding + ref.ParallelEngramEmbedding = HostTableEmbedding + try: + with torch.device(self.device): + block = ref.Block(layer, self.args, self.layout) + finally: + ref.ParallelEngramEmbedding = real + block.load_state_dict(state, strict=True) + if block.engram is not None: + block.engram.embed.rows = self.engram_rows[layer] + return block + + # -- per-record shared_attn ------------------------------------------------ + + def restore(self, item: Item) -> None: + shared = self.ref.shared_attn + item.restored = {} + for name in SHARED_FIELDS: + value = item.shared.get(name) + if name == "candidates" and value is not None: + value = unpack_bits(value) + item.restored[name] = value + setattr(shared, name, value) + + def snapshot(self, item: Item, block: Any) -> None: + shared = self.ref.shared_attn + own = {id(b) for b in block.buffers()} + for name in SHARED_FIELDS: + value = getattr(shared, name) + if value is None or value is item.restored.get(name): + continue # unchanged by this layer: keep the stored copy + if id(value) in own: # a view of this block's cache: the next record overwrites it + value = value.clone() + item.shared[name] = pack_bits(value) if name == "candidates" else value + for name in SHARED_FIELDS: + setattr(shared, name, None) + + # -- the forward ------------------------------------------------------------- + + def start(self, items: list[Item]) -> None: + for item in items: + h = self.embed(item.tokens) + item.h = h.unsqueeze(2).repeat(1, 1, self.hc, 1) + item.pre_mix = self.ref.make_identity_pre_mix(item.h, self.hc) + if self.hash_state is not None: + item.hashes = self.hash_state(item.tokens, 0) + + def run_layer(self, block: Any, items: list[Item], before_moe: Callable[[int, Any, Any], None] | None) -> None: + for item in items: # attention: one record at a time, whole sequence + if block.engram is not None: + item.h = block.engram(item.h, item.hashes[:, :, block.engram.layer_hash_index, :], None) + self.restore(item) + residual = item.h + attn_pre, attn_post, attn_comb = block.hc_mixes(item.h, block.hc_attn_fn, block.hc_attn_scale, block.hc_attn_base) + x = block.attn_norm(block.hc_pre(item.h, item.pre_mix)) + x = block.attn(x, 0) + item.h = block.hc_post(x, residual, attn_post, attn_comb) + item.scratch["attn_pre"] = attn_pre + self.snapshot(item, block) + + groups: dict[Any, list[Item]] = {} + for item in items: + groups.setdefault(item.group, []).append(item) + for key, members in groups.items(): # FFN: the whole group through the MoE at once + xs = [] + for item in members: + ffn_pre, ffn_post, ffn_comb = block.hc_mixes(item.h, block.hc_ffn_fn, block.hc_ffn_scale, block.hc_ffn_base) + xs.append(block.ffn_norm(block.hc_pre(item.h, item.scratch.pop("attn_pre")))) + item.scratch["ffn"] = (ffn_pre, ffn_post, ffn_comb) + if before_moe is not None: + before_moe(block.layer_id, block, key) + y = block.ffn(torch.cat(xs, dim=1) if len(xs) > 1 else xs[0], None) + for item, part in zip(members, y.split([x.size(1) for x in xs], dim=1)): + ffn_pre, ffn_post, ffn_comb = item.scratch.pop("ffn") + item.h = block.hc_post(part, item.h, ffn_post, ffn_comb) + item.pre_mix = ffn_pre + + def forward(self, items: list[Item], before_moe: Callable[[int, Any, Any], None] | None = None) -> None: + """Run ``items`` through every backbone layer; leaves each ``item.h`` / ``item.pre_mix``.""" + + n_layers = int(self.args.n_layers) + skip = (".engram.embed.",) + pending = self._pool.submit(self.load_layer, "layers.0.", skip) + self.start(items) + for layer in range(n_layers): + state = pending.result() + if layer + 1 < n_layers: + pending = self._pool.submit(self.load_layer, f"layers.{layer + 1}.", skip) + block = self.build_block(layer, state) + del state + self.run_layer(block, items, before_moe) + del block + for item in items: + item.shared.clear() + + def final_hidden(self, item: Item) -> torch.Tensor: + """``self.norm(layer.hc_pre(h, pre_mix))``: the head's input, [1, s, dim].""" + + return self.norm(self.ref.Block.hc_pre(None, item.h, item.pre_mix)) diff --git a/scripts/runpod/v41_reap/tau384_pod.sh b/scripts/runpod/v41_reap/tau384_pod.sh new file mode 100644 index 0000000..d8f1aab --- /dev/null +++ b/scripts/runpod/v41_reap/tau384_pod.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# Tau sweep over the pool of 384 on one GPU pod. +# +# 1. Fit 2-bit E8P codes for all 384 experts of every layer (keep_quant_torch.py, the fit the Mac +# serves; uploading the Mac's 130 GB of codes would take 12 h at its 3 MB/s uplink). +# 2. Simulate the expert cache that omlx_lattice.py serves (local_eval.py: 144 residents, re-route +# on a miss) at several admission thresholds, scored against the teacher's logits on the Mac's +# 22 eval sessions (first 1,024 tokens each). +# +# On the volume before the GPU starts (a CPU pod put it there): +# /workspace/v41-hf, /workspace/pack, /workspace/wheels, /workspace/keep +# /workspace/pool384/{imp-all.npz,reap-stats.npz,sessions.json} +# /workspace/pool384/stats/eval/sessions/NAME.npz teacher logits +# Output on the volume: pool384/codes/e8p-384 (the codes), pool384/tau-sweep/summary.json, DONE. +set -euo pipefail +W=/workspace +P=$W/pool384 +KIT=$W/keep/scripts/runpod/v41_reap +LOCAL=${LOCAL:-/root/v41} +CODES=$LOCAL/codes/e8p-384 +export PYTHONPATH=$W/keep/src PIP_BREAK_SYSTEM_PACKAGES=1 +log() { printf '[%s] %s\n' "$(date -u +%FT%TZ)" "$*" | tee -a "$P/pod.log"; } +trap 'status=$?; [[ $status -ne 0 ]] && { log "FAILED (status $status)"; touch "$P/FAILED"; }' EXIT + +# Never run from the volume (one stream is 120 MB/s): the checkpoint goes to the container's +# disk with 16 streams (about 480 GB). +log "copy checkpoint to $LOCAL/hf" +mkdir -p "$LOCAL/hf" +(cd "$W/v41-hf" && find . -type f -not -path './.cache/*' -print0 | xargs -0 -P 16 -I{} cp --parents {} "$LOCAL/hf/") +[[ $(cd "$W/v41-hf" && find . -type f -not -path './.cache/*' -printf '%s\n' | awk '{s+=$1} END {print s}') == $(cd "$LOCAL/hf" && find . -type f -printf '%s\n' | awk '{s+=$1} END {print s}') ]] +log "copy done: $(du -sh "$LOCAL/hf" | cut -f1)" + +pip install -q --no-index --find-links "$W/wheels" -r "$KIT/requirements-ref.txt" \ + || pip install -q -r "$KIT/requirements-ref.txt" +nvidia-smi --query-gpu=name,memory.total --format=csv | tee -a "$P/pod.log" + +fit() { python3 "$KIT/keep_quant_torch.py" fit --hf-path "$LOCAL/hf" --importance "$P/imp-all.npz" "$@"; } +log "fit smoke: four experts of layer 0" +rm -rf "$LOCAL/fit-smoke" +fit --out "$LOCAL/fit-smoke" --layers 0 --experts 4 | tee -a "$P/pod.log" +hours=$(python3 -c "import json,sys; print(round(json.loads(open(sys.argv[1]).readline())['seconds'] / 4 * 384 * 40 / 3600, 2))" "$LOCAL/fit-smoke/metrics.jsonl") +log "fit of all 15,360 experts projected at $hours h" +python3 -c "import sys; sys.exit(float(sys.argv[1]) > float(sys.argv[2]))" "$hours" "${MAX_FIT_H:-4}" || { log "over MAX_FIT_H, stopping"; exit 1; } +fit --out "$CODES" > "$P/fit.log" 2>&1 +[[ $(find "$CODES" -name 'layer-??.npz' | wc -l) == 40 ]] +log "fit done: $(tail -1 "$P/fit.log")" +mkdir -p "$P/codes/e8p-384" # keep the codes for later pod runs; the copy runs beside the sweep +(cd "$CODES" && find . -type f -print0 | xargs -0 -P 8 -I{} cp {} "$P/codes/e8p-384/") & +keep=$! + +run() { + python3 "$KIT/local_eval.py" --device cuda --hf-path "$LOCAL/hf" --pack "$W/pack/v41-pack.jsonl" --stats "$P/stats" \ + --max-rss-gb 150 --cache-warm "$P/reap-stats.npz" --dynamic "|$CODES" "$@" +} +v() { echo --variant "$1=all" --tau "$1=$3" --cache-k "$1=$2" --miss "$1=reroute" ${4:+--cache-policy "$1=$4"}; } + +log "sweep smoke: one session, 128 tokens" +echo '["eval-00000"]' > "$P/smoke-sessions.json" +# shellcheck disable=SC2046 +run --sessions-file "$P/smoke-sessions.json" --max-tokens 128 $(v q384 384 0) $(v t20 144 0.2) --out "$P/smoke-out" 2>&1 | tail -5 | tee -a "$P/pod.log" + +# q384: every expert resident (2-bit coding alone). s144: the 144 hottest, never changed. +# tNN: the cache of 144 at tau 0.NN. o144: each record's own best 144 (the policy's ceiling). +log "sweep: 22 sessions, 1,024 tokens, 7 variants" +# shellcheck disable=SC2046 +run --sessions-file "$P/sessions.json" --max-tokens 1024 \ + $(v q384 384 0) $(v s144 144 1e9) $(v t10 144 0.1) $(v t20 144 0.2) $(v t30 144 0.3) $(v t45 144 0.45) $(v o144 144 0.2 oracle) \ + --out "$P/tau-sweep" > "$P/sweep.log" 2>&1 +python3 "$KIT/routed_report.py" --eval "$P/tau-sweep" | tee -a "$P/pod.log" +wait $keep +[[ $(find "$P/codes/e8p-384" -name 'layer-??.npz' | wc -l) == 40 ]] +log "all done" +touch "$P/DONE" diff --git a/scripts/runpod/v41_reap/torch_kernel.py b/scripts/runpod/v41_reap/torch_kernel.py new file mode 100644 index 0000000..47f7610 --- /dev/null +++ b/scripts/runpod/v41_reap/torch_kernel.py @@ -0,0 +1,140 @@ +"""Plain-torch stand-ins for the reference ``inference/kernel.py`` (tilelang, GPU only). + +Install as ``sys.modules["kernel"]`` before importing DeepSeek's ``model.py``. Two uses: + +- CPU tests at toy size, with bf16 weights. The in-place quantizers stay no-ops there, because toy + head sizes are not multiples of the quantization blocks. +- The real FP8/FP4 checkpoint on a machine without the GPU kernels (the M4 Max). Set + ``exact_quant = True``. Every function is then a port of the kernel's math: the same block + scales, power-of-2 rounding and clamps. The GEMMs dequantize their operands and multiply in + fp32, so only the summation order differs from the kernels. + +``sparse_attn`` and ``hc_split_sinkhorn`` are always exact ports. +""" + +from __future__ import annotations + +import torch + + +exact_quant = False # True: in-place quantizers fake-quantize like the GPU kernels (needed for real weights) + +FP8_MAX = 448.0 +FP4_MAX = 6.0 +# e2m1 magnitudes by 3-bit code; bit 3 of the nibble is the sign +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + + +def _pow2_ceil(v: torch.Tensor) -> torch.Tensor: + """2 ** ceil(log2(v)) for positive fp32 v, by the kernel's bit trick (fast_round_scale).""" + bits = v.float().contiguous().view(torch.int32) + exp = ((bits >> 23) & 0xFF) - 127 + ((bits & 0x7FFFFF) != 0).to(torch.int32) + return ((exp + 127) << 23).view(torch.float32) + + +def _blocks(x: torch.Tensor, block_size: int) -> torch.Tensor: + assert x.size(-1) % block_size == 0, (x.shape, block_size) + return x.float().unflatten(-1, (-1, block_size)) # [..., n_blocks, block] + + +def act_quant(x, block_size=128, scale_fmt=None, scale_dtype=torch.float32, inplace=False): + """Block-wise FP8 e4m3 quantization of activations (``act_quant_kernel``).""" + if inplace and not exact_quant: + return x + xb = _blocks(x, block_size) + amax = xb.abs().amax(dim=-1).clamp_min(1e-4) + s = amax * torch.tensor(1 / FP8_MAX, dtype=torch.float32) + if scale_fmt is not None: + s = _pow2_ceil(s) + y = (xb / s[..., None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + if inplace: + x.copy_((y.float() * s[..., None]).flatten(-2).to(x.dtype)) + return x + return y.flatten(-2), s.to(scale_dtype) + + +def _e2m1_round(v: torch.Tensor) -> torch.Tensor: + """Nearest e2m1 value, ties to the even code (``T.Cast(FP4, ...)``); v already in [-6, 6].""" + mag = v.abs() + grid = E2M1.to(v.device) + dist = (mag[..., None] - grid).abs() + best = dist.amin(dim=-1, keepdim=True) + codes = torch.arange(8, device=v.device) + # among equally near codes take the even one (its mantissa bit is 0) + pick = torch.where(dist == best, codes % 2, 99).argmin(dim=-1) # ties are adjacent codes: one even + return torch.copysign(grid[pick], v) + + +def fp4_act_quant(x, block_size=32, inplace=False, scale_dtype=torch.float8_e8m0fnu): + """Block-wise FP4 fake quantization (``fp4_quant_kernel``); only the in-place form is used.""" + assert inplace, "model.py only calls fp4_act_quant in place" + if not exact_quant: + return x + xb = _blocks(x, block_size) + amax = xb.abs().amax(dim=-1) + if scale_dtype == torch.float8_e4m3fn: + s = (amax.clamp_min(6 * 2**-9) / FP4_MAX).to(torch.float8_e4m3fn).float() + else: + s = _pow2_ceil(amax.clamp_min(6 * 2**-126) * torch.tensor(1 / FP4_MAX, dtype=torch.float32)) + y = _e2m1_round((xb / s[..., None]).clamp(-FP4_MAX, FP4_MAX)) + x.copy_((y * s[..., None]).flatten(-2).to(x.dtype)) + return x + + +def _dequant_rows(values: torch.Tensor, scale: torch.Tensor, block: int) -> torch.Tensor: + """values [R, K] fp32 times a per-(row, K-block) scale [R, K/block].""" + return (values.unflatten(-1, (-1, block)) * scale.float()[..., None]).flatten(-2) + + +def decode_e2m1(b: torch.Tensor) -> torch.Tensor: + """float4_e2m1fn_x2 [N, K/2] to fp32 [N, K]; the low nibble is the first element.""" + codes = b.contiguous().view(torch.uint8) + nib = torch.stack((codes & 0x0F, codes >> 4), dim=-1).flatten(-2).long() + mag = E2M1.to(b.device)[nib & 0x7] + return torch.where((nib & 0x8) != 0, -mag, mag) + + +def fp8_gemm(a, a_s, b, b_s, scale_dtype=torch.float32, block_size=128): + """C[M,N] = A[M,K] @ B[N,K]^T; A scaled per (row, K block), B per (N block, K block).""" + K = a.size(-1) + A = _dequant_rows(a.reshape(-1, K).float(), a_s.reshape(-1, K // block_size), block_size) + sb = b_s.float().repeat_interleave(block_size, dim=0)[: b.size(0)] + B = _dequant_rows(b.float(), sb, block_size) + return (A @ B.T).to(torch.get_default_dtype()).view(*a.shape[:-1], b.size(0)) + + +def fp4_gemm(a, a_s, b, b_s, scale_dtype=torch.float32, act_block_size=128): + """C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T; B scaled per 32 along K (e8m0).""" + K = a.size(-1) + A = _dequant_rows(a.reshape(-1, K).float(), a_s.reshape(-1, K // act_block_size), act_block_size) + B = _dequant_rows(decode_e2m1(b), b_s, 32) + return (A @ B.T).to(torch.get_default_dtype()).view(*a.shape[:-1], b.size(0)) + + +def sparse_attn(q, kv, attn_sink, topk_idxs, softmax_scale): + """q [b,m,h,d], kv [b,n,d] (one shared head), topk_idxs [b,m,k] with -1 = empty slot. + Softmax over the gathered slots plus a per-head sink logit; a row with no valid slot is 0.""" + + b = q.shape[0] + valid = topk_idxs >= 0 + gathered = kv.float()[torch.arange(b)[:, None, None], topk_idxs.clamp_min(0).long()] # [b,m,k,d] + scores = torch.einsum("bmhd,bmkd->bmhk", q.float(), gathered) * softmax_scale + scores = scores.masked_fill(~valid[:, :, None, :], -torch.inf) + row_max = scores.amax(dim=-1, keepdim=True).clamp_min(-1e30) # the kernel's finite floor + p = torch.exp(scores - row_max) + denom = p.sum(dim=-1) + torch.exp(attn_sink.float()[None, None, :] - row_max[..., 0]) + o = torch.einsum("bmhk,bmkd->bmhd", p, gathered) / denom[..., None] + return o.to(q.dtype) + + +def hc_split_sinkhorn(mixes, hc_scale, hc_base, hc_mult=4, sinkhorn_iters=20, eps=1e-6): + hc = hc_mult + pre = torch.sigmoid(mixes[..., :hc] * hc_scale[0] + hc_base[:hc]) + eps + post = 2 * torch.sigmoid(mixes[..., hc : 2 * hc] * hc_scale[1] + hc_base[hc : 2 * hc]) + comb = (mixes[..., 2 * hc :] * hc_scale[2] + hc_base[2 * hc :]).unflatten(-1, (hc, hc)) + comb = comb.softmax(dim=-1) + eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + return pre, post, comb diff --git a/scripts/runpod/v41_reap/x1_pod.sh b/scripts/runpod/v41_reap/x1_pod.sh new file mode 100755 index 0000000..de7617e --- /dev/null +++ b/scripts/runpod/v41_reap/x1_pod.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# X1 on one GPU pod: "prune twice" at the first cut (384 -> 288 -> 192 -> 96), scored against the +# shipped REAP tables. Plan: docs/deepseek-v4-flash/2026-11-09-routed-precision-plan.md, section 7. +# +# The volume already holds everything (a CPU pod put it there, so the GPU starts work at once): +# /workspace/v41-hf the pinned checkpoint, fetch verified +# /workspace/pack the calibration + eval pack +# /workspace/wheels pip wheels for inference/requirements.txt +# /workspace/x1/*.json fm-288, fm-192 (total mass on full-model routing), survivors-192 (shipped +# REAP-50), pm-96 (the Mac's pick inside REAP-50) +set -euo pipefail +W=/workspace +KIT=$W/keep/scripts/runpod/v41_reap +HF=${HF:-$W/v41-hf} +X=$W/x1 +CAL_SHARD=${CAL_SHARD:-0/12} # 2 of 24 calibration sessions, about 256K tokens, every domain +EVAL_SHARD=${EVAL_SHARD:-0/2} # 54 of 108 eval sessions +CHUNK=${CHUNK_TOKENS:-262144} +export PYTHONPATH=$W/keep/src MAX_GPU_HOURS=${MAX_GPU_HOURS:-9} +export PIP_BREAK_SYSTEM_PACKAGES=1 # Ubuntu 24.04 images mark the system Python as externally managed +log() { printf '[%s] %s\n' "$(date -u +%FT%TZ)" "$*" | tee -a "$X/x1.log"; } +trap 'status=$?; [[ $status -ne 0 ]] && { log "FAILED (status $status)"; touch "$X/FAILED"; }' EXIT + +# The network volume reads one stream at about 120 MB/s and a random Engram row in about 3 ms, so the +# smoke test sat at 0% GPU for 50 minutes. Sixteen parallel streams read 3 GB/s: copy the checkpoint +# to the container's own disk first (LOCAL_HF, needs ~480 GB there) and run from that. +if [[ -n "${LOCAL_HF:-}" ]]; then + log "copy checkpoint to $LOCAL_HF" + mkdir -p "$LOCAL_HF" + (cd "$W/v41-hf" && find . -type f -not -path './.cache/*' -print0 | xargs -0 -P 16 -I{} cp --parents {} "$LOCAL_HF/") + [[ $(cd "$W/v41-hf" && find . -type f -not -path './.cache/*' -printf '%s\n' | awk '{s+=$1} END {print s}') == $(cd "$LOCAL_HF" && find . -type f -printf '%s\n' | awk '{s+=$1} END {print s}') ]] + HF=$LOCAL_HF + log "copy done" +fi +export HF + +pip install -q --no-index --find-links "$W/wheels" -r "$KIT/requirements-ref.txt" \ + || pip install -q -r "$KIT/requirements-ref.txt" +pip freeze > "$W/pip-freeze-gpu.txt" # pod.sh's gpu_setup then skips its own install +nvidia-smi --query-gpu=name,memory.total --format=csv | tee -a "$X/x1.log" + +log "smoke" +bash "$KIT/pod.sh" smoke + +calibrate() { # name, mask + log "calibration masked to $2 ($CAL_SHARD)" + python3 "$KIT/collect.py" --hf-path "$HF" --pack "$W/pack/v41-pack.jsonl" --out "$X/$1" --sweep calibration \ + --mask "$2" --shard "$CAL_SHARD" --chunk-tokens "$CHUNK" > "$X/cal-$1.log" 2>&1 + log "calibration $1 done: $(tail -1 "$X/cal-$1.log")" +} +calibrate m288 "$X/fm-288.json" +python3 "$KIT/rerank.py" --stats "$X/m288/calibration/sessions" --within "$X/fm-288.json" --keep 192 --out "$X/pt-192.json" | tee -a "$X/x1.log" +calibrate m192 "$X/pt-192.json" +python3 "$KIT/rerank.py" --stats "$X/m192/calibration/sessions" --within "$X/pt-192.json" --keep 96 --out "$X/pt-96.json" | tee -a "$X/x1.log" + +log "eval ($EVAL_SHARD): teacher + reap192, fm192, pt192, pm96, pt96" +python3 "$KIT/collect.py" --hf-path "$HF" --pack "$W/pack/v41-pack.jsonl" --out "$X/eval" --sweep eval \ + --shard "$EVAL_SHARD" --chunk-tokens "$CHUNK" \ + --survivors reap192="$X/survivors-192.json" --survivors fm192="$X/fm-192.json" --survivors pt192="$X/pt-192.json" \ + --survivors pm96="$X/pm-96.json" --survivors pt96="$X/pt-96.json" > "$X/eval.log" 2>&1 +python3 "$KIT/score.py" --teacher "$X/eval" --student reap192="$X/eval/student-reap192" --student fm192="$X/eval/student-fm192" \ + --student pt192="$X/eval/student-pt192" --student pm96="$X/eval/student-pm96" --student pt96="$X/eval/student-pt96" > "$X/scores.json" +log "scores written" +touch "$X/DONE" diff --git a/scripts/runpod/v41_reap/x1b_pod.sh b/scripts/runpod/v41_reap/x1b_pod.sh new file mode 100755 index 0000000..7686874 --- /dev/null +++ b/scripts/runpod/v41_reap/x1b_pod.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# X1 follow-up on the same pod, after x1_pod.sh: is 96 better picked in one pass on the full model's +# routing (fm96), or in two with the routing re-measured on the 192 it comes from (pf96)? +# Needs /workspace/x1/fm-96.json and the checkpoint already copied to $HF by x1_pod.sh. +set -euo pipefail +W=/workspace; KIT=$W/keep/scripts/runpod/v41_reap; X=$W/x1; HF=${HF:-/root/v41-hf} +export PYTHONPATH=$W/keep/src HF +log() { printf '[%s] %s\n' "$(date -u +%FT%TZ)" "$*" | tee -a "$X/x1.log"; } +trap 'status=$?; [[ $status -ne 0 ]] && { log "x1b FAILED (status $status)"; touch "$X/FAILED-B"; }' EXIT +log "x1b: calibration masked to fm-192" +python3 "$KIT/collect.py" --hf-path "$HF" --pack "$W/pack/v41-pack.jsonl" --out "$X/mfm192" --sweep calibration \ + --mask "$X/fm-192.json" --shard "${CAL_SHARD:-0/12}" --chunk-tokens 262144 > "$X/cal-mfm192.log" 2>&1 +python3 "$KIT/rerank.py" --stats "$X/mfm192/calibration/sessions" --within "$X/fm-192.json" --keep 96 --out "$X/pf-96.json" | tee -a "$X/x1.log" +log "x1b: eval teacher + fm96, pf96" +python3 "$KIT/collect.py" --hf-path "$HF" --pack "$W/pack/v41-pack.jsonl" --out "$X/eval2" --sweep eval \ + --shard "${EVAL_SHARD:-0/2}" --chunk-tokens 262144 \ + --survivors fm96="$X/fm-96.json" --survivors pf96="$X/pf-96.json" > "$X/eval2.log" 2>&1 +python3 "$KIT/score.py" --teacher "$X/eval2" --student fm96="$X/eval2/student-fm96" --student pf96="$X/eval2/student-pf96" > "$X/scores-b.json" +log "x1b scores written" +touch "$X/DONE-B" diff --git a/scripts/v41_build/bench_expert_formats.py b/scripts/v41_build/bench_expert_formats.py new file mode 100644 index 0000000..2bcf05f --- /dev/null +++ b/scripts/v41_build/bench_expert_formats.py @@ -0,0 +1,143 @@ +"""K2: one decode-step MoE at V4.1 shape, by routed-expert format, in one process. + +- oMLX's fused decode kernels (``decode_moe.fused_decode_moe``: router, top-k routed experts and + the shared expert in three dispatches) with routed experts as MXFP4 (the checkpoint's own + format) and as oMLX's affine codes at 2 and 3 bits, group 64. Random weights. +- This repo's E8P lattice kernel (``ramp.kernels.e8p_m1``; routed experts only, no router and + no shared expert) on real 2-bit codes of one layer. + +Each sample chains ``--chain`` dependent MoEs before one ``mx.eval``, as a decode step does. +Running each at top-k 6 and top-k 1 separates the cost per routed expert from the fixed cost. + + PYTHONPATH=src:benchmarks ~/Development/_refs/omlx/.venv/bin/python \ + scripts/v41_build/bench_expert_formats.py --config ~/Models/v41-reap50/config.json \ + --e8p ~/Models/v41-reap-run/quant/e8p-podimp/layer-20.npz +""" + +import argparse +import json +import statistics +import time +from dataclasses import replace +from pathlib import Path + +import mlx.core as mx +import numpy as np + +EXPERTS = 64 + + +def median_ms(step, chain, iterations=30, warmup=10): + def sample(): + started = time.perf_counter() + step(chain) + return (time.perf_counter() - started) * 1e3 / chain + + for _ in range(warmup): + sample() + return statistics.median(sample() for _ in range(iterations)) + + +def omlx_moe(config, mode, bits, group, rng): + from omlx.patches.deepseek_v41.language import MoE + from omlx.patches.deepseek_v41.quantization import QuantizedProjection + + def projection(out, inp, lead, mode, bits, group): + weight = mx.array(rng.integers(0, 1 << 32, size=(*lead, out, inp * bits // 32), dtype=np.uint32)) + if mode == "affine": + shape = (*lead, out, inp // group) + return QuantizedProjection(weight, mx.full(shape, 0.01, mx.bfloat16), bits, mode, + biases=mx.full(shape, -0.005 * (2**bits - 1), mx.bfloat16), group_size=group) + return QuantizedProjection(weight, mx.full((*lead, out, inp // 32), 118, mx.uint8), bits, mode) # 2^-9 + + moe = MoE(config) + dim, inter = config.dim, config.moe_inter_dim + for name, out, inp in (("w1", inter, dim), ("w3", inter, dim), ("w2", dim, inter)): + setattr(moe.experts, name, projection(out, inp, (EXPERTS,), mode, bits, group)) + setattr(moe.shared_experts, name, projection(out, inp, (), "mxfp8", 8, 32)) + moe.gate.weight = mx.array(rng.normal(size=(EXPERTS, dim)).astype(np.float32)).astype(mx.bfloat16) + moe.gate.bias = mx.zeros((EXPERTS,)) + arrays = [moe.gate.weight, moe.gate.bias] + nbytes = 0 + for name in ("w1", "w3", "w2"): + p, s = getattr(moe.experts, name), getattr(moe.shared_experts, name) + fields = [p.weight, p.scales] + ([p.biases] if "biases" in p else []) + arrays += fields + [s.weight, s.scales] + nbytes += sum(a.nbytes for a in fields) // EXPERTS + mx.eval(arrays) + return moe, nbytes + + +def bench_omlx(config, mode, bits, group, top_k, chain, seed=0): + from omlx.patches.deepseek_v41 import decode_moe + + rng = np.random.default_rng(seed) + c = replace(config, n_routed_experts=EXPERTS, n_activated_experts=top_k) + moe, expert_bytes = omlx_moe(c, mode, bits, group, rng) + x0 = mx.array(rng.normal(size=(1, 1, c.dim)).astype(np.float32)).astype(mx.bfloat16) + assert decode_moe.supported(moe, x0, c), (mode, bits, group) + state = {"x": x0} + + def step(n): + x = state["x"] + for _ in range(n): + out = decode_moe.fused_decode_moe(moe, x, c)[0].astype(mx.float32) + x = (out / (mx.max(mx.abs(out)) + 1e-6) * 3).astype(mx.bfloat16) # a new input, so new routes + mx.eval(x) + state["x"] = x + + ms = median_ms(step, chain) + assert bool(mx.all(mx.isfinite(state["x"].astype(mx.float32)))), "non-finite activations" + return ms, expert_bytes + + +def bench_e8p(path, top_k, chain, sweep=((4, 1), (8, 1), (4, 4), (2, 4))): + from bench_e8p_m1 import PROJECTIONS, projection_bytes_per_expert, run + from keep.vq.e8 import e8p_packed_abs_grid + from ramp.nn.switch_linear import QuantizedVQSwitchLinear + + z, layers = np.load(path), {} + for name, key in zip(PROJECTIONS, ("w1", "w3", "w2")): + codes, scales = z[f"{key}_codes"][:EXPERTS], z[f"{key}_scales"][:EXPERTS] + layers[name] = QuantizedVQSwitchLinear( + input_dims=codes.shape[2] * 8, output_dims=codes.shape[1], num_experts=EXPERTS, + codes=mx.array(codes), scales=mx.array(scales), codebook=mx.array(e8p_packed_abs_grid()), + group_size=codes.shape[2] * 8 // scales.shape[2], code_bits=16) + mx.eval([layer.parameters() for layer in layers.values()]) + result = run(layers, top_k=top_k, iterations=30, warmup=10, seed=0, sweep=sweep, ffns_per_sync=chain) + assert result["outputs_agree"], result["configs"] + best = result["best"] + return (best["median_ms_per_ffn"], sum(projection_bytes_per_expert(layers[p]) for p in PROJECTIONS), + (best["rows_per_simdgroup"], best["simdgroups"]), result["production_median_ms_per_ffn"]) + + +def main(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--config", type=Path, required=True) + p.add_argument("--e8p", type=Path, required=True) + p.add_argument("--chain", type=int, default=40) + p.add_argument("--out", type=Path) + args = p.parse_args() + from omlx.patches.deepseek_v41.config import ModelConfig + + config = ModelConfig.from_dict(json.loads(args.config.expanduser().read_text())) + rows = [] + for label, mode, bits, group in (("omlx mxfp4", "mxfp4", 4, 32), ("omlx affine 2-bit g64", "affine", 2, 64), + ("omlx affine 3-bit g64", "affine", 3, 64)): + (ms6, nbytes), (ms1, _) = (bench_omlx(config, mode, bits, group, k, args.chain) for k in (6, 1)) + rows.append({"format": label, "ms_top6": ms6, "ms_top1": ms1, "expert_mb": nbytes / 1e6}) + mx.clear_cache() + (ms6, nbytes, best, production), (ms1, *_) = (bench_e8p(args.e8p.expanduser(), k, args.chain) for k in (6, 1)) + rows.append({"format": "e8p lattice 2-bit (routed experts only)", "ms_top6": ms6, "ms_top1": ms1, + "expert_mb": nbytes / 1e6, "rows_per_simdgroup,simdgroups": best, "generic_kernel_ms_top6": production}) + for row in rows: + row["ms_per_expert"] = (row["ms_top6"] - row["ms_top1"]) / 5 + row["ms_fixed"] = row["ms_top1"] - row["ms_per_expert"] + row["expert_gb_per_s"] = row["expert_mb"] / 1e3 / (row["ms_per_expert"] / 1e3) + print(json.dumps({k: round(v, 4) if isinstance(v, float) else v for k, v in row.items()})) + if args.out: + args.out.write_text(json.dumps({"mlx": mx.__version__, "chain": args.chain, "experts": EXPERTS, "rows": rows}, indent=1)) + + +if __name__ == "__main__": + main() diff --git a/scripts/v41_build/d3_pressure.py b/scripts/v41_build/d3_pressure.py new file mode 100644 index 0000000..d7ab7ec --- /dev/null +++ b/scripts/v41_build/d3_pressure.py @@ -0,0 +1,95 @@ +"""D3: decode speed of a running server while another process takes memory. + +Sends one request after another to an OpenAI-compatible endpoint. After a baseline it allocates +incompressible memory in steps, holds it, lets it go and keeps measuring. Stepping stops early at +a critical pressure level or at ``--max-swap-gb`` of swap. Writes ``requests.jsonl`` (one line per +request) and ``samples.jsonl`` (every 5 s: memory taken here, swap, pageouts, compressor, free +pages, pressure level, the server's footprint) into ``--out``. Standard library only. + + python d3_pressure.py --pid SERVER_PID --request REQ.json --out DIR +""" + +import argparse +import json +import os +import re +import subprocess +import threading +import time +import urllib.request +from pathlib import Path + + +def sh(*cmd: str) -> str: + return subprocess.run(cmd, capture_output=True, text=True, check=False).stdout + + +def system(pid: int) -> dict: + vm = {k.strip('" '): int(v) for k, v in re.findall(r'^(.+?):\s+(\d+)\.$', sh("vm_stat"), re.MULTILINE)} + page = int(sh("sysctl", "-n", "hw.pagesize")) + swap = float(re.search(r"used = ([\d.]+)M", sh("sysctl", "-n", "vm.swapusage"))[1]) / 1024 + return {"level": int(sh("sysctl", "-n", "kern.memorystatus_vm_pressure_level")), "swap_gb": round(swap, 2), + "free_gb": round(vm["Pages free"] * page / 2**30, 1), + "compressor_gb": round(vm["Pages occupied by compressor"] * page / 2**30, 1), + "pageouts": vm["Pageouts"], "swapouts": vm["Swapouts"], + "server": sh("top", "-l", "1", "-pid", str(pid), "-stats", "mem").split()[-1]} + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--pid", type=int, required=True) + p.add_argument("--request", type=Path, required=True) + p.add_argument("--out", type=Path, required=True) + p.add_argument("--url", default="http://127.0.0.1:8017/v1/chat/completions") + p.add_argument("--step-gb", type=int, default=4) + p.add_argument("--max-gb", type=int, default=40) + p.add_argument("--max-swap-gb", type=float, default=24) + p.add_argument("--every", type=int, default=45, help="seconds per step") + p.add_argument("--rest", type=int, default=120, help="seconds of baseline, of hold and of recovery") + args = p.parse_args() + args.out.mkdir(parents=True, exist_ok=True) + body, t0, state, held = args.request.read_bytes(), time.time(), {"phase": "baseline", "run": True}, [] + + def requests() -> None: + with (args.out / "requests.jsonl").open("w") as f: + while state["run"]: + start, phase, taken = time.time() - t0, state["phase"], len(held) + req = urllib.request.Request(args.url, body, {"Content-Type": "application/json"}) + try: + usage = json.load(urllib.request.urlopen(req, timeout=1800))["usage"] + except Exception as e: # noqa: BLE001 (a dead server is a result) + usage = {"error": f"{type(e).__name__}: {e}"[:200]} + time.sleep(5) + f.write(json.dumps({"t": round(start, 1), "end": round(time.time() - t0, 1), "phase": phase, "taken_gb": taken, + **{k: usage.get(k) for k in ("completion_tokens", "generation_tokens_per_second", + "time_to_first_token", "error") if k in usage}}) + "\n") + f.flush() + + def samples() -> None: + with (args.out / "samples.jsonl").open("w") as f: + while state["run"]: + state["last"] = system(args.pid) + f.write(json.dumps({"t": round(time.time() - t0, 1), "phase": state["phase"], "taken_gb": len(held), + **state["last"]}) + "\n") + f.flush() + time.sleep(5) + + threads = [threading.Thread(target=fn, daemon=True) for fn in (requests, samples)] + for t in threads: + t.start() + time.sleep(args.rest) + state["phase"] = "step" + while len(held) < args.max_gb and state["last"]["level"] < 4 and state["last"]["swap_gb"] < args.max_swap_gb: + held.extend(os.urandom(1 << 30) for _ in range(args.step_gb)) # random: the compressor gains nothing + time.sleep(args.every) + state["phase"] = "hold" + time.sleep(args.rest) + held.clear() + state["phase"] = "recovery" + time.sleep(args.rest) + state["run"] = False + threads[0].join(timeout=300) + + +if __name__ == "__main__": + main() diff --git a/scripts/v41_build/k5_pagecache.py b/scripts/v41_build/k5_pagecache.py new file mode 100644 index 0000000..e2ee61a --- /dev/null +++ b/scripts/v41_build/k5_pagecache.py @@ -0,0 +1,122 @@ +"""K5: can macOS's page cache hold expert bytes that MLX computes on in place? + +The idea under test: map one file holding every expert, alias it as an MLX array, and let the +kernel decide which pages stay in memory. That only works if + +1. MLX aliases a file mapping (no copy into its own buffers), +2. using a small part of the aliased buffer on the GPU leaves the rest out of memory. + +Result on the M4 Max, mlx 0.32.4 (2026-11-11): 1 holds for every mapping mode (same address, no +page touched by aliasing). 2 fails: a kernel that reads 1 MiB of a 2 GiB aliased buffer brings +every page of it into memory. One buffer over all experts would therefore page in the whole +pool on first use, so the design is dead in this form; the cost of a cold read was not measured +because no read is cold after the first. + +Each trial writes its own file with F_NOCACHE, so its pages start out of the page cache +(checked with mincore). Prints one JSON line per measurement. + + ~/Development/_refs/omlx/.venv/bin/python scripts/v41_build/k5_pagecache.py DIR [--gb 2] +""" + +import argparse +import ctypes +import fcntl +import json +import mmap +import os +import time +from pathlib import Path + +import mlx.core as mx +import numpy as np + +PAGE = mmap.PAGESIZE +libc = ctypes.CDLL(None, use_errno=True) + + +def cold_file(path: Path, nbytes: int) -> None: + """``nbytes`` of uint16 noise, written past the page cache.""" + chunk = np.random.default_rng(0).integers(0, 1 << 16, 1 << 23, dtype=np.uint16).tobytes() + assert nbytes % len(chunk) == 0 and len(chunk) % PAGE == 0 + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC) + fcntl.fcntl(fd, fcntl.F_NOCACHE, 1) + for _ in range(nbytes // len(chunk)): + assert os.write(fd, chunk) == len(chunk) + os.fsync(fd) + os.close(fd) + + +def in_memory(m: np.memmap) -> np.ndarray: + """One flag per page of the mapping: is it in memory (mincore)?""" + vec = (ctypes.c_char * (m.nbytes // PAGE))() + if libc.mincore(ctypes.c_void_p(m.ctypes.data), ctypes.c_size_t(m.nbytes), vec): + raise OSError(ctypes.get_errno(), "mincore") + return (np.frombuffer(vec, np.uint8) & 1).astype(bool) + + +def alias(m: np.memmap) -> mx.array: + a = mx.asarray(m, copy=False) # raises if MLX would copy + mx.eval(a) + return a + + +def timed_sum(a: mx.array, start: int, count: int) -> float: + t = time.perf_counter() + mx.eval(a[start : start + count].astype(mx.float32).sum()) + return (time.perf_counter() - t) * 1e3 + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("dir", type=Path) + p.add_argument("--gb", type=int, default=2) + args = p.parse_args() + args.dir.mkdir(parents=True, exist_ok=True) + nbytes, mb = args.gb << 30, (1 << 20) // 2 # mb: uint16 elements per MiB + say = lambda **kw: print(json.dumps(kw), flush=True) + say(mlx=mx.__version__, page=PAGE, file_gb=args.gb, wired_limit=mx.set_wired_limit(0)) + mx.eval(mx.zeros(8).sum()) # Metal is up before anything is timed + + # 1. Aliasing, for each way of mapping the file. + modes = {} + for mode in ("r", "r+", "c"): + path = args.dir / f"alias-{mode.replace('+', 'w')}.bin" + cold_file(path, nbytes) + m = np.memmap(path, dtype=np.uint16, mode=mode) + cold = float(in_memory(m).mean()) + active = mx.get_active_memory() + try: + a = alias(m) + except Exception as e: # noqa: BLE001 (the message is the result) + say(test="alias", mode=mode, ok=False, error=f"{type(e).__name__}: {e}"[:200]) + continue + same = np.asarray(memoryview(a)).ctypes.data == m.ctypes.data + say(test="alias", mode=mode, ok=True, same_address=bool(same), cold_before=cold, + in_memory_after_alias=float(in_memory(m).mean()), mlx_active_added_mb=(mx.get_active_memory() - active) >> 20) + if same: + modes[mode] = path + del a, m + if not modes: + say(verdict="dead: MLX copies every mapping from Python; try allocator::make_buffer natively") + return + mode = next(iter(modes)) + + # 2. Does touching 1 MiB on the GPU bring the whole buffer into memory? + for where in (0, nbytes // 4): # element offsets: the start and the middle of the file + path = args.dir / f"touch-{where}.bin" # a new file each time: a rewritten one keeps its cached pages + cold_file(path, nbytes) + m = np.memmap(path, dtype=np.uint16, mode=mode) + a = alias(m) + before = in_memory(m) + ms = timed_sum(a, where, mb) + after = in_memory(m) + say(test="touch-1MiB", at_gb=where * 2 / (1 << 30), ms=round(ms, 2), in_memory_before=float(before.mean()), + in_memory_after=float(after.mean()), pages_after=int(after.sum()), pages_needed=(mb * 2) // PAGE) + del a, m + + for f in args.dir.glob("*.bin"): + f.unlink() + + +if __name__ == "__main__": + main() diff --git a/scripts/v41_build/omlx_compact.py b/scripts/v41_build/omlx_compact.py new file mode 100644 index 0000000..aa50072 --- /dev/null +++ b/scripts/v41_build/omlx_compact.py @@ -0,0 +1,112 @@ +"""A fixed resident expert set under stock oMLX, with no copy of the checkpoint (K1's speed floor). + +``make`` writes a model directory that links every file of a DeepSeek-V4.1 source checkpoint and +replaces two: ``config.json`` (``n_routed_experts`` = C) and the tensor index, which lists only +each layer's C hottest experts under the names 0..C-1. ``compact-map.json`` records which stored +tensor each of those names means and which router rows to keep. + +Any other first argument runs the oMLX CLI with three patches, all read-side: shard headers gain +the new names as aliases, and the router weight and biases are cut to the kept rows as they load. +The model then routes each token to its top experts among the C resident ones, every expert is in +memory, and oMLX's fused decode kernels apply (they do not with its own expert offload). + + python omlx_compact.py make SRC DST --experts 48 --mass reap50-pruned.npz + OMLX_COMPACT=DST/compact-map.json python omlx_compact.py serve --model-dir ... +""" + +import argparse +import json +import os +import re +import sys +from pathlib import Path + +EXPERT = re.compile(r"(layers\.(\d+)\.ffn\.experts)\.(\d+)\.(.+)$") +GATE = re.compile(r"layers\.(\d+)\.ffn\.gate\.(weight|bias|bias_vl)$") + + +def make(src: Path, dst: Path, experts: int, mass_file: Path) -> None: + import numpy as np + + mass = np.load(mass_file)["reap__gate_norm_sum"] # [layer, expert] routed gate mass + order = [np.argsort(-row, kind="stable").tolist() for row in mass] # every expert, hottest first + by_mass = [row[:experts] for row in order] + hot = [sorted(row) for row in by_mass] + rank = [[by.index(e) for e in row] for by, row in zip(by_mass, hot)] # 0 is the hottest + dst.mkdir(parents=True, exist_ok=True) + for item in src.iterdir(): + if item.name not in ("config.json", "model.safetensors.index.json") and not (dst / item.name).exists(): + (dst / item.name).symlink_to(item.resolve()) + config = json.loads((src / "config.json").read_text()) + text = config.get("text_config", config) + assert mass.shape == (text["num_hidden_layers"], text["n_routed_experts"]), mass.shape + text["n_routed_experts"] = experts + (dst / "config.json").write_text(json.dumps(config, indent=1)) + index = json.loads((src / "model.safetensors.index.json").read_text()) + weight_map, alias = {}, {} + for key, filename in index["weight_map"].items(): + m = EXPERT.match(key) + if m is None: + weight_map[key] = filename + elif int(m[3]) in hot[int(m[2])]: + new = f"{m[1]}.{hot[int(m[2])].index(int(m[3]))}.{m[4]}" + weight_map[new], alias[new] = filename, key + index["weight_map"] = weight_map + (dst / "model.safetensors.index.json").write_text(json.dumps(index)) + (dst / "compact-map.json").write_text(json.dumps( + {"experts": experts, "source": str(src.resolve()), "hot": hot, "rank": rank, "order": order, + "alias": alias})) + print(f"{dst}: {experts} experts per layer, {len(alias)} tensors renamed") + + +def patch(spec: dict, full: dict | None = None) -> None: + import mlx.core as mx + from omlx.patches.deepseek_v41 import loading, residency, storage + + alias, hot = spec["alias"], [mx.array(rows) for rows in spec["hot"]] + + def aliased(header): # a fresh dict: the stored names are read before any is overwritten + return {**header, **{new: header[old] for new, old in alias.items() if old in header}} + + opened, read_header, source_weights = storage.TensorFile.__init__, residency.header_with_offset, loading.iter_source_weights + headers = {} + + def tensor_file(self, path): + opened(self, path) + self.header = aliased(self.header) + + def header_with_offset(*args): + if args not in headers: + header, offset = read_header(*args) + headers[args] = aliased(header), offset + return headers[args] + + def iter_source_weights(*args, **kwargs): + for tensors, specs in source_weights(*args, **kwargs): + for key in tensors: + m = GATE.search(key) + if m: + if full is not None: # the whole router, for callers that track demand + full[int(m[1]), m[2]] = tensors[key] + tensors[key] = tensors[key][hot[int(m[1])]] + yield tensors, specs + + storage.TensorFile.__init__ = tensor_file + residency.header_with_offset = header_with_offset + loading.iter_source_weights = iter_source_weights + + +if __name__ == "__main__": + if sys.argv[1:2] == ["make"]: + p = argparse.ArgumentParser() + p.add_argument("src", type=Path) + p.add_argument("dst", type=Path) + p.add_argument("--experts", type=int, required=True) + p.add_argument("--mass", type=Path, required=True) + a = p.parse_args(sys.argv[2:]) + make(a.src.expanduser(), a.dst.expanduser(), a.experts, a.mass.expanduser()) + else: + patch(json.loads(Path(os.environ["OMLX_COMPACT"]).read_text())) + from omlx.cli import main + + sys.exit(main()) diff --git a/scripts/v41_build/omlx_counters.py b/scripts/v41_build/omlx_counters.py new file mode 100644 index 0000000..20adaf4 --- /dev/null +++ b/scripts/v41_build/omlx_counters.py @@ -0,0 +1,41 @@ +"""Run the oMLX CLI with the V4.1 expert-slot counters written to a file. + +oMLX's server logs expert-cache hits and misses per request only for its generic offload wrapper; +the V4.1 wrapper keeps the same counters but nothing reads them. This starts ``omlx`` unchanged +and appends one JSON line to ``$OMLX_COUNTERS`` whenever the totals move: + + OMLX_COUNTERS=/tmp/counters.jsonl python omlx_counters.py serve --model-dir ... +""" + +import json +import os +import sys +import threading +import time + +from omlx.cli import main +from omlx.patches.deepseek_v41 import moe_offload + +slots, made = [], moe_offload._ExpertSlots.__init__ + + +def init(self, *args, **kwargs): + made(self, *args, **kwargs) + slots.append(self) + + +def watch(path): + last = None + while True: + time.sleep(0.5) + now = [sum(getattr(s, f) for s in slots) for f in ("hits", "misses", "fetched_bytes")] + if slots and now != last: + last = now + with open(path, "a") as f: + f.write(json.dumps({"t": round(time.time(), 2), "layers": len(slots), "capacity": slots[0].plan.capacity, + "hits": now[0], "misses": now[1], "fetched_mb": round(now[2] / 1e6, 1)}) + "\n") + + +moe_offload._ExpertSlots.__init__ = init +threading.Thread(target=watch, args=(os.environ["OMLX_COUNTERS"],), daemon=True).start() +sys.exit(main()) diff --git a/scripts/v41_build/omlx_lattice.py b/scripts/v41_build/omlx_lattice.py new file mode 100644 index 0000000..e9100a6 --- /dev/null +++ b/scripts/v41_build/omlx_lattice.py @@ -0,0 +1,494 @@ +"""The no-wait expert cache (``omlx_routed.py``) with 2-bit lattice experts, under stock oMLX. + +The routed experts are this repo's E8P codes (2.04 bits per weight, 9.0 MB an expert against +18.8 MB at MXFP4), so about twice as many stay resident. oMLX loads a small model +(``omlx_compact.py make --experts 6``); after the load each layer's routed experts are replaced: + +- Storage. A slot is one row of six stacked arrays (codes and scales of w1, w3, w2). The bytes + come straight out of the quantizer's ``layer-NN.npz`` files, which are stored uncompressed, so + an install is a positional read into the slot's rows, as in the MXFP4 cache. +- Router. It keeps every row of the pool and scores all of it once per token (oMLX's gate + kernel). One small kernel then takes two top-k's of those scores: among the resident experts + (the bias of every other expert is -1e30), which is what runs, and in the whole pool, which + is the demand the cache policy reads later. Installing or evicting an expert changes its bias + and its entry in the expert-to-slot table; the router rows never move. +- Decode step. Router, then this repo's ``e8p_m1`` kernels for the routed experts, then oMLX's + own kernels for the shared expert, which also add the routed sum. +- Prefill. oMLX's unfused MoE, with the routed experts on this repo's sorted-route kernels. + + OMLX_COMPACT=DIR/compact-map.json OMLX_LATTICE=CODES_DIR python omlx_lattice.py serve ... + python omlx_lattice.py check --config CONFIG.json --codes CODES_DIR + +The pool can be wider than the loaded checkpoint. ``OMLX_LATTICE`` is then a JSON file:: + + {"codes": [[CODES_DIR, IDS.json], ...], "router": "router.safetensors", "mass": "reap-stats.npz"} + +``IDS.json`` (``{"survivors": {layer: [...]}}``) gives the pool number of each expert of its +directory, in file order; the first directory holds the loaded checkpoint's experts. ``router`` has +every router row of the pool and ``mass`` (``reap__gate_norm_sum``) orders the starting residents. +Paths in the file may be relative to it. + +Environment: ``OMLX_LATTICE_SLOTS`` (resident experts per layer; default ``auto``, from this +machine's RAM, see ``auto_slots``), ``OMLX_LATTICE_BUDGET_GIB`` (the footprint ``auto`` aims at, +default 70% of RAM) and the policy settings of ``omlx_routed.py``. ``check`` compares one layer's decode step with dense matrices +decoded on the host, the prefill path with the decode path, and slots after installs with their +stored bytes, then times the decode-step MoE. +""" + +import json +import os +import struct +import sys +import time +import zipfile +from dataclasses import replace +from functools import cache +from pathlib import Path +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +from mlx import nn + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src")) + +from omlx.patches.deepseek_v41 import decode_moe as DM +from omlx.patches.deepseek_v41 import language, loading +from omlx.patches.deepseek_v41.activation import quantize_swiglu_activation +from omlx.patches.deepseek_v41.quantization import quantize_activation +from omlx_compact import patch as compact_patch +from omlx_routed import OFF, Cache, Store + +from keep.vq.e8 import decode_weight_matrix, e8p_packed_abs_grid +from ramp.kernels.e8p_m1 import e8p_m1_per_route, e8p_m1_token_layout +from ramp.nn.switch_linear import QuantizedVQSwitchLinear + +TENSORS = [(proj, field) for proj in ("w1", "w3", "w2") for field in ("codes", "scales")] + +# One threadgroup per token. Top-k among the residents (slots and routing weights, as oMLX's +# up kernel computes them) and top-k in the whole pool (expert numbers and raw scores). +_TOPK = r""" + threadgroup int chosen[8]; + const uint tid = thread_index_in_threadgroup, b = threadgroup_position_in_grid.x; + const device float* sc = scores + b * NE; + topk_select(sc, bias, chosen, tid); + if (tid == 0) { + float wsum = 0.0f; + for (int j = 0; j < TOPK; ++j) wsum += sc[chosen[j]]; + for (int j = 0; j < TOPK; ++j) { + slot[b * TOPK + j] = slot_of[chosen[j]]; + weight[b * TOPK + j] = sc[chosen[j]] / (wsum + 1e-20f) * route_scale[0]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + topk_select(sc, pool_bias, chosen, tid); + if (tid == 0) + for (int j = 0; j < TOPK; ++j) { + want[b * TOPK + j] = uint(chosen[j]); + want_score[b * TOPK + j] = sc[chosen[j]]; + } +""" + + +@cache +def _topk_kernel(): + return mx.fast.metal_kernel( + name="v41_lattice_topk", input_names=["scores", "bias", "pool_bias", "slot_of", "route_scale"], + output_names=["slot", "weight", "want", "want_score"], source=_TOPK, header=DM._HEADER) + + +@cache +def _unused(dtype): + """Stands in for the routed MXFP4 inputs of oMLX's up kernel, which runs no routed column.""" + a = mx.zeros((8,), dtype) + mx.eval(a) + return a + + +class Codes: + """Where each expert's codes and scales sit inside the ``layer-NN.npz`` files of ``roots``. + ``where[layer][expert]`` is (directory, position in its file); without it, (0, expert).""" + + def __init__(self, roots: list[Path], layers: int, where=None): + self.files, self.where = [[] for _ in range(layers)], where # [layer][directory] + for i, root in ((i, root) for i in range(layers) for root in roots): + path = root / f"layer-{i:02d}.npz" + fd, members = os.open(path, os.O_RDONLY), {} + with zipfile.ZipFile(path) as z: + for info in z.infolist(): + assert info.compress_type == zipfile.ZIP_STORED, f"{path}: {info.filename} is compressed" + name, extra = struct.unpack(" tuple[int, int, int]: + d, expert = self.where[layer][expert] if self.where else (0, expert) + fd, members = self.files[layer][d] + start, shape, dtype = members[name] + size = int(np.prod(shape[1:])) * dtype.itemsize + return fd, start + expert * size, size + + def read(self, layer: int, name: str, experts, pool=None) -> np.ndarray: + _, shape, dtype = self.files[layer][0][1][name] + host = np.empty((len(experts), *shape[1:]), dtype) + jobs = [(*self.slab(layer, name, e), host[s]) for s, e in enumerate(experts)] + for done in [pool.submit(Store.pread_into, *job) for job in jobs] if pool else jobs: + done.result() if pool else Store.pread_into(*done) + return host + + +class Experts(nn.Module): + """oMLX's routed ``Expert`` for the unfused path; ``indices`` are slots.""" + + quantizes_input = True + + def __init__(self, limit): + super().__init__() + self._limit = limit + + def __call__(self, x, indices, weights, sorted_indices=False, *, input_quantized=False): + # sorted: x [routes, 1, dim], indices [routes]; otherwise x [..., 1, 1, dim], indices [..., k] + lead = indices.shape if sorted_indices else indices.shape[:-1] + rows = x.reshape(*lead, x.shape[-1]) + gate = self.w1(rows, indices, sorted_indices=sorted_indices) + up = self.w3(rows, indices, sorted_indices=sorted_indices) + h = quantize_swiglu_activation(gate, up, weights, x.dtype, self._limit) + return self.w2(h, indices, sorted_indices=sorted_indices)[..., None, :].astype(x.dtype) + + +class Layer: + def __init__(self, index, moe, experts, weight, bias, order, slots): + assert weight.dtype == mx.bfloat16, weight.dtype # the gate kernel reads bf16 rows + self.index, self.gate, self.experts = index, moe.gate, experts + self.full_weight, self.full_bias = weight.astype(mx.float32), bias.astype(mx.float32) + mx.eval(self.full_weight, self.full_bias) + self.pool_bias = np.array(self.full_bias) + resident = [int(e) for e in order[:slots]] # hottest first + self.slot_expert = np.array(resident, dtype=np.int64) # -1: unrouted + self.slot_of = {e: s for s, e in enumerate(resident)} + self.last_use = -1.0 - np.arange(slots, dtype=np.float64) # colder starting experts go first + self.cooling: list[tuple[int, int]] = [] + self.free: list[int] = [] + self.reading: dict[int, tuple[int, list]] = {} + masked, slot_map = np.full(self.pool_bias.size, OFF, np.float32), np.zeros(self.pool_bias.size, np.uint32) + masked[resident], slot_map[resident] = self.pool_bias[resident], np.arange(slots) + self.slot_map = mx.array(slot_map) # expert -> slot, read only for experts whose bias is on + moe.gate.weight, moe.gate.bias, moe.experts = weight, mx.array(masked), experts + self.views = {} + for proj, field in TENSORS: + lin = getattr(experts, proj) + self.views[proj, field] = (lin, lin[field], np.asarray(memoryview(lin[field]))) + assert self.views_current() + + def views_current(self) -> bool: + """The host views still alias the arrays the kernels read.""" + return all(lin[field] is owner and np.frombuffer(memoryview(owner), np.uint8).ctypes.data == view.ctypes.data + for (_, field), (lin, owner, view) in self.views.items()) + + +class LatticeCache(Cache): + def __init__(self, model, config, roots: list[Path], order, full, slots: int, where=None): + super().__init__(config, mx) + blocks = model.language_model.layers + self.codes = Codes(roots, len(blocks), where) + self.codebook = mx.array(e8p_packed_abs_grid()) + layers = [] + for i, blk in enumerate(blocks): + experts = Experts(config.swiglu_limit) + for proj in ("w1", "w3", "w2"): + codes, scales = (mx.array(self.codes.read(i, f"{proj}_{field}", order[i][:slots], self.pool)) + for field in ("codes", "scales")) + mx.eval(codes, scales) + setattr(experts, proj, QuantizedVQSwitchLinear( + input_dims=codes.shape[2] * 8, output_dims=codes.shape[1], num_experts=slots, codes=codes, + scales=scales, codebook=self.codebook, group_size=codes.shape[2] * 8 // scales.shape[2], + code_bits=16)) + layers.append(Layer(i, blk.ffn, experts, full[i, "weight"], full[i, "bias"], order[i], slots)) + self.pool_config = replace(config, n_routed_experts=layers[0].pool_bias.size) + self.ready(layers) + + def slabs(self, layer, expert, slot): + for (proj, field), (_, _, view) in layer.views.items(): + yield (*self.codes.slab(layer.index, f"{proj}_{field}", expert), view[slot]) + + def route(self, layer, expert, slot): + layer.slot_map[expert] = slot + layer.gate.bias[expert] = float(layer.pool_bias[expert]) + + def unroute(self, layer, expert, slot): + layer.gate.bias[expert] = OFF + + +def decode(cache, layer, moe, x, config, mma, stock_gate): + """One MoE for the decode step. x: [tokens, dim] bf16. Returns ([tokens, dim] bf16, slots, weights).""" + if layer.index == 0: + cache.begin() + nb, dim, inter, k = x.shape[0], config.dim, config.moe_inter_dim, config.n_activated_experts + scores, xq = stock_gate(moe, x, cache.pool_config) # every expert of the pool; act_quant(x) + slot, weight, want, want_score = _topk_kernel()( + inputs=[scores, moe.gate.bias, layer.full_bias, layer.slot_map, DM._const((config.route_scale,))], + template=[("NE", cache.pool_config.n_routed_experts), ("TOPK", k)], + grid=(nb * 32, 1, 1), threadgroup=(32, 1, 1), + output_shapes=[(nb * k,)] * 4, output_dtypes=[mx.uint32, mx.float32, mx.uint32, mx.float32]) + cache.pending.setdefault(cache.step, []).append((want, want_score)) # as evaluated: reading them later costs no eval + e, s = layer.experts, moe.shared_experts + # 16-bit activations: the lattice kernel is 1.7x slower on fp32 rows, and act_quant values are exact in bf16. + rows, slots = xq.reshape(nb, dim).astype(mx.bfloat16), slot.reshape(nb, k) + gate = e8p_m1_token_layout(rows, e.w1.codes, e.w1.scales, cache.codebook, slots, output_dims=inter) + up = e8p_m1_token_layout(rows, e.w3.codes, e.w3.scales, cache.codebook, slots, output_dims=inter) + h = quantize_swiglu_activation(gate, up, weight, mx.bfloat16, config.swiglu_limit) + down = e8p_m1_per_route(h.reshape(nb * k, inter), e.w2.codes, e.w2.scales, cache.codebook, slot, output_dims=dim) + routed = down.astype(mx.float32).reshape(nb, k, dim).sum(1).reshape(-1) # bf16 per expert, fp32 sum + # The shared expert on oMLX's kernels: its up kernel with no routed column (TOPK 0), then its + # shared down kernel, which adds the routed sum and rounds to bf16. + shared = tuple(DM.weight_format(getattr(s, n)) for n in ("w1", "w3", "w2")) + up_k, _, down_s = DM._kernels_nb(mma and all(f[0] in ("mxfp8", "affine") for f in shared), (DM.MXFP4,) * 3 + shared) + no_w, no_s, no_b = _unused(mx.uint32), _unused(mx.uint8), _unused(mx.float32) + hqs = up_k( + inputs=[no_b, no_b, no_b, DM._const((config.swiglu_limit,)), xq, no_w, no_s, no_b, no_w, no_s, no_b, + s.w1.weight, s.w1.scales, DM._biases(s.w1), s.w3.weight, s.w3.scales, DM._biases(s.w3)], + template=[("NE", 1), ("TOPK", 0), ("DIM", dim), ("INTER", inter), ("NB", nb)], + grid=(256, inter // 32, 1), threadgroup=(256, 1, 1), + output_shapes=[(1,), (nb * inter,), (1,), (1,)], + output_dtypes=[mx.float32, mx.float32, mx.int32, mx.float32])[1] + out = down_s( + inputs=[hqs, routed, s.w2.weight, s.w2.scales, DM._biases(s.w2)], + template=[("DIM", dim), ("INTER", inter), ("NB", nb)], + grid=(dim // 8 * 64, 1, 1), threadgroup=(64, 1, 1), + output_shapes=[(nb * dim,)], output_dtypes=[mx.bfloat16])[0] + return out.reshape(nb, dim), slots, weight.reshape(nb, k) + + +def hook(state: dict) -> None: + """Send the layers of ``state["cache"]`` through the lattice paths; every other MoE is oMLX's.""" + stock = SimpleNamespace(one=DM.fused_decode_moe, many=DM.fused_decode_moe_nb, supported=DM.supported, + gate=DM._gate, gate_call=language.Gate.__call__) + + def find(gate): + cache = state.get("cache") + return cache, cache.by_gate.get(id(gate)) if cache else None + + def supported(moe, x, config): + layer = find(moe.gate)[1] + return stock.supported(moe, x, config) if layer is None else x.size == config.dim and x.dtype == mx.bfloat16 + + def one(moe, x, config): + cache, layer = find(moe.gate) + if layer is None: + return stock.one(moe, x, config) + out, slots, weights = decode(cache, layer, moe, x.reshape(1, -1), config, False, stock.gate) + return out.reshape(x.shape), slots, weights + + def many(moe, x, config, mma=None): + cache, layer = find(moe.gate) + if layer is None: + return stock.many(moe, x, config, mma) + return decode(cache, layer, moe, x, config, DM.use_mma(x.shape[0]) if mma is None else mma, stock.gate) + + def gate_plain(self, x, image_mask): # prefill and every unfused path + cache, layer = find(self) + if layer is None: + return stock.gate_call(self, x, image_mask) + if layer.index == 0: + cache.begin() + index, weights = stock.gate_call(self, x, image_mask) # pool numbers of resident experts + return cache.observe(self, x, layer.slot_map[index]), weights + + DM.fused_decode_moe, DM.fused_decode_moe_nb, DM.supported = one, many, supported + language.Gate.__call__ = gate_plain + + +def wide_pool(path: Path, loaded: dict): + """(directories, order, router rows, where) of the pool file ``path``. ``loaded`` is the + router of the loaded checkpoint, whose rows must be the first directory's.""" + spec = json.loads(path.read_text()) + + def at(p: str) -> Path: # paths in the pool file are relative to it + return path.parent / Path(p).expanduser() + + ids = [json.loads(at(f).read_text())["survivors"] for _, f in spec["codes"]] + where = [{e: (d, j) for d, per in enumerate(ids) for j, e in enumerate(per[str(i)])} for i in range(len(ids[0]))] + mass = np.load(at(spec["mass"]))["reap__gate_norm_sum"] + assert all(sorted(w) == list(range(mass.shape[1])) for w in where), "the directories do not cover the pool once" + router = mx.load(str(at(spec["router"]))) + full = {(i, f): router[f"layers.{i}.ffn.gate.{f}"] for i in range(len(where)) for f in ("weight", "bias")} + for (i, f), rows in loaded.items(): + if f in ("weight", "bias"): + assert mx.array_equal(full[i, f][mx.array(ids[0][str(i)])], rows), f"layer {i}: router {f} differs" + order = [np.argsort(-row, kind="stable").tolist() for row in mass] + return [at(d) for d, _ in spec["codes"]], order, full, where + + +FLOOR = 120 # residents per layer below which sampled code came out misspelled (build plan, C2) + + +def auto_slots(ram_bytes: int, drafter: bool, budget_gib: float | None = None) -> int: + """Resident experts per layer for a footprint of ``budget_gib`` (default: 70% of RAM). + + Measured on an M4 Max, 128 GB (2026-11-12): the server's footprint peaks near 20 GiB plus + 0.336 GiB per resident, and 8 GiB more with the drafter (87 GiB at 200 residents, 88 GiB at + 180 with the drafter). Other chips are unmeasured; the budget is the knob.""" + budget = 0.7 * ram_bytes / 2**30 if budget_gib is None else budget_gib + return int((budget - 20 - 8 * drafter) / 0.336) + + +def drafter_on(argv: list[str]) -> bool: + """Does ``--base-path``'s model_settings.json switch the drafter on for any model?""" + if "--base-path" not in argv: + return False + settings = Path(argv[argv.index("--base-path") + 1]).expanduser() / "model_settings.json" + return settings.exists() and any(m.get("mtp_enabled") for m in json.loads(settings.read_text())["models"].values()) + + +def patch(spec: dict, root: Path, slots: int) -> None: + full, state = {}, {} + compact_patch(spec, full) + load = loading.load + + def load_lattice(*args, **kwargs): + model, processor = load(*args, **kwargs) + config = model.language_model.layers[0].ffn.gate._config + if root.is_file(): + roots, order, wide, where = wide_pool(root, full) + state["cache"] = LatticeCache(model, config, roots, order, wide, min(slots, len(order[0])), where) + else: + state["cache"] = LatticeCache(model, config, [root], spec["order"], full, min(slots, len(spec["order"][0]))) + full.clear() + return model, processor + + loading.load = load_lattice + hook(state) + + +def check(config_path: Path, root: Path, slots: int = 32, chain: int = 40) -> None: + from bench_expert_formats import median_ms, omlx_moe + from omlx.patches.deepseek_v41.config import ModelConfig + + rng = np.random.default_rng(0) + c = replace(ModelConfig.from_dict(json.loads(config_path.read_text())), n_routed_experts=64) + moes = [omlx_moe(c, "mxfp4", 4, 32, rng)[0] for _ in range(2)] + for moe in moes: # shared experts: random MXFP8 bytes below the codes for inf and NaN + for name in ("w1", "w3", "w2"): + shape = getattr(moe.shared_experts, name).weight.shape + codes = rng.integers(0, 0x60, size=(*shape, 4), dtype=np.uint8) | (rng.integers(0, 2, size=(*shape, 4), dtype=np.uint8) << 7) + getattr(moe.shared_experts, name).weight = mx.array(codes.view(np.uint32).reshape(shape)) + pool = 192 + full = {} + for i in range(2): + full[i, "weight"] = mx.array(rng.normal(size=(pool, c.dim)).astype(np.float32)).astype(mx.bfloat16) + full[i, "bias"] = mx.array(rng.normal(size=(pool,)).astype(np.float32) * 0.1) + os.environ.setdefault("OMLX_ROUTED_TAU", "inf") + model = SimpleNamespace(language_model=SimpleNamespace(layers=[SimpleNamespace(ffn=m) for m in moes])) + roots, where = [root], None + if root.is_file(): # a pool file: both directories, the first two layers + roots, _, _, where = wide_pool(root, {}) + pool = len(where[0]) + for i in range(2): + full[i, "weight"] = mx.array(rng.normal(size=(pool, c.dim)).astype(np.float32)).astype(mx.bfloat16) + full[i, "bias"] = mx.array(rng.normal(size=(pool,)).astype(np.float32) * 0.1) + state = {"cache": LatticeCache(model, c, roots, [rng.permutation(pool).tolist() for _ in moes], full, slots, where)} + cache = state["cache"] + hook(state) + + def token(): + return mx.array(rng.normal(size=(c.dim,)).astype(np.float32)).astype(mx.bfloat16) + + def dense(layer, proj, expert): + return decode_weight_matrix(*(cache.codes.read(layer, f"{proj}_{f}", [expert])[0] for f in ("codes", "scales")), + code_bits=16) + + def host(a): + return np.array(a.astype(mx.float32)) + + def reference(i, x, slots_, weights): + """The same MoE from matrices decoded on the host (and oMLX's shared expert, unfused), + rounded to bf16 where oMLX's kernels round.""" + xq, limit, total = host(quantize_activation(x.reshape(1, -1)))[0], c.swiglu_limit, 0.0 + for slot, weight in zip(slots_, weights): + expert = int(cache.layers[i].slot_expert[slot]) + g, u = bf16(dense(i, "w1", expert) @ xq), bf16(dense(i, "w3", expert) @ xq) + if limit: + g, u = np.minimum(g, limit), np.clip(u, -limit, limit) + h = bf16(g / (1 + np.exp(-g)) * u * weight) + total = total + bf16(dense(i, "w2", expert) @ host(quantize_activation(mx.array(h[None])))[0]) + shared = moes[i].shared_experts(quantize_activation(x.reshape(1, 1, -1)), input_quantized=True) + return total + host(shared).reshape(-1) + + def bf16(a): + return host(mx.array(a.astype(np.float32)).astype(mx.bfloat16)) + + def gap(a, b): + return float(np.max(np.abs(a - b)) / np.max(np.abs(b))) + + def against_reference(i): + x = token() + out, slots_, weights = DM.fused_decode_moe(moes[i], x.reshape(1, 1, -1), c) + return gap(host(out).reshape(-1), reference(i, x, np.array(slots_).reshape(-1), np.array(weights).reshape(-1))) + + worst = max(against_reference(i) for i in (0, 1) for _ in range(2)) + print(f"decode step against host-decoded matrices: worst relative gap {worst:.4f}") + assert worst < 2e-2 + + x = mx.stack([token() for _ in range(40)])[None] # [1, 40, dim]: the sorted-route prefill path + decoded = host(DM.fused_decode_moe_nb(moes[0], x[0, :8], c)[0]) + worst = max(gap(host(moes[0](x[:, :t], None))[0, :8][: min(t, 8)], decoded[: min(t, 8)]) for t in (40, 4)) + print(f"prefill path against decode path: worst relative gap {worst:.4f}") + assert worst < 2e-2 + + cache.tau = 0.0 # admit everything demanded + for step in range(24): + for moe in moes: + mx.eval(DM.fused_decode_moe(moe, token().reshape(1, 1, -1), c)[0]) + time.sleep(0.05) + print("after 24 tokens at tau 0:", {k: cache.stats[k] for k in ("wanted", "admitted", "installed", "evicted", "no_slot")}) + time.sleep(0.5) + assert cache.stats["installed"] > 0 and cache.mismatched() == 0 + cache.tau = float("inf") + worst = max(against_reference(i) for i in (0, 1) for _ in range(3)) + print(f"after installs, against host-decoded matrices: worst relative gap {worst:.4f}; no slot differs from its bytes") + assert worst < 2e-2 + + state_x = {"x": token().reshape(1, 1, -1)} + + def step(n): + x = state_x["x"] + for j in range(n): + out = DM.fused_decode_moe(moes[j % 2], x, c)[0].astype(mx.float32) + x = (out / (mx.max(mx.abs(out)) + 1e-6) * 3).astype(mx.bfloat16) + mx.eval(x) + state_x["x"] = x + + print(f"decode-step MoE (router over {pool}, {c.n_activated_experts} lattice experts, shared expert): " + f"{median_ms(step, chain):.3f} ms") + + +if __name__ == "__main__": + if sys.argv[1:2] == ["check"]: + import argparse + + p = argparse.ArgumentParser() + p.add_argument("--config", type=Path, required=True) + p.add_argument("--codes", type=Path, required=True) + a = p.parse_args(sys.argv[2:]) + check(a.config.expanduser(), a.codes.expanduser()) + else: + assert auto_slots(128 << 30, True) == 183 and auto_slots(128 << 30, False, 87) == 199 + slots = os.environ.get("OMLX_LATTICE_SLOTS", "auto") + if slots == "auto": + budget = os.environ.get("OMLX_LATTICE_BUDGET_GIB") + slots = auto_slots(os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE"), drafter_on(sys.argv), + float(budget) if budget else None) + slots = int(slots) + print(f"lattice cache: {slots} resident experts per layer" + ( + f"; below {FLOOR}, generated code was unreliable in our tests" if slots < FLOOR else ""), file=sys.stderr) + assert slots >= 6, "not enough memory for the router's top-6" + patch(json.loads(Path(os.environ["OMLX_COMPACT"]).read_text()), Path(os.environ["OMLX_LATTICE"]).expanduser(), slots) + from omlx.cli import main + + sys.exit(main()) diff --git a/scripts/v41_build/omlx_routed.py b/scripts/v41_build/omlx_routed.py new file mode 100644 index 0000000..34dcde6 --- /dev/null +++ b/scripts/v41_build/omlx_routed.py @@ -0,0 +1,330 @@ +"""An expert cache that never waits, under stock oMLX: a miss re-routes, the fetch lands later. + +Starts from ``omlx_compact.py``'s model (C resident experts per layer as plain stacked arrays, the +router cut to their rows, so oMLX's fused decode kernels apply) and adds three things: + +- Demand. Each layer also scores the whole pool's router rows on the same input; the top experts + and their gate weights ride along in the token's graph (``mx.depends``), so reading them later + costs no wait. +- A cache policy, run once per forward call before its first layer: an expert demanded with gate + weight >= tau that is not resident is admitted into a free slot; the least recently demanded + residents are evicted to keep ``spare`` slots free. +- No-wait installs. A slot's row in the stacked weight arrays is overwritten in place by reader + threads (the stored bytes are the array's bytes). The router is the switch: an evicted slot's + bias goes to -1e30 first, its bytes are overwritten only once every forward built before that + has finished, and its router row is set to the new expert only after the bytes are in. A token + never waits for a read and never runs a half-written expert. + + OMLX_COMPACT=DIR/compact-map.json OMLX_ROUTED_LOG=log.jsonl python omlx_routed.py serve ... + +Environment: ``OMLX_ROUTED_TAU`` (default 0.2; ``inf`` keeps the starting set and must reproduce +the fixed set exactly), ``OMLX_ROUTED_SPARE`` (free slots kept per layer, default 4), +``OMLX_ROUTED_LAG`` (forward calls between a demand and its read-out, default 2: the call in +flight is never waited for), ``OMLX_ROUTED_VERIFY`` (each log line also counts routed slots whose +bytes differ from their expert's stored bytes; slow, for testing). +""" + +import json +import os +import struct +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np + +from omlx_compact import patch as compact_patch + +PROJECTIONS = ("w1", "w2", "w3") +OFF = -1e30 # a router bias no score can overcome + + +class Store: + """Positional reads of one expert's stored tensors from the source checkpoint.""" + + def __init__(self, source: Path): + self.source = source + self.index = json.loads((source / "model.safetensors.index.json").read_text())["weight_map"] + self.files: dict[str, tuple[int, dict, int]] = {} + + def slab(self, key: str) -> tuple[int, int, int]: + name = self.index[key] + if name not in self.files: + fd = os.open(self.source / name, os.O_RDONLY) + size = struct.unpack(" int: + """Fill ``view`` from the file. Positional, so reader threads may call it.""" + target = memoryview(view.reshape(-1).view(np.uint8)) + assert target.nbytes == nbytes, (target.nbytes, nbytes) + done = 0 + while done < nbytes: + count = os.preadv(fd, [target[done:]], offset + done) + assert count > 0, "truncated tensor" + done += count + return nbytes + + +class Layer: + def __init__(self, index, moe, full_weight, full_bias, hot, rank, mx): + self.index, self.gate, self.experts = index, moe.gate, moe.experts + self.full_weight = full_weight.astype(mx.float32) # [pool, dim] + self.full_bias = full_bias.astype(mx.float32) + self.typed_weight, self.typed_bias = full_weight, full_bias # the rows a slot's router entry is copied from + mx.eval(self.full_weight, self.full_bias) + self.slot_expert = np.array(hot, dtype=np.int64) # -1: unrouted + self.slot_of = {e: s for s, e in enumerate(hot)} + self.last_use = -1.0 - np.array(rank, dtype=np.float64) # colder starting experts go first + self.cooling: list[tuple[int, int]] = [] # (slot, forward call that unrouted it) + self.free: list[int] = [] + self.reading: dict[int, tuple[int, list]] = {} # expert -> (slot, futures) + self.views = {} + for proj in PROJECTIONS: + lin = getattr(self.experts, proj) + assert lin["weight"].dtype == mx.uint32 and lin["scales"].dtype == mx.uint8, proj + self.views[proj] = (lin, lin["weight"], lin["scales"], + np.asarray(memoryview(lin["weight"])), np.asarray(memoryview(lin["scales"]))) + + def views_current(self) -> bool: + """The host views still alias the arrays the kernels read.""" + for lin, weight, scales, weight_view, scales_view in self.views.values(): + for owner, name, view in ((weight, "weight", weight_view), (scales, "scales", scales_view)): + if lin[name] is not owner or np.frombuffer(memoryview(owner), np.uint8).ctypes.data != view.ctypes.data: + return False + return True + + +class Cache: + """The policy. A subclass sets ``layers`` and says where a slot's bytes live (``slabs``) and + how the router switches an expert on and off (``route``, ``unroute``).""" + + def __init__(self, config, mx): + self.mx, self.config = mx, config + self.tau = float(os.environ.get("OMLX_ROUTED_TAU", "0.2")) + self.spare = int(os.environ.get("OMLX_ROUTED_SPARE", "4")) + self.lag = int(os.environ.get("OMLX_ROUTED_LAG", "2")) + self.log_path = os.environ.get("OMLX_ROUTED_LOG") + assert config.score_func == "sqrtsoftplus" and self.lag >= 1 + self.pool = ThreadPoolExecutor(max_workers=8, thread_name_prefix="routed-io") + self.step = 0 + self.pending: dict[int, list] = {} # forward call -> per-layer (index, weight) demand arrays + self.stats = dict.fromkeys(("calls", "tokens", "demand", "resident", "wanted", "admitted", "installed", + "evicted", "bytes", "no_slot", "policy_s"), 0) + self.started = time.time() + + def ready(self, layers): + self.layers = layers + self.by_gate = {id(layer.gate): layer for layer in layers} + # A resident expert's rows must equal its stored bytes: installs rely on it. + for layer in layers[:: max(1, len(layers) // 4)]: + assert not self.differing(layer, 0), f"layer {layer.index}: slot 0 differs from its stored bytes" + + def slabs(self, layer, expert, slot): + """(fd, offset, nbytes, the slot's host view) for each stored tensor of ``expert``.""" + raise NotImplementedError + + def route(self, layer, expert, slot): + raise NotImplementedError + + def unroute(self, layer, expert, slot): + raise NotImplementedError + + def differing(self, layer, slot) -> int: + """Tensors of a routed slot whose bytes differ from its expert's stored bytes.""" + bad = 0 + for fd, offset, nbytes, view in self.slabs(layer, int(layer.slot_expert[slot]), slot): + stored = np.empty(nbytes, np.uint8) + Store.pread_into(fd, offset, nbytes, stored) + bad += not np.array_equal(stored, view.reshape(-1).view(np.uint8)) + return bad + + # --- demand, recorded while a forward call is built ------------------------------------- + + def demand(self, layer, x): + mx, c = self.mx, self.config + raw = (x.reshape(-1, x.shape[-1]).astype(mx.float32) @ layer.full_weight.T) / c.gate_temp + scores = mx.sqrt(mx.logaddexp(raw, 0)) + index = mx.argpartition(-(scores + layer.full_bias), kth=c.n_activated_experts - 1, axis=-1) + index = index[..., : c.n_activated_experts] + return index, mx.take_along_axis(scores, index, -1) + + def observe(self, gate, x, tied): + """Record this layer's demand and tie it to ``tied`` so one eval covers both.""" + layer = self.by_gate[id(gate)] + index, weight = self.demand(layer, x) + self.pending.setdefault(self.step, []).append((index, weight)) + tied = self.mx.depends(tied, [index, weight]) + return tied[0] if isinstance(tied, (list, tuple)) else tied + + # --- policy, run before the first layer of each forward call ---------------------------- + + def begin(self): + self.step += 1 + ready = self.step - self.lag + if self.tau != float("inf") and ready in self.pending: + started = time.perf_counter() + self.harvest(ready, self.pending.pop(ready)) + self.stats["policy_s"] += time.perf_counter() - started + for old in [s for s in self.pending if s < ready]: + del self.pending[old] + if self.log_path and self.step % 64 == 0: + self.write_log() + + def harvest(self, call, demand): + c = self.config + self.stats["calls"] += 1 + # Reading this call's demand means every earlier call has left the GPU. + for layer, (index, weight) in zip(self.layers, demand): + index = np.asarray(index).reshape(-1, c.n_activated_experts) + weight = np.asarray(weight, dtype=np.float64).reshape(index.shape) + weight = weight / (weight.sum(-1, keepdims=True) + 1e-20) * c.route_scale + if layer.index == 0: + self.stats["tokens"] += index.shape[0] + self.stats["demand"] += index.size + strong = weight >= self.tau + wanted: dict[int, float] = {} + for expert, w in zip(index[strong].tolist(), weight[strong].tolist()): + wanted[expert] = max(w, wanted.get(expert, 0.0)) + self.stats["resident"] += sum(e in layer.slot_of for e in index.reshape(-1).tolist()) + self.install(layer) + missing = [] + for expert, w in wanted.items(): + slot = layer.slot_of.get(expert) + if slot is not None: + layer.last_use[slot] = call + elif expert not in layer.reading: + missing.append((w, expert)) + self.stats["wanted"] += len(missing) + layer.free += [slot for slot, since in layer.cooling if since <= call] + layer.cooling = [(slot, since) for slot, since in layer.cooling if since > call] + if missing and layer.free and not layer.views_current(): + raise RuntimeError(f"layer {layer.index}: expert arrays moved; cannot install in place") + for _, expert in sorted(missing, reverse=True): + if not layer.free: + self.stats["no_slot"] += 1 + continue + slot = layer.free.pop() + futures = [self.pool.submit(Store.pread_into, *slab) for slab in self.slabs(layer, expert, slot)] + layer.reading[expert] = (slot, futures) + self.stats["admitted"] += 1 + self.evict(layer, set(wanted)) + + def install(self, layer) -> None: + """Route the experts whose bytes have landed. The router arrays are replaced, never + written in place and never evaluated here: calls already built keep the old ones, and + the next call's own eval computes the new ones without a wait.""" + done = [e for e, (_, futures) in layer.reading.items() if all(f.done() for f in futures)] + for expert in done: + slot, futures = layer.reading.pop(expert) + self.stats["bytes"] += sum(f.result() for f in futures) + self.route(layer, expert, slot) + layer.slot_expert[slot], layer.slot_of[expert] = expert, slot + layer.last_use[slot] = self.step + self.stats["installed"] += 1 + + def evict(self, layer, protected) -> None: + """Unroute the least recently demanded residents until ``spare`` slots are free or coming.""" + short = self.spare - len(layer.free) - len(layer.cooling) + if short <= 0: + return + order = np.argsort(layer.last_use, kind="stable") + victims = [int(s) for s in order if layer.slot_expert[s] >= 0 and int(layer.slot_expert[s]) not in protected] + for slot in victims[:short]: + expert = int(layer.slot_expert[slot]) + del layer.slot_of[expert] + layer.slot_expert[slot] = -1 + layer.last_use[slot] = np.inf + self.unroute(layer, expert, slot) + # Calls built before this one may still route here; its bytes may change only once + # the call being built now is known to have finished. + layer.cooling.append((slot, self.step)) + self.stats["evicted"] += 1 + + def mismatched(self) -> int: + """Routed slots (first and last layer) whose rows differ from their expert's stored bytes.""" + return sum(self.differing(layer, slot) for layer in (self.layers[0], self.layers[-1]) + for slot, expert in enumerate(layer.slot_expert.tolist()) if expert >= 0) + + def write_log(self): + row = {"t": round(time.time() - self.started, 2), "step": self.step, **self.stats, + "reading": sum(len(layer.reading) for layer in self.layers)} + if os.environ.get("OMLX_ROUTED_VERIFY"): + row["mismatched"] = self.mismatched() + with open(self.log_path, "a") as f: + f.write(json.dumps(row) + "\n") + + +class FP4Cache(Cache): + """Slots are rows of oMLX's stacked MXFP4 arrays; the router holds one row per slot.""" + + def __init__(self, model, config, spec, full, mx): + super().__init__(config, mx) + self.store = Store(Path(spec["source"])) + for key in {name: key for key, name in self.store.index.items() if ".ffn.experts." in key}.values(): + self.store.slab(key) # open every expert shard now, not during a token + self.ready([Layer(i, blk.ffn, full[i, "weight"], full[i, "bias"], spec["hot"][i], spec["rank"][i], mx) + for i, blk in enumerate(model.language_model.layers)]) + + def slabs(self, layer, expert, slot): + for proj in PROJECTIONS: + base = f"layers.{layer.index}.ffn.experts.{expert}.{proj}." + for field, view in (("weight", layer.views[proj][3]), ("scale", layer.views[proj][4])): + yield (*self.store.slab(base + field), view[slot]) + + def route(self, layer, expert, slot): + layer.gate.weight[slot] = layer.typed_weight[expert] + layer.gate.bias[slot] = layer.typed_bias[expert] + + def unroute(self, layer, expert, slot): + layer.gate.bias[slot] = OFF + + +def patch(spec: dict) -> None: + import mlx.core as mx + from omlx.patches.deepseek_v41 import decode_moe, language, loading + + full: dict = {} + compact_patch(spec, full) + state: dict = {} + load, fused_gate, gate_call = loading.load, decode_moe._gate, language.Gate.__call__ + + def load_routed(*args, **kwargs): + model, processor = load(*args, **kwargs) + state["cache"] = FP4Cache(model, model.language_model.layers[0].ffn.gate._config, spec, full, mx) + full.clear() + return model, processor + + def gate_fused(moe, x, config): # one-token and batched fused decode + cache = state.get("cache") + if cache is None or id(moe.gate) not in cache.by_gate: + return fused_gate(moe, x, config) + if cache.by_gate[id(moe.gate)].index == 0: + cache.begin() + scores, quantized = fused_gate(moe, x, config) + return cache.observe(moe.gate, x, scores), quantized + + def gate_plain(self, x, image_mask): # prefill and every unfused path + cache = state.get("cache") + if cache is None or id(self) not in cache.by_gate: + return gate_call(self, x, image_mask) + if cache.by_gate[id(self)].index == 0: + cache.begin() + index, weights = gate_call(self, x, image_mask) + return cache.observe(self, x, index), weights + + loading.load = load_routed + decode_moe._gate = gate_fused + language.Gate.__call__ = gate_plain + + +if __name__ == "__main__": + patch(json.loads(Path(os.environ["OMLX_COMPACT"]).read_text())) + from omlx.cli import main + + sys.exit(main()) diff --git a/scripts/v41_build/run_v41.sh b/scripts/v41_build/run_v41.sh new file mode 100755 index 0000000..473e0f4 --- /dev/null +++ b/scripts/v41_build/run_v41.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# Serve DeepSeek-V4.1-Flash-1017 from a pool bundle with the no-wait expert cache, under oMLX at a +# pinned commit: every one of the 384 experts per layer reachable, as many resident as RAM allows. +# +# run_v41.sh BUNDLE [PORT] +# +# BUNDLE holds models/v41-lattice (slim_checkpoint.py output), pool.json and what it names. The +# first run clones oMLX into OMLX_DIR (default ~/.cache/v41/omlx) and builds its venv with uv. +# OMLX_LATTICE_SLOTS and OMLX_LATTICE_BUDGET_GIB override the resident count (omlx_lattice.py). +set -euo pipefail +REPO=$(cd "$(dirname "$0")/../.." && pwd) +B=$(cd "${1:?usage: run_v41.sh BUNDLE [PORT]}" && pwd); PORT=${2:-8017} +OMLX=${OMLX_DIR:-$HOME/.cache/v41/omlx}; PIN=4213c3284abc7b3746a91464a7af7ccc3732b759 +[ "$(uname -sm)" = "Darwin arm64" ] || { echo "needs an Apple-silicon Mac" >&2; exit 1; } +RAM=$(( $(sysctl -n hw.memsize) >> 30 )) +[ "$RAM" -ge 96 ] || echo "warning: $RAM GB of RAM; generated code was unreliable below 96 GB in our tests" >&2 +for f in models/v41-lattice/compact-map.json pool.json; do + [ -f "$B/$f" ] || { echo "$B/$f is missing" >&2; exit 1; } +done +if [ ! -f "$B/base/model_settings.json" ]; then # the drafter costs 8 GiB: on from 128 GB of RAM + mkdir -p "$B/base" + [ "$RAM" -ge 120 ] && MTP=', "mtp_enabled": true, "mtp_fixed_depth": 3' || MTP= + echo "{\"version\": 2, \"models\": {\"v41-lattice\": {\"deepseek_v41_engram_ssd_offload\": true$MTP}}}" > "$B/base/model_settings.json" +fi +if [ ! -x "$OMLX/.venv/bin/python" ]; then + command -v git > /dev/null && command -v uv > /dev/null || { echo "needs git and uv (https://docs.astral.sh/uv/)" >&2; exit 1; } + [ -d "$OMLX/.git" ] || git clone --quiet https://github.com/jundot/omlx "$OMLX" + git -C "$OMLX" checkout --quiet "$PIN" + (cd "$OMLX" && uv sync --quiet) +fi +echo "endpoint http://127.0.0.1:$PORT/v1, model v41-lattice" >&2 +cd "$OMLX" +# tau 0.3 is the admission threshold every measurement in the build plan used (the code's default is 0.2) +OMLX_COMPACT=$B/models/v41-lattice/compact-map.json OMLX_LATTICE=$B/pool.json OMLX_ROUTED_TAU=${OMLX_ROUTED_TAU:-0.3} \ + exec .venv/bin/python "$REPO/scripts/v41_build/omlx_lattice.py" serve --model-dir "$B/models" --base-path "$B/base" \ + --port "$PORT" --memory-guard-gb $((RAM * 3 / 4)) --max-concurrent-requests 1 --no-cache --headless diff --git a/scripts/v41_build/slim_checkpoint.py b/scripts/v41_build/slim_checkpoint.py new file mode 100644 index 0000000..1a97267 --- /dev/null +++ b/scripts/v41_build/slim_checkpoint.py @@ -0,0 +1,102 @@ +"""A compact model directory as real files, without the routed experts its loader never reads. + +SRC is an ``omlx_compact.py make`` directory: config, tensor index and ``compact-map.json`` beside +links to a checkpoint's shards. The lattice cache (``omlx_lattice.py``) loads only the few experts +the map names and takes every other routed expert from its codes, so the rest of the FP4 experts +are dead weight in a download. DST gets every file of SRC; a shard that loses nothing is cloned +(``cp -c``, free on APFS) and the others are rewritten with the tensors that stay, names and bytes +untouched. Each rewritten shard is read back and compared with the source by CRC. Standard library +only. + + python slim_checkpoint.py SRC DST +""" + +import json +import os +import re +import shutil +import struct +import subprocess +import sys +import zlib +from pathlib import Path + +EXPERT = re.compile(r"layers\.\d+\.ffn\.experts\.\d+\.") # a routed expert of a decoder layer (not the drafter's) +CHUNK = 1 << 24 + + +def header(fd: int) -> tuple[dict, int]: + (size,) = struct.unpack(" int: + """CRC of ``size`` bytes at ``begin``; with ``out``, the bytes are also written to it.""" + value, done = 0, 0 + while done < size: + chunk = os.pread(fd, min(size - done, CHUNK), begin + done) + assert chunk, "truncated shard" + value, done = zlib.crc32(chunk, value), done + len(chunk) + if out: + out.write(chunk) + return value + + +def slim(src: Path, dst: Path, keep: set[str]) -> tuple[int, int]: + """Write shard ``src`` to ``dst`` without its routed experts outside ``keep``. (bytes in, bytes out).""" + fd = os.open(src, os.O_RDONLY) + entries, start = header(fd) + names = [n for n in entries if n != "__metadata__" and (n in keep or not EXPERT.match(n))] + if len(names) == len(entries) - ("__metadata__" in entries): + os.close(fd) + if subprocess.run(["cp", "-c", str(src), str(dst)], capture_output=True, check=False).returncode: # not APFS: a plain copy + shutil.copyfile(src, dst) + return src.stat().st_size, src.stat().st_size + names.sort(key=lambda n: entries[n]["data_offsets"][0]) + out_header, offset = {k: entries[k] for k in ("__metadata__",) if k in entries}, 0 + for n in names: + begin, end = entries[n]["data_offsets"] + out_header[n] = {**entries[n], "data_offsets": [offset, offset + end - begin]} + offset += end - begin + blob = json.dumps(out_header, separators=(",", ":")).encode() + blob += b" " * (-len(blob) % 8) + with open(dst, "wb") as out: + out.write(struct.pack(" None: + src, dst = Path(sys.argv[1]).expanduser(), Path(sys.argv[2]).expanduser() + spec = json.loads((src / "compact-map.json").read_text()) + keep = set(spec["alias"].values()) # stored names of the experts the loader reads + dst.mkdir(parents=True) + total_in = total_out = 0 + for item in sorted(src.iterdir()): + real = item.resolve() + if item.name in (".cache", "SHA256SUMS"): # download bookkeeping and sums of the full shards + continue + if item.name == "compact-map.json": + (dst / item.name).write_text(json.dumps(spec | {"source": "."})) + elif item.name.endswith(".safetensors") and item.name.startswith("model-"): + a, b = slim(real, dst / item.name, keep) + total_in, total_out = total_in + a, total_out + b + print(f"{item.name}: {a / 1e9:.2f} GB -> {b / 1e9:.2f} GB", flush=True) + elif real.is_dir(): + shutil.copytree(real, dst / item.name) + else: + shutil.copyfile(real, dst / item.name) + print(f"{dst}: {total_in / 1e9:.1f} GB -> {total_out / 1e9:.1f} GB") + + +if __name__ == "__main__": + main() diff --git a/src/keep/convert/fp8_block.py b/src/keep/convert/fp8_block.py index a0f0b83..031945f 100644 --- a/src/keep/convert/fp8_block.py +++ b/src/keep/convert/fp8_block.py @@ -19,13 +19,27 @@ Pure numpy so it runs headless; the MLX fast path can come later if decode throughput matters. + +:func:`repack_fp8_block_as_mxfp8` is the lossless alternative to decoding: it +re-lays the same e4m3 bytes and ue8m0 exponents out in MLX's ``mxfp8`` layout +(group 32 along the input axis), so a resident can stay at 8.25 bits per weight +instead of expanding to bf16, with every dequantized value bit-identical. """ from __future__ import annotations import numpy as np -__all__ = ["decode_e4m3", "decode_ue8m0", "dequantize_fp8_block"] +__all__ = [ + "MXFP8_GROUP_SIZE", + "decode_e4m3", + "decode_ue8m0", + "dequantize_fp8_block", + "repack_fp8_block_as_mxfp8", +] + +#: MLX ``mxfp8`` supports exactly one group size: 32 elements per E8M0 scale. +MXFP8_GROUP_SIZE = 32 _E4M3_TABLE: np.ndarray | None = None @@ -130,3 +144,112 @@ def dequantize_fp8_block( c0, c1 = j * block_cols, min((j + 1) * block_cols, cols) decoded[r0:r1, c0:c1] *= row_scales[j] return decoded + + +def _power_of_two_scales_as_ue8m0(scales: np.ndarray) -> np.ndarray: + """Encode literal power-of-two multipliers as ue8m0 bytes, or refuse. + + Only exact powers of two in ``2**-127 .. 2**127`` have an E8M0 encoding. + Anything else would need rounding, and rounding is exactly what a lossless + repack must never do, so it raises instead. + """ + + values = scales.astype(np.float64) + if not np.all(np.isfinite(values)) or np.any(values <= 0.0): + raise ValueError("power-of-two block scales must be finite and positive") + mantissa, exponent = np.frexp(values) + if np.any(mantissa != 0.5): + raise ValueError( + "block scales are not all exact powers of two; an mxfp8 repack " + "would have to round them, so it is refused" + ) + unbiased = exponent.astype(np.int64) - 1 + if np.any(unbiased < -127) or np.any(unbiased > 127): + raise ValueError("power-of-two block scales are outside the E8M0 range") + return (unbiased + 127).astype(np.uint8) + + +def repack_fp8_block_as_mxfp8( + codes: np.ndarray, + scales: np.ndarray, + *, + block_size: tuple[int, int] = (128, 128), + group_size: int = MXFP8_GROUP_SIZE, +) -> tuple[np.ndarray, np.ndarray]: + """Re-lay block-scaled e4m3 weights out as MLX ``mxfp8`` without changing a bit. + + Returns ``(packed, group_scales)``: + + - ``packed``: ``uint32`` of shape ``(rows, cols // 4)``, the unchanged e4m3 + code bytes read little-endian four to a word, which is the element order + ``mx.dequantize(..., mode="mxfp8")`` expects. + - ``group_scales``: ``uint8`` of shape ``(rows, cols // 32)``, each block's + ue8m0 exponent repeated over the rows and the 32-wide column groups the + block covers. + + Dequantizing the result reproduces :func:`dequantize_fp8_block` exactly: + e4m3 times a power of two is representable in float32 and bf16 alike. + Building the same weights with ``mx.quantize(..., mode="mxfp8")`` instead is + *not* exact -- it re-derives each group's exponent by rounding -- so this + function is the only supported way to produce an mxfp8 resident. + + Requires ``cols % 32 == 0`` and ``block_size[1] % 32 == 0`` so every + 32-wide group sits inside one source block. ue8m0 byte 255 (NaN) is + refused rather than bound. + """ + + if not isinstance(codes, np.ndarray): + raise ValueError(f"codes must be a numpy ndarray, got {type(codes).__name__}") + if not isinstance(scales, np.ndarray): + raise ValueError(f"scales must be a numpy ndarray, got {type(scales).__name__}") + if codes.dtype != np.uint8: + raise ValueError(f"e4m3 codes must be uint8, got {codes.dtype}") + if codes.ndim != 2: + raise ValueError(f"codes must be 2-D, got shape {codes.shape}") + if group_size != MXFP8_GROUP_SIZE: + raise ValueError( + f"mxfp8 has a fixed group size of {MXFP8_GROUP_SIZE}, got {group_size!r}" + ) + block_rows, block_cols = block_size + for label, value in (("block_size[0]", block_rows), ("block_size[1]", block_cols)): + if ( + isinstance(value, bool) + or not isinstance(value, (int, np.integer)) + or value <= 0 + ): + raise ValueError(f"{label} must be a positive int, got {value!r}") + rows, cols = codes.shape + if cols % group_size != 0: + raise ValueError( + f"input dimension {cols} is not a multiple of the mxfp8 group size {group_size}" + ) + if block_cols % group_size != 0: + raise ValueError( + f"block width {block_cols} is not a multiple of the mxfp8 group size " + f"{group_size}, so a group would straddle two source scales" + ) + expected = (-(-rows // block_rows), -(-cols // block_cols)) + if scales.shape != expected: + raise ValueError( + f"scales shape {scales.shape} does not match " + f"{expected} blocks for weight {codes.shape} " + f"at block size {block_size}" + ) + if scales.dtype == np.uint8: + if np.any(scales == 255): + raise ValueError("ue8m0 byte 255 is NaN; refusing to repack a NaN block scale") + exponents = scales + elif np.issubdtype(scales.dtype, np.floating): + exponents = _power_of_two_scales_as_ue8m0(scales) + else: + raise ValueError( + f"scales must be uint8 (ue8m0-encoded) or a floating dtype holding " + f"powers of two, got dtype {scales.dtype}" + ) + + # ``astype`` copies, so the result never aliases the caller's code buffer. + packed = np.ascontiguousarray(codes).view(" 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/e8_1bit_m1.metal b/src/ramp/kernels/e8_1bit_m1.metal new file mode 100644 index 0000000..2fc8efc --- /dev/null +++ b/src/ramp/kernels/e8_1bit_m1.metal @@ -0,0 +1,90 @@ +// Few-token gather matvec for 8-bit E8 codes (the 1-bit residual of a nested +// tier), the companion of e8p_m1.metal: one threadgroup column per route, +// route r multiplies activation row r / X_DIV by expert rhs_indices[r]. A +// route whose index is 0xFFFFFFFF has no residual and writes zeros. +// +// A codeword is one byte that indexes a 256-entry table of 8 values. The +// table arrives with nibbles k and k+4 holding dims 2k and 2k+1 (see +// pair_packed in e8_1bit_m1.py), so each threadgroup decodes it into half2 +// pairs once and a codeword then costs one threadgroup read. A lane takes +// CHUNK (4 or 8) consecutive codewords with one 4- or 8-byte load per row. +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint expert = uint(rhs_indices[route]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = (route / uint(X_DIV)) * codewords * 8u; +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); + +if (expert == 0xFFFFFFFFu) { + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + if (simd_lane == 0u && row0 + r < out_dim) { + out[route * out_dim + row0 + r] = OUT_T(0.0f); + } + } + return; +} + +// (1024 + n) * 0.5 - 516 = (n - 8) / 2, exact in half. +threadgroup uint4 decoded[256]; +for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + uint packed = codebook[idx]; + uint4 entry; + for (uint k = 0; k < 4u; ++k) { + half2 h = as_type(((packed >> (4u * k)) & 0x000F000Fu) | 0x64006400u); + entry[k] = as_type(fma(h, half2(0.5h), half2(-516.0h))); + } + decoded[idx] = entry; +} +threadgroup_barrier(mem_flags::mem_threadgroup); + +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float2 xv[CHUNK * 4]; + for (uint i = 0; i < uint(CHUNK); ++i) { + for (uint k = 0; k < 4u; ++k) { + uint xi = x_row + (cw0 + i) * 8u + 2u * k; + xv[i * 4u + k] = float2(float(x[xi]), float(x[xi + 1u])); + } + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint2 words = uint2(0u); + if (uint(CHUNK) == 8u) { + words = ((const device uint2*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.x = ((const device uint*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint4 entry = decoded[(words[i >> 2u] >> ((i & 3u) * 8u)) & 0xFFu]; + for (uint k = 0; k < 4u; ++k) { + part += dot(float2(as_type(entry[k])), xv[i * 4u + k]); + } + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/src/ramp/kernels/e8_1bit_m1.py b/src/ramp/kernels/e8_1bit_m1.py new file mode 100644 index 0000000..ad882d4 --- /dev/null +++ b/src/ramp/kernels/e8_1bit_m1.py @@ -0,0 +1,97 @@ +"""Few-token kernels for 8-bit E8 codes, the 1-bit residual of a nested tier. + +The companion of ``e8p_m1``: the same two layouts over ``uint8`` codes, for a residual that is +added to the output of the 16-bit base. A route whose index is ``NO_RESIDUAL`` writes zeros, so +one call serves a token whose routes are partly at the higher tier. + +- :func:`e8_1bit_m1_token_layout` for gate/up (one activation row, top-k experts). +- :func:`e8_1bit_m1_per_route` for down (one activation row per route). +""" + +from __future__ import annotations + +from functools import cache +from pathlib import Path + +import numpy as np + +from ramp.kernels.e8p_m1 import _chunk + +__all__ = ["NO_RESIDUAL", "e8_1bit_m1_per_route", "e8_1bit_m1_token_layout", "pair_packed"] + +NO_RESIDUAL = 0xFFFFFFFF + + +def pair_packed(packed: np.ndarray) -> np.ndarray: + """The packed 1-bit table (nibble d is dim d) with dims 2k and 2k+1 in nibbles k and k+4, + the layout the kernel decodes into half2 pairs.""" + + p, out = np.asarray(packed, dtype=np.uint32), np.zeros(256, np.uint32) + for k in range(4): + out |= ((p >> np.uint32(8 * k)) & np.uint32(0xF)) << np.uint32(4 * k) + out |= ((p >> np.uint32(8 * k + 4)) & np.uint32(0xF)) << np.uint32(16 + 4 * k) + return out + + +@cache +def _kernel(): + import mlx.core as mx + + return mx.fast.metal_kernel( + name="ramp_e8_1bit_m1", + input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], + source=(Path(__file__).resolve().parent / "e8_1bit_m1.metal").read_text(), + ) + + +def _launch(x, codes, scales, codebook, rhs_indices, *, output_dims, rows_per_simdgroup, simdgroups, x_div): + import mlx.core as mx + + from ramp.kernels.vq_qmv import _qmv_output_dtype + + if codes.dtype != mx.uint8: + raise ValueError(f"8-bit E8 codes must be uint8, got {codes.dtype}") + if rows_per_simdgroup < 1 or simdgroups < 1 or simdgroups > 32: + raise ValueError("rows_per_simdgroup must be >= 1 and simdgroups in 1..32") + routes, rows_per_tile = int(rhs_indices.size), rows_per_simdgroup * simdgroups + output_dtype = _qmv_output_dtype(x.dtype) + (out,) = _kernel()( + inputs=[mx.contiguous(x), codes, scales, codebook, rhs_indices.reshape(-1)], + template=[ + ("ROWS_PER_SG", rows_per_simdgroup), + ("SIMDGROUPS", simdgroups), + ("CHUNK", _chunk(int(codes.shape[2]), int(scales.shape[2]))), + ("X_DIV", int(x_div)), + ("OUT_T", output_dtype), + ], + grid=(32 * simdgroups, (output_dims + rows_per_tile - 1) // rows_per_tile, routes), + threadgroup=(32 * simdgroups, 1, 1), + output_shapes=[(1, routes, output_dims)], + output_dtypes=[output_dtype], + ) + return out + + +def e8_1bit_m1_token_layout(x, codes, scales, codebook, rhs_indices, *, output_dims: int, + rows_per_simdgroup: int = 4, simdgroups: int = 1): + """Token ``t``'s row against its ``top_k`` experts: ``[tokens, top_k, out]``.""" + + if x.ndim != 2 or rhs_indices.ndim != 2 or rhs_indices.shape[0] != x.shape[0]: + raise ValueError("token-layout E8 needs x [tokens, in] and rhs_indices [tokens, top_k]") + tokens, top_k = int(rhs_indices.shape[0]), int(rhs_indices.shape[1]) + out = _launch(x, codes, scales, codebook, rhs_indices, output_dims=output_dims, + rows_per_simdgroup=rows_per_simdgroup, simdgroups=simdgroups, x_div=top_k) + return out.reshape(tokens, top_k, output_dims) + + +def e8_1bit_m1_per_route(x, codes, scales, codebook, rhs_indices, *, output_dims: int, + rows_per_simdgroup: int = 4, simdgroups: int = 1): + """Route ``r``'s activation row against expert ``rhs_indices[r]``: ``[routes, out]``.""" + + routes = int(rhs_indices.size) + if x.ndim != 2 or x.shape[0] != routes: + raise ValueError(f"per-route E8 needs x [{routes}, in], got {tuple(x.shape)}") + out = _launch(x, codes, scales, codebook, rhs_indices, output_dims=output_dims, + rows_per_simdgroup=rows_per_simdgroup, simdgroups=simdgroups, x_div=1) + return out.reshape(routes, output_dims) diff --git a/src/ramp/kernels/e8p_m1.metal b/src/ramp/kernels/e8p_m1.metal new file mode 100644 index 0000000..e3c2797 --- /dev/null +++ b/src/ramp/kernels/e8p_m1.metal @@ -0,0 +1,104 @@ +// Few-token (M=1 decode, M<=32 verify) gather matvec for 16-bit E8P codes on +// GPUs without NAX. One threadgroup column per route: route r multiplies +// activation row r / X_DIV by expert rhs_indices[r]. +// +// Each simdgroup owns ROWS_PER_SG output rows. A lane takes CHUNK (4 or 8) +// consecutive codewords: one 8- or 16-byte load per row, so a simdgroup reads +// a contiguous 32 * CHUNK * 2-byte span per row per step, and the lane's +// activations are loaded once and reused for every row. +// +// E8P layout (see mlx_vq_decode_e8p_value in common_e8.metal): codeword bits +// 0-7 are signs, bits 8-15 index the 256-entry abs grid of packed nibbles. +// Dimension d reads nibble {0,4,1,5,2,6,3,7}[d], so nibbles k and k+4 are +// dims 2k and 2k+1 and decode together as one half2. Bit p of the effective +// signs negates nibble p, and every value gets +-0.25 by parity. +// +// Decode is the bottleneck (loads alone run at ~410 GB/s on an M4 Max), so +// each threadgroup first decodes the abs grid into half2 pairs, and a +// codeword then costs one threadgroup read plus one XOR per pair for signs. +// +// X_DIV=top_k: token t's routes share activation row t (gate/up). +// X_DIV=1: every route has its own activation row (down). +uint simd_lane = thread_index_in_simdgroup; +uint simd_id = simdgroup_index_in_threadgroup; +uint tg_lane = thread_position_in_threadgroup.x; +uint row_tile = threadgroup_position_in_grid.y; +uint route = threadgroup_position_in_grid.z; + +uint expert = uint(rhs_indices[route]); +uint out_dim = uint(codes_shape[1]); +uint codewords = uint(codes_shape[2]); +uint groups = uint(scales_shape[2]); +uint codewords_per_group = codewords / groups; +uint x_row = (route / uint(X_DIV)) * codewords * 8u; + +// (1024 + n) * 0.5 - 516 = (n - 8) / 2, exact in half. +threadgroup uint4 decoded[256]; +for (uint idx = tg_lane; idx < 256u; idx += threads_per_threadgroup.x) { + uint packed = codebook[idx]; + uint4 entry; + for (uint k = 0; k < 4u; ++k) { + half2 h = as_type(((packed >> (4u * k)) & 0x000F000Fu) | 0x64006400u); + entry[k] = as_type(fma(h, half2(0.5h), half2(-516.0h))); + } + decoded[idx] = entry; +} +threadgroup_barrier(mem_flags::mem_threadgroup); + +uint row0 = (row_tile * uint(SIMDGROUPS) + simd_id) * uint(ROWS_PER_SG); +float acc[ROWS_PER_SG]; +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + acc[r] = 0.0f; +} + +for (uint cw0 = simd_lane * uint(CHUNK); cw0 < codewords; cw0 += 32u * uint(CHUNK)) { + float2 xv[CHUNK * 4]; + float xs[CHUNK]; + for (uint i = 0; i < uint(CHUNK); ++i) { + float sum = 0.0f; + for (uint k = 0; k < 4u; ++k) { + uint xi = x_row + (cw0 + i) * 8u + 2u * k; + xv[i * 4u + k] = float2(float(x[xi]), float(x[xi + 1u])); + sum += xv[i * 4u + k].x + xv[i * 4u + k].y; + } + xs[i] = sum; + } + uint group = cw0 / codewords_per_group; + for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + if (row >= out_dim) { + break; + } + uint base = expert * out_dim + row; + uint4 words; + if (uint(CHUNK) == 8u) { + words = ((const device uint4*)(codes + base * codewords))[cw0 / 8u]; + } else { + words.xy = ((const device uint2*)(codes + base * codewords))[cw0 / 4u]; + } + float part = 0.0f; + for (uint i = 0; i < uint(CHUNK); ++i) { + uint codeword = (words[i >> 1u] >> ((i & 1u) * 16u)) & 0xFFFFu; + uint signs = codeword & 0xFFu; + uint parity = popcount(signs) & 1u; + uint effective = signs ^ parity; + uint4 entry = decoded[codeword >> 8u]; + // Sign bits k and k+4 move to bits 15 and 31 of pair k. + uint spread = (effective & 0x0Fu) | ((effective & 0xF0u) << 12u); + for (uint k = 0; k < 4u; ++k) { + half2 v = as_type(entry[k] ^ ((spread << (15u - k)) & 0x80008000u)); + part += dot(float2(v), xv[i * 4u + k]); + } + part += (parity != 0u ? -0.25f : 0.25f) * xs[i]; + } + acc[r] += part * float(scales[base * groups + group]); + } +} + +for (uint r = 0; r < uint(ROWS_PER_SG); ++r) { + uint row = row0 + r; + float total = simd_sum(acc[r]); + if (simd_lane == 0u && row < out_dim) { + out[route * out_dim + row] = OUT_T(total); + } +} diff --git a/src/ramp/kernels/e8p_m1.py b/src/ramp/kernels/e8p_m1.py new file mode 100644 index 0000000..43ca060 --- /dev/null +++ b/src/ramp/kernels/e8p_m1.py @@ -0,0 +1,183 @@ +"""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. + +On a Metal GPU without the NAX kernels (M1-M4), M=1 calls with 16-bit codes +use ``e8p_m1.metal`` by default, since ``benchmarks/bench_e8p_m1.py`` passed +on the M4 Max (2026-11-05). ``RAMP_E8P_M1=1`` forces it on, ``=0`` off: + +- :func:`e8p_m1_token_layout` for gate/up (one activation row, top-k + experts). +- :func:`e8p_m1_per_route` for down (one activation row *per route*). + +The kernel gives each simdgroup several rows, loads 4-8 codewords per lane at +once, and decodes the abs grid once per threadgroup. On the M4 Max that is +1.6-2.0x the generic M=1 kernel's E8P branch, which it replaces here; the +production 8-bit M=1 kernel is untouched. A loads-only stub reaches ~410 GB/s, +so decode ALU is still the limit. +""" + +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", + "should_use_e8p_m1", +] + +E8P_M1_ENV = "RAMP_E8P_M1" +#: Token rows the kernel serves: M=1 decode and MTP verify windows. Wider +#: calls (prefill) keep the sorted-route kernels, which share expert decode. +E8P_MAX_TOKENS = 32 +_KERNEL_DIR = Path(__file__).resolve().parent +_ON = {"1", "on", "true", "yes"} +_OFF = {"0", "off", "false", "no"} + + +def default_enabled() -> bool: + """On for a Metal GPU without NAX; NAX GPUs keep their tensor-unit kernels. + + Not cached: ``nax.is_available`` already is, and tests patch it. + """ + + import mlx.core as mx + + from ramp.kernels import nax + + return bool(mx.metal.is_available()) and not nax.is_available() + + +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 + if value in ("", "auto"): + return default_enabled() + raise ValueError(f"{E8P_M1_ENV} must be on/off/auto (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 1 <= tokens <= E8P_MAX_TOKENS + and not sorted_indices + and route_strategy in ("auto", "direct") + and e8p_m1_enabled() + ) + + +@cache +def _kernel(): + import mlx.core as mx + + return mx.fast.metal_kernel( + name="ramp_e8p_m1", + input_names=["x", "codes", "scales", "codebook", "rhs_indices"], + output_names=["out"], + source=(_KERNEL_DIR / "e8p_m1.metal").read_text(), + ) + + +def _chunk(codewords: int, groups: int) -> int: + """Codewords per wide load: a chunk must not straddle a row or scale group.""" + + for chunk in (8, 4): + if codewords % groups == 0 and (codewords // groups) % chunk == 0: + return chunk + raise ValueError(f"E8P M=1 needs codewords per scale group divisible by 4, got {codewords}/{groups}") + + +def _launch(x, codes, scales, codebook, rhs_indices, *, output_dims, rows_per_simdgroup, simdgroups, x_div): + import mlx.core as mx + + from ramp.kernels.gather_vqmm import _qmv_output_dtype + + if rows_per_simdgroup < 1 or simdgroups < 1 or simdgroups > 32: + raise ValueError("rows_per_simdgroup must be >= 1 and simdgroups in 1..32") + routes = int(rhs_indices.size) + rows_per_tile = rows_per_simdgroup * simdgroups + output_dtype = _qmv_output_dtype(x.dtype) + (out,) = _kernel()( + inputs=[mx.contiguous(x), codes, scales, codebook, rhs_indices.reshape(-1)], + template=[ + ("ROWS_PER_SG", rows_per_simdgroup), + ("SIMDGROUPS", simdgroups), + ("CHUNK", _chunk(int(codes.shape[2]), int(scales.shape[2]))), + ("X_DIV", int(x_div)), + ("OUT_T", output_dtype), + ], + grid=(32 * simdgroups, (output_dims + rows_per_tile - 1) // rows_per_tile, routes), + threadgroup=(32 * simdgroups, 1, 1), + output_shapes=[(1, routes, output_dims)], + output_dtypes=[output_dtype], + ) + return out + + +def e8p_m1_token_layout( + x, + codes, + scales, + codebook, + rhs_indices, + *, + output_dims: int, + rows_per_simdgroup: int = 4, + simdgroups: int = 1, +): + """Token ``t``'s row against its ``top_k`` experts: ``[tokens, top_k, out]``.""" + + if x.ndim != 2 or rhs_indices.ndim != 2 or rhs_indices.shape[0] != x.shape[0]: + raise ValueError("token-layout E8P needs x [tokens, in] and rhs_indices [tokens, top_k]") + tokens, top_k = int(rhs_indices.shape[0]), int(rhs_indices.shape[1]) + out = _launch( + x, codes, scales, codebook, rhs_indices, + output_dims=output_dims, + rows_per_simdgroup=rows_per_simdgroup, + simdgroups=simdgroups, + x_div=top_k, + ) + return out.reshape(tokens, top_k, output_dims) + + +def e8p_m1_per_route( + x, + codes, + scales, + codebook, + rhs_indices, + *, + output_dims: int, + rows_per_simdgroup: int = 4, + simdgroups: int = 1, +): + """Route ``r``'s activation row against expert ``rhs_indices[r]``: ``[routes, out]``.""" + + 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)}") + out = _launch( + x, codes, scales, codebook, rhs_indices, + output_dims=output_dims, + rows_per_simdgroup=rows_per_simdgroup, + simdgroups=simdgroups, + x_div=1, + ) + 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/models/glm4_moe_adapter.py b/src/ramp/models/glm4_moe_adapter.py index 2afdd13..510f3a5 100644 --- a/src/ramp/models/glm4_moe_adapter.py +++ b/src/ramp/models/glm4_moe_adapter.py @@ -9,6 +9,7 @@ import mlx.nn as nn from ramp.kernels import nax +from ramp.kernels.e8p_m1 import should_use_e8p_m1 from ramp.nn.switch_linear import HighPrecisionSwitchLinear, QuantizedVQSwitchLinear from ramp.ops.vq_switch import ( _expert_block_tile_descriptors, @@ -157,6 +158,13 @@ def num_experts(self) -> int: return self.gate_proj.num_experts def _can_use_shared_sorted_route_prefill(self, x: mx.array, indices: mx.array) -> bool: + if indices.ndim >= 2 and should_use_e8p_m1( + code_bits=self.gate_proj.code_bits, + tokens=indices.size // indices.shape[-1], + sorted_indices=False, + route_strategy=self.gate_proj.route_strategy, + ): + return False # few-token 16-bit E8P: the per-projection kernel is faster return ( indices.size >= 64 and x.ndim >= 2 diff --git a/src/ramp/nn/switch_linear.py b/src/ramp/nn/switch_linear.py index aadaf41..0b98441 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, + ): + # 16-bit codes at M=1 on GPUs without the NAX tensor-unit + # kernels (default there; RAMP_E8P_M1=0/1 overrides). + 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=flat_indices.shape[0] // route_shape[-1], + 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/fixtures/v41_reference/LICENSE b/tests/fixtures/v41_reference/LICENSE new file mode 100644 index 0000000..d62e3be --- /dev/null +++ b/tests/fixtures/v41_reference/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 DeepSeek + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/tests/fixtures/v41_reference/README.md b/tests/fixtures/v41_reference/README.md new file mode 100644 index 0000000..0cb4f7c --- /dev/null +++ b/tests/fixtures/v41_reference/README.md @@ -0,0 +1,3 @@ +`model.py` and `engram.py` are copied unchanged from `deepseek-ai/DeepSeek-V4.1-Flash-1017`, in the `inference/` directory at revision `1bae0e1088993b9b82de08f59de5be5a34e42170`. They are MIT-licensed; the license is in `LICENSE`. `model.py` has sha256 `e5decd76fe70fd56ef2f10124d245c6f64fb6a0003cc9068a1d3182a168b8d2d`. + +The CPU tests in `tests/test_runpod_v41_stream.py` run them at toy size. They swap in `scripts/runpod/v41_reap/torch_kernel.py` for the reference's tilelang `kernel.py`, and stub out the modules used only for vision. diff --git a/tests/fixtures/v41_reference/engram.py b/tests/fixtures/v41_reference/engram.py new file mode 100644 index 0000000..a36bc6c --- /dev/null +++ b/tests/fixtures/v41_reference/engram.py @@ -0,0 +1,184 @@ +from dataclasses import dataclass + +import numpy as np +import torch +from sympy import isprime +from torch import nn + + +def find_next_prime(start: int, seen_primes: set[int]) -> int: + """The smallest prime above `start` that has not been handed out yet.""" + candidate = start + 1 + while not isprime(candidate) or candidate in seen_primes: + candidate += 1 + return candidate + + +def build_compressed_token_map(tokenizer) -> tuple[list[int], int]: + """Map every token id onto a smaller id space where tokens that normalize alike collapse together. + + N-grams are hashed over these compressed ids, so " The", "the" and "THE" all hash the same way. + Returns the lookup plus the size of the compressed vocab -- and that size matters beyond bounds + checking, because every hash multiplier is derived from it. + """ + from tokenizers import Regex, normalizers + + # a private-use char, so a token that is exactly one space survives Strip() instead of + # collapsing to the empty string and merging with unrelated tokens + sentinel = "\ue000" + normalizer = normalizers.Sequence( + [ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ] + ) + + # the raw Rust tokenizer, matching what training decodes with (no clean_up_tokenization_spaces) + backend = tokenizer.backend_tokenizer + key_to_new: dict[str, int] = {} + lookup = [0] * len(tokenizer) + for token_id in range(len(tokenizer)): + text = backend.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + # a partial UTF-8 byte token: nothing to normalize, so key it by its raw form + key = backend.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + + return lookup, len(key_to_new) + + +def compute_hash_multipliers( + layer_ids: tuple[int, ...], max_ngram_size: int, tokenizer_vocab_size: int +) -> torch.Tensor: + """One multiplier per (layer, lookback), from a per-layer RNG so layers hash differently. + + Kept odd, and bounded so that `token_id * multiplier` cannot overflow int64. + """ + max_long = np.iinfo(np.int64).max + multiplier_bound = max(1, (max_long // tokenizer_vocab_size) // 2) + rows = [] + for layer_id in layer_ids: + generator = np.random.default_rng(10007 * layer_id) + values = generator.integers( + low=0, + high=multiplier_bound, + size=(max_ngram_size,), + dtype=np.int64, + ) + rows.append(torch.tensor(values * 2 + 1)) + return torch.stack(rows) + + +@dataclass(frozen=True) +class EngramLayout: + """Bucket layout of the n-gram hash tables. + + A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max_ngram_size-gram), each split + over `n_heads` heads. Every (n-gram size, head) pair owns its own prime-sized bucket range in the + layer's table; the primes are drawn in order and never reused, which keeps the ranges disjoint. + """ + + max_ngram_size: int + layer_ids: tuple[int, ...] + num_embeddings: tuple[int, ...] # table rows, per engram layer + primes: tuple[tuple[tuple[int, ...], ...], ...] # [layer][n-gram size][head] bucket modulus + n_heads: int + head_dim: int + + @classmethod + def from_args(cls, args) -> "EngramLayout | None": + layer_ids = tuple(args.engram_layer_ids) + if not layer_ids: + return None + max_ngram_size, n_heads = args.engram_max_ngram_size, args.engram_n_heads + primes, seen = [], set() + for _ in layer_ids: + per_ngram = [] + for _ in range(max_ngram_size - 1): + sizes, current = [], args.engram_vocab_size - 1 + for _ in range(n_heads): + current = find_next_prime(current, seen) + seen.add(current) + sizes.append(current) + per_ngram.append(tuple(sizes)) + primes.append(tuple(per_ngram)) + return cls( + max_ngram_size=max_ngram_size, + layer_ids=layer_ids, + num_embeddings=tuple(args.engram_num_embeddings), + primes=tuple(primes), + n_heads=n_heads, + head_dim=args.engram_head_dim, + ) + + +class NgramHashState(nn.Module): + """Maps each position to the hash ids of the n-grams ending there. + + Ids go through the compressed table, then each position is hashed with the `max_ngram_size - 1` + tokens before it. Look-back stops at the start of the sequence and at any dead token (an image + span, cached as DEAD), so an n-gram never spans one. The cache carries all of this across the + prefill/decode split. + """ + + DEAD = -1 + + def __init__(self, args, layout: EngramLayout, tokenizer): + super().__init__() + self.layout = layout + # every hash multiplier derives from the compressed vocab size, so a mismatch there would + # silently rehash the whole table + token_map, vocab_size = build_compressed_token_map(tokenizer) + assert vocab_size == args.engram_compressed_vocab_size, (vocab_size, args.engram_compressed_vocab_size) + self.pad_id = token_map[args.engram_pad_id] + flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout.primes] + offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat] + multipliers = compute_hash_multipliers(layout.layer_ids, layout.max_ngram_size, vocab_size) + self.register_buffer("primes", torch.tensor(layout.primes), persistent=False) + self.register_buffer("offsets", torch.tensor(np.array(offsets)), persistent=False) + self.register_buffer("multipliers", multipliers, persistent=False) + self.register_buffer("token_map", torch.tensor(token_map), persistent=False) + self.register_buffer( + "cache", torch.empty(args.max_batch_size, args.max_seq_len, dtype=torch.int64), persistent=False + ) + + @torch.inference_mode() + def forward(self, input_ids: torch.Tensor, start_pos: int, token_mask: torch.Tensor | None = None) -> torch.Tensor: + """token_mask: [B, L], False for tokens that take no part in an n-gram (image spans). + Returns the hash ids, shaped [B, L, n_engram_layers, n_hash_cols].""" + batch, seqlen = input_ids.shape + compressed = self.token_map[input_ids] + if token_mask is not None: + compressed = torch.where(token_mask, compressed, self.DEAD) + self.cache[:batch, start_pos : start_pos + seqlen] = compressed + + positions = torch.arange(start_pos, start_pos + seqlen, device=input_ids.device).expand(batch, seqlen) + tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool) + for shift in range(self.layout.max_ngram_size): + source = self.cache[:batch].gather(1, (positions - shift).clamp_min(0)) + blocked = blocked | (positions < shift) | (source == self.DEAD) + tokens.append(torch.where(blocked, self.pad_id, source)) + tokens = torch.stack(tokens, dim=-1) # [B, L, max_ngram_size] + + # XOR the multiplied ids together one lookback at a time, so the running value after step i + # is the hash of the (i+1)-gram; each lands in its own prime-sized bucket range + products = tokens.unsqueeze(2) * self.multipliers # [B, L, n_engram_layers, max_ngram_size] + rolling, hashes = products[..., 0], [] + for i in range(1, self.layout.max_ngram_size): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1]) + return torch.cat(hashes, dim=-1) + self.offsets diff --git a/tests/fixtures/v41_reference/model.py b/tests/fixtures/v41_reference/model.py new file mode 100644 index 0000000..46cc831 --- /dev/null +++ b/tests/fixtures/v41_reference/model.py @@ -0,0 +1,1312 @@ +import math +from contextlib import contextmanager +from dataclasses import dataclass +from functools import lru_cache +from typing import Literal + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from engram import EngramLayout, NgramHashState +from image_processor import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START +from kernel import ( + act_quant, + fp4_act_quant, + fp4_gemm, + fp8_gemm, + hc_split_sinkhorn, + sparse_attn, +) +from torch import nn +from vision import Aligner, ViT + +# Set once by Transformer.__init__; one model per process, so layers just read them. +world_size = 1 +rank = 0 +default_dtype = torch.float8_e4m3fn # storage dtype for Linear weights, from ModelArgs.dtype +fp8_block_size = 32 # one fp8 scale per 32x32 weight block / 32 activations +fp4_block_size = 32 # one fp4 scale per 32 elements along K +scale_fmt = "ue8m0" +scale_dtype = torch.float8_e8m0fnu + + +@contextmanager +def set_dtype(dtype): + """Temporarily override torch's default dtype, restoring it even if the body raises.""" + prev = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(prev) + + +@dataclass +class ModelArgs: + """Field names are exactly the config JSON keys. The defaults are a small model that + `python model.py` can run, not the released shapes -- though the scale-independent + values (norm_eps, score_func, hc_*, engram_*) do match it.""" + + # runtime limits rather than model shape: they size the KV caches + max_batch_size: int = 4 + max_seq_len: int = 4096 + temperature: float = 1 + dtype: Literal["bf16", "fp8"] = "fp8" + expert_dtype: Literal["fp4"] | None = "fp4" + vocab_size: int = 129280 + dim: int = 1024 + moe_inter_dim: int = 1024 + n_layers: int = 5 + n_mtp_layers: int = 1 # extra draft layers appended after the backbone, indices n_layers.. + n_heads: int = 16 + # moe + n_routed_experts: int = 8 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + gate_temp: float = 1.0 + norm_topk_prob: bool = True + route_scale: float = 1.0 + swiglu_limit: float = 0.0 + # attention: latent q/kv projections, plus a LoRA-factorised output projection over o_groups + q_lora_rank: int = 256 + head_dim: int = 128 + rope_head_dim: int = 32 + norm_eps: float = 1e-20 + o_groups: int = 8 + o_lora_rank: int = 256 + # sparse attention: every layer attends over a sliding window, and may add compressed KV on top + window_size: int = 128 + # one entry per layer, MTP layers included: 0 = sliding window only, r = KV compressed r-to-1 + compress_ratios: tuple[int, ...] = (0, 2, 2, 1, 1, 0) + # layers sharing a ratio also share one compressed KV and one indexer, produced by the first + kv_source_layers: tuple[int, ...] = (1, 3) + index_source_layers: tuple[int, ...] = (1, 3) + # rope, with YaRN extrapolation when original_seq_len > 0. Compressed KV rotates at its own + # theta because one latent stands for compress_ratio tokens, so its positions are further apart. + compress_rope_theta: float = 40000.0 + original_seq_len: int = 0 + rope_theta: float = 10000.0 + rope_factor: float = 40 + beta_fast: int = 32 + beta_slow: int = 1 + # the indexer: a small extra attention that scores compressed positions, so each query can keep + # just `index_topk` of them. Names match DeepSeek-V3.2-Exp, where this mechanism first appeared. + index_n_heads: int = 16 + index_head_dim: int = 64 + index_topk: int = 64 + # candidate pre-filtering: candidate_source_layer < 0 turns it off and the other two are unused + candidate_source_layer: int = -1 + candidate_topk_blocks: int = 0 + candidate_block_size: int = 0 + # hyper-connections: the residual stream is carried as hc_mult parallel copies + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + # engram: n-gram hash lookups added into the residual stream at a few layers + engram_layer_ids: tuple[int, ...] = () + engram_num_embeddings: tuple[int, ...] = () # unpadded table rows; each rank allocates ceil(rows / world_size) + engram_max_ngram_size: int = 1 + engram_vocab_size: int = 0 # bucket size each (n-gram size, head) starts searching primes from + engram_n_heads: int = 0 + engram_head_dim: int = 0 + engram_pad_id: int = 2 # token that fills n-gram slots with no history; matches training + # size of the compressed tokenizer vocab; every hash multiplier is derived from it + engram_compressed_vocab_size: int = 0 + # vision (VL); vision_n_layers == 0 disables the vision path + vision_n_layers: int = 0 + vision_dim: int = 1024 + vision_n_heads: int = 16 + vision_inter_dim: int = 2816 + vision_patch_size: int = 14 + vision_rope_theta: float = 10000.0 + vision_downsample_ratio: int = 3 + vision_max_n_token: int = 1024 + vision_min_pixels: int = 544 * 544 + vision_max_wh_ratio: int | None = None + # raw id of <｜deepseek_image｜>; every position of an image span carries this id in input_ids + image_token_id: int = 129264 + # dspark draft head. Only the forward pass is implemented here -- nothing calls forward_spec, + # so these are read but the speculative-decoding loop itself is out of scope for this repo. + dspark_block_size: int = 0 + dspark_noise_token_id: int = 0 + dspark_target_layer_ids: tuple[int, ...] = () + dspark_markov_rank: int = 256 + dspark_n_routed_experts: int = 0 + dspark_n_activated_experts: int = 0 + + @property + def vision_enabled(self) -> bool: + return self.vision_n_layers > 0 + + def get_moe_config(self, layer_id: int) -> tuple[int, int]: + """Return the routed/activated expert counts for a given layer.""" + if layer_id < self.n_layers: + return self.n_routed_experts, self.n_activated_experts + return ( + self.dspark_n_routed_experts or self.n_routed_experts, + self.dspark_n_activated_experts or self.n_activated_experts, + ) + + +class ParallelEmbedding(nn.Module): + """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows. + Out-of-range indices are zero-masked before all_reduce to combine partial embeddings.""" + + def __init__(self, vocab_size: int, dim: int): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + assert vocab_size % world_size == 0, ( + f"Vocabulary size must be divisible by world size (world_size={world_size})" + ) + self.part_vocab_size = vocab_size // world_size + self.vocab_start_idx = rank * self.part_vocab_size + self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if world_size > 1: + # ids off this rank read row 0 then get zeroed, so the all_reduce sums one real row + mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx) + x = x - self.vocab_start_idx + x[mask] = 0 + y = F.embedding(x, self.weight) + if world_size > 1: + y[mask] = 0 + dist.all_reduce(y) + return y + + +def linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + """Pick a GEMM from the weight dtype. Quantized weights need a quantized activation, and both + fp4 and fp8 weights take an fp8 one -- for fp4 the kernel handles the mixed precision.""" + assert bias is None + + if weight.dtype == torch.float4_e2m1fn_x2: + x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype) + return fp4_gemm( + x, + s, + weight, + weight.scale, + scale_dtype, + act_block_size=fp8_block_size, + ) + elif weight.dtype == torch.float8_e4m3fn: + x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype) + return fp8_gemm( + x, + s, + weight, + weight.scale, + scale_dtype, + block_size=fp8_block_size, + ) + else: + return F.linear(x, weight) + + +class Linear(nn.Module): + """bf16, fp8 or fp4 weights. Quantized ones get a `scale`, also attached to `.weight` so that + `linear()` can reach it.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + super().__init__() + self.in_features = in_features + self.out_features = out_features + dtype = dtype or default_dtype + if dtype == torch.float4_e2m1fn_x2: + # two values per byte: [out, in] logically, [out, in//2] stored + self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2)) + self.weight.scale = self.scale = nn.Parameter( + torch.empty(out_features, in_features // fp4_block_size, dtype=torch.float8_e8m0fnu) + ) + elif dtype == torch.float8_e4m3fn: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.weight.scale = self.scale = nn.Parameter( + torch.empty( + (out_features + fp8_block_size - 1) // fp8_block_size, + (in_features + fp8_block_size - 1) // fp8_block_size, + dtype=torch.float8_e8m0fnu, + ) + ) + else: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.register_parameter("scale", None) + if bias: + self.bias = nn.Parameter(torch.empty(out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class ColumnParallelLinear(Linear): + """Splits the output dim across ranks; each rank's slice of the output is already complete.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + assert out_features % world_size == 0, ( + f"Output features must be divisible by world size (world_size={world_size})" + ) + self.part_out_features = out_features // world_size + super().__init__(in_features, self.part_out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class RowParallelLinear(Linear): + """Splits the reduction dim, so each rank holds a partial sum: hence the fp32 all_reduce, with + the bias added only after it.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + assert in_features % world_size == 0, ( + f"Input features must be divisible by world size (world_size={world_size})" + ) + self.part_in_features = in_features // world_size + super().__init__(self.part_in_features, out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = linear(x, self.weight, None) + if world_size > 1: + y = y.float() + dist.all_reduce(y) + if self.bias is not None: + y += self.bias + return y.type_as(x) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +class ParallelEngramEmbedding(nn.Module): + """The n-gram hash table, sharded over its rows. Stays fp8: rows are dequantized on lookup.""" + + def __init__(self, num_embeddings: int, dim: int): + super().__init__() + self.num_embeddings = num_embeddings + self.dim = dim + self.part_num_embeddings = (num_embeddings + world_size - 1) // world_size + self.vocab_start_idx = rank * self.part_num_embeddings + self.vocab_end_idx = self.vocab_start_idx + self.part_num_embeddings + + self.block_size = fp8_block_size + # the table stays fp8 as stored: rows are dequantized with `scale` on lookup + self.weight = nn.Parameter(torch.empty(self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn)) + self.scale = nn.Parameter(torch.empty(self.part_num_embeddings, dim // self.block_size, dtype=scale_dtype)) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + mask = (indices < self.vocab_start_idx) | (indices >= self.vocab_end_idx) + local_indices = indices - self.vocab_start_idx + local_indices = local_indices.masked_fill(mask, 0) + + values = F.embedding(local_indices, self.weight) + scales = F.embedding(local_indices, self.scale) + values = values.float().unflatten(-1, (-1, self.block_size)) * scales.float().unsqueeze(-1) + values = values.flatten(-2).to(torch.bfloat16) + values = values.masked_fill(mask.unsqueeze(-1), 0) + + if world_size > 1: + dist.all_reduce(values) + return values + + +class Engram(nn.Module): + """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream. + + The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per hc copy plus a shared + value. The gate is a normalized dot product of stream against key. + """ + + def __init__(self, args: ModelArgs, layer_id: int, layout: EngramLayout): + super().__init__() + self.layer_id = layer_id + self.layer_hash_index = layout.layer_ids.index(layer_id) + self.dim = args.dim + self.hc_mult = args.hc_mult + self.clamp_value = 1e-6 + + self.embed = ParallelEngramEmbedding(layout.num_embeddings[self.layer_hash_index], layout.head_dim) + n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads + self.wkv = Linear(n_hash_cols * layout.head_dim, args.dim * (args.hc_mult + 1)) + self.eps = args.norm_eps + self.q_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim)) + self.k_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim)) + + def forward(self, x: torch.Tensor, hash_ids: torch.Tensor, token_mask: torch.Tensor | None = None) -> torch.Tensor: + """x: [B, L, hc_mult, dim]; hash_ids: [B, L, n_hash_cols]; token_mask: [B, L], False shuts + the gate so those positions pass through untouched.""" + kv = self.wkv(self.embed(hash_ids).flatten(-2)) + key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1) + key = key.float().unflatten(-1, (self.hc_mult, self.dim)) + weight = self.q_weight.float() * self.k_weight.float() # only ever used as a product + h, eps = x.float(), self.eps + # normalized per (token, hc copy) over `dim`, NOT jointly over the copies + rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(key.square().mean(-1) + eps) + dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5 + # signed sqrt before the sigmoid, matching the training kernel + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot)) + if token_mask is not None: + gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0) + return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype) + + +@lru_cache(2) +def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor: + """Rotary frequencies as complex exponentials, one row per position. + + With original_seq_len > 0 this applies YaRN: dimensions whose wavelength already fits inside the + training context keep their frequency, those far beyond it are divided by `factor`, and the + `beta_fast`..`beta_slow` band in between is faded across with a linear ramp. + """ + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + # the dim whose wavelength completes `rotations` turns over the training context + def corrected_dim(rotations): + return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base)) + + low = max(math.floor(corrected_dim(beta_fast)), 0) + high = min(math.ceil(corrected_dim(beta_slow)), dim - 1) + ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1) + smooth = 1 - ramp + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + freqs = torch.outer(torch.arange(seqlen), freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Rotate `x` in place, taking adjacent element pairs as complex numbers. Accepts [b, s, d] and + [b, s, h, d]; `inverse` conjugates the rotation, which is how the attention output gets the + query's rotation removed again so the cache can stay in one shared rotated form.""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +@lru_cache(1) +def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int): + """Which sliding-window cache slots each query attends to; -1 marks a slot holding nothing. + + The cache is a ring of `window_size` slots. Prefill needs one row per query, each seeing its own + causal window. A decode step has a single query that sees the whole ring, listed oldest first. + Order within a row does not matter to `sparse_attn`, which handles every slot independently. + """ + if start_pos == 0: + end = torch.arange(seqlen).unsqueeze(1) + idxs = (end - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size)) + idxs = torch.where(idxs > end, -1, idxs) # before the sequence started + else: + oldest = start_pos % window_size + 1 + idxs = torch.cat([torch.arange(oldest, window_size), torch.arange(oldest)]) + idxs = torch.where(idxs > start_pos, -1, idxs) # ring still filling + # sparse_attn needs real [b, m, topk] int32 memory, hence the materializing expand + return idxs.int().unsqueeze(0).expand(bsz, -1, -1).contiguous() + + +class Compressor(nn.Module): + """Pools `compress_ratio` consecutive tokens into one KV latent with a learned softmax gate. + + Returns the latent before RoPE, or None while a group is still filling up -- so during decode it + only yields every `compress_ratio` steps, holding the partial group in `kv_state`/`score_state`. + Pre-RoPE is deliberate: the indexer needs the unrotated form, so Attention rotates afterwards. + """ + + def __init__(self, args: ModelArgs, layer_id: int): + super().__init__() + compress_ratio = args.compress_ratios[layer_id] + head_dim = args.head_dim + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.norm = RMSNorm(head_dim, args.norm_eps) + # ratio 1 is a plain projection, so it stays in the checkpoint's bf16; the softmax pooling + # above ratio 1 runs in fp32, so those weights are promoted to fp32 to match + self.wkv = Linear(args.dim, head_dim, dtype=torch.float32 if compress_ratio > 1 else torch.bfloat16) + if compress_ratio > 1: + self.wgate = Linear(args.dim, head_dim, dtype=torch.float32) + # tail of an incomplete group, carried across decode steps + self.kv_state: torch.Tensor + self.score_state: torch.Tensor + state_shape = (args.max_batch_size, compress_ratio, head_dim) + self.register_buffer("kv_state", torch.zeros(state_shape, dtype=torch.float32), persistent=False) + self.register_buffer( + "score_state", torch.full(state_shape, -torch.inf, dtype=torch.float32), persistent=False + ) + + def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor | None: + bsz, seqlen, _ = x.size() + ratio, dtype = self.compress_ratio, x.dtype + if ratio == 1: # one token per group: nothing to pool, so no gate and no fp32 + return self.norm(self.wkv(x)) + + x = x.float() + kv, score = self.wkv(x), self.wgate(x) + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + if remainder: # trailing partial group waits in the state + kv, self.kv_state[:bsz, :remainder] = kv.split([cutoff, remainder], dim=1) + score, self.score_state[:bsz, :remainder] = score.split([cutoff, remainder], dim=1) + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: # one token per step: fill a slot, and pool only when the group just completed + should_compress = (start_pos + 1) % ratio == 0 + slot = start_pos % ratio + self.kv_state[:bsz, slot] = kv.squeeze(1) + self.score_state[:bsz, slot] = score.squeeze(1) + if should_compress: + kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True) + if not should_compress: + return None + return self.norm(kv.to(dtype)) + + +class Indexer(torch.nn.Module): + """Keeps the `index_topk` best compressed positions per query. + + A small side attention: fp4 query heads against one shared key per compressed position, scores + rectified then combined by `weights_proj`. With a candidate source this is the second of two + levels; `select_candidate_blocks` is the first. + """ + + def __init__(self, args: ModelArgs, layer_id: int): + super().__init__() + # the index keys are derived from the compressor's latent, so only a layer that compresses + # its own KV can produce them; every other indexer reads them from that layer's cache + self.owns_k = layer_id in args.kv_source_layers + self.compress_ratio = args.compress_ratios[layer_id] + self.is_candidate_source = layer_id == args.candidate_source_layer + self.uses_candidates = 0 <= args.candidate_source_layer < layer_id + self.candidate_topk_blocks = args.candidate_topk_blocks + self.candidate_block_size = args.candidate_block_size + self.dim = args.dim + self.n_heads = args.index_n_heads + self.n_local_heads = args.index_n_heads // world_size + self.index_head_dim = args.index_head_dim + self.rope_head_dim = args.rope_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.softmax_scale = self.index_head_dim**-0.5 + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.index_head_dim) + self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16) + self.freqs_cis: torch.Tensor | None = None + if self.owns_k: + self.wk = Linear(args.head_dim, self.index_head_dim, dtype=torch.bfloat16) + self.k_norm = RMSNorm(self.index_head_dim, args.norm_eps) + self.k_cache: torch.Tensor + self.register_buffer( + "k_cache", + torch.zeros(args.max_batch_size, args.max_seq_len // self.compress_ratio, args.index_head_dim), + persistent=False, + ) + + def forward(self, x: torch.Tensor, qr: torch.Tensor, latent: torch.Tensor, start_pos: int, offset: int): + """`latent` is this layer's RoPE-free compressed latent, None when this layer does not + compress or when its current group is still incomplete. An index-key owner turns it into + index keys here, which has to happen before Attention overwrites that same storage with + the RoPE'd, quantized values.""" + assert self.freqs_cis is not None + bsz, seqlen, _ = x.size() + ratio, rd, end_pos = self.compress_ratio, self.rope_head_dim, start_pos + seqlen + + # Append index keys when a compression group completes. + if self.owns_k and latent is not None: + # a latent stands for the first token of its group, so group j takes position j * ratio + freqs = ( + self.freqs_cis[: seqlen - seqlen % ratio : ratio] + if start_pos == 0 + else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0) + ) + k = self.k_norm(self.wk(latent)) + apply_rotary_emb(k[..., -rd:], freqs) + fp4_act_quant(k, fp4_block_size, True) + self.k_cache[:bsz, start_pos // ratio : start_pos // ratio + k.size(1)] = k + if self.owns_k: + shared_attn.index_k = self.k_cache + + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim)) + apply_rotary_emb(q[..., -rd:], self.freqs_cis[start_pos:end_pos]) + fp4_act_quant(q, fp4_block_size, True) + + index_k = shared_attn.index_k[:bsz, : end_pos // ratio] + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5) + index_score = torch.einsum("bshd,btd->bsht", q, index_k) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + if world_size > 1: + dist.all_reduce(index_score) + + # how many compressed positions each query can see: a block becomes visible once the query + # has passed its last token. One query per decode step, so there it is just a number. + if start_pos == 0: + compress_lens = (torch.arange(1, seqlen + 1, device=x.device) // ratio).unsqueeze(-1) + index_score.masked_fill_(torch.arange(seqlen // ratio, device=x.device) >= compress_lens, -torch.inf) + else: + compress_lens = end_pos // ratio + + if self.is_candidate_source: + shared_attn.candidates = select_candidate_blocks( + index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size + ) + elif self.uses_candidates: + # level two: score with our own weights, but only inside the source's candidate blocks + index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf) + + # top-k by score, re-sorted into position order; unreachable -> -1, rest shifted by offset + topk = min(self.index_topk, end_pos // ratio) + idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + return torch.where(idxs < compress_lens, idxs + offset, -1).int() + + +def select_candidate_blocks( + logits: torch.Tensor, + compress_lens: torch.Tensor | int, + topk_blocks: int, + block_size: int, +) -> torch.Tensor: + """Level one of the two-level top-k: keep the `topk_blocks` highest-scoring blocks per query. + + `logits` is [..., n_positions] with positions the query cannot reach already at -inf, which is + what makes a block score of -inf mean "not reachable yet". `compress_lens` is a plain int during + decode, or broadcasts against logits' leading dims during prefill. Returns a bool mask shaped + like `logits`, so the layers consuming it just mask and never think about blocks again. + """ + width = logits.size(-1) + # score each block by its best position; -inf pads the last one out to block_size + scores = F.pad(logits, (0, -width % block_size), value=-torch.inf) + scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1) + num_blocks = scores.size(-1) + + # the block with this query's newest position is only partly filled, so pin it in: it holds the + # most recent tokens but could otherwise be outscored by an older, full block + last = (compress_lens - 1) // block_size + scores = scores.masked_fill(torch.arange(num_blocks, device=logits.device) == last, torch.inf) + + top = scores.topk(min(topk_blocks, num_blocks), dim=-1) + # fewer reachable blocks than topk_blocks means leftover picks came back -inf: drop them + keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, top.indices, top.values > -torch.inf) + return keep.repeat_interleave(block_size, dim=-1)[..., :width] + + +class Attention(nn.Module): + """Latent attention over two KV sources at once, concatenated into one `sparse_attn` call: a + sliding window of raw KV, plus -- when compress_ratio > 0 -- `index_topk` compressed positions + reaching further back. Q and the output projection are both low-rank, the latter grouped. + + compress_ratio > 0 does not mean the layer compresses its own KV: only kv_source_layers do, + the rest read that same cache. + """ + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.dim = args.dim + self.n_heads = args.n_heads + self.n_local_heads = args.n_heads // world_size + self.q_lora_rank = args.q_lora_rank + self.o_lora_rank = args.o_lora_rank + self.head_dim = args.head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = args.head_dim - args.rope_head_dim + self.n_groups = args.o_groups + self.n_local_groups = self.n_groups // world_size + self.window_size = args.window_size + self.compress_ratio = args.compress_ratios[layer_id] + self.eps = args.norm_eps + + self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32)) + self.wq_a = Linear(self.dim, self.q_lora_rank) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.wkv = Linear(self.dim, self.head_dim) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * args.o_lora_rank, + dtype=torch.bfloat16, + ) + self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim) + self.softmax_scale = self.head_dim**-0.5 + + is_backbone = layer_id < args.n_layers + self.is_kv_source = is_backbone and layer_id in args.kv_source_layers + self.is_index_source = is_backbone and layer_id in args.index_source_layers + self.compressor: Compressor | None = None + self.indexer: Indexer | None = None + if self.is_kv_source: + self.compressor = Compressor(args, layer_id) + if self.is_index_source: + self.indexer = Indexer(args, layer_id) + + self.window_kv_cache: torch.Tensor + self.register_buffer( + "window_kv_cache", + torch.zeros(args.max_batch_size, args.window_size, self.head_dim), + persistent=False, + ) + if self.is_kv_source: + self.compress_kv_cache: torch.Tensor + self.register_buffer( + "compress_kv_cache", + torch.zeros( + args.max_batch_size, + args.max_seq_len // self.compress_ratio, + self.head_dim, + ), + persistent=False, + ) + if self.compress_ratio: + original_seq_len, rope_theta = ( + args.original_seq_len, + args.compress_rope_theta, + ) + else: + # disable YaRN and use base rope_theta in pure sliding-window attention + original_seq_len, rope_theta = 0, args.rope_theta + freqs_cis = precompute_freqs_cis( + self.rope_head_dim, + args.max_seq_len, + original_seq_len, + rope_theta, + args.rope_factor, + args.beta_fast, + args.beta_slow, + ) + self.freqs_cis: torch.Tensor + self.register_buffer("freqs_cis", freqs_cis, persistent=False) + + def _window_kv(self, x, freqs_cis, start_pos): + """This layer's sliding-window K and the window positions every query may attend to. The K + stays fp8, quantized over the whole post-RoPE vector, RoPE tail included.""" + bsz, seqlen, _ = x.size() + win = self.window_size + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + if start_pos == 0: # prefill: attend over this chunk, seeding the ring buffer for decode + if seqlen <= win: + self.window_kv_cache[:bsz, :seqlen] = kv + else: + cutoff = seqlen % win + self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = kv[:, -win:].split( + [win - cutoff, cutoff], dim=1 + ) + window_kv = kv + else: # decode: one token into the ring buffer, attend over the whole window + self.window_kv_cache[:bsz, start_pos % win] = kv.squeeze(1) + window_kv = self.window_kv_cache[:bsz] + return window_kv, get_window_topk_idxs(win, bsz, seqlen, start_pos) + + def _compress_topk_idxs(self, x, qr, latent, start_pos, offset, compress_len): + """Which compressed positions each query attends to. Index sources run their own indexer; + the layers in between reuse the result their source published.""" + if not self.is_index_source: + return shared_attn.topk_idxs + + bsz, seqlen, _ = x.size() + if compress_len == 0: + idxs = torch.empty(bsz, seqlen, 0, dtype=torch.int32, device=x.device) + else: + assert self.indexer is not None + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + idxs = self.indexer(x, qr, latent, start_pos, offset) + shared_attn.topk_idxs = idxs + return idxs + + def _compress_kv(self, x, qr, start_pos, offset): + """The shared compressed KV and the compressed positions every query may attend to. This + layer compresses its own KV only when it is a source; otherwise it just reads the cache.""" + bsz, seqlen, _ = x.size() + ratio = self.compress_ratio + compress_len = (start_pos + seqlen) // ratio + latent = None + if self.is_kv_source: + latent = self.compressor(x, start_pos) + shared_attn.compress_kv = self.compress_kv_cache + # the indexer needs the latent before RoPE, so it runs before the cache is written + idxs = self._compress_topk_idxs(x, qr, latent, start_pos, offset, compress_len) + if latent is not None: + # a latent stands for the first token of its group, so group j takes position j * ratio + freqs = ( + self.freqs_cis[: seqlen - seqlen % ratio : ratio] + if start_pos == 0 + else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0) + ) + apply_rotary_emb(latent[..., -self.rope_head_dim :], freqs) + # Compressed KV uses groups of 16 with E4M3 scales; the indexer uses 32 with E8M0. + fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn) + self.compress_kv_cache[:bsz, start_pos // ratio : start_pos // ratio + latent.size(1)] = latent + # read after the write, so this does not depend on the slice aliasing the cache + return shared_attn.compress_kv[:bsz, :compress_len], idxs + + def forward(self, x: torch.Tensor, start_pos: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + rd = self.rope_head_dim + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + kv, topk_idxs = self._window_kv(x, freqs_cis, start_pos) + if self.compress_ratio: + compress_kv, compress_idxs = self._compress_kv(x, qr, start_pos, kv.size(1)) + kv = torch.cat([kv, compress_kv], dim=1) + topk_idxs = torch.cat([topk_idxs, compress_idxs], dim=-1) + + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + # wo_a is block-diagonal over groups (each projects only its own heads), hence einsum not + # Linear. convert.py dequantizes it to bf16; quantize/dequantize o before einsum. + o = o.view(bsz, seqlen, self.n_local_groups, -1) + act_quant(o, fp8_block_size, scale_fmt, scale_dtype, True) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + +class Gate(nn.Module): + """MoE gating. The correction bias steers expert selection only; the routing weights come from the + unbiased scores. Image-span tokens use a separate bias (training `noaux_tc_for_vl`).""" + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + n_routed_experts, n_activated_experts = args.get_moe_config(layer_id) + self.dim = args.dim + self.topk = n_activated_experts + self.score_func = args.score_func + self.gate_temp = args.gate_temp + self.norm_topk_prob = args.norm_topk_prob + self.route_scale = args.route_scale + self.weight = nn.Parameter(torch.empty(n_routed_experts, args.dim)) + self.bias = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) + self.bias_vl = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) if args.vision_enabled else None + + def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: + """x: [n, dim]; image_mask: [n] bool, True for tokens inside an image span.""" + scores = linear(x.float(), self.weight.float()) / self.gate_temp + if self.score_func == "softmax": + scores = scores.softmax(dim=-1) + elif self.score_func == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + bias = self.bias + if image_mask is not None and self.bias_vl is not None: + bias = torch.where(image_mask.unsqueeze(-1), self.bias_vl, bias) + # the bias picks experts but does not scale them: weights come from the raw scores + indices = (scores + bias).topk(self.topk, dim=-1)[1] + weights = scores.gather(1, indices) + if self.norm_topk_prob and self.topk > 1: + weights /= weights.sum(dim=-1, keepdim=True) + 1e-20 # not norm_eps, matches training + weights *= self.route_scale + return weights, indices + + +class Expert(nn.Module): + """One SwiGLU FFN. The clamps come straight from training, where they keep fp8/fp4 activations in + range: the up branch is clamped on both sides, the gate branch only from above.""" + + def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0.0): + super().__init__() + self.w1 = Linear(dim, inter_dim, dtype=dtype) + self.w2 = Linear(inter_dim, dim, dtype=dtype) + self.w3 = Linear(dim, inter_dim, dtype=dtype) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor: + dtype = x.dtype + gate = self.w1(x).float() + up = self.w3(x).float() + if self.swiglu_limit > 0: + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + gate = torch.clamp(gate, max=self.swiglu_limit) + x = F.silu(gate) * up + if weights is not None: + x = weights * x + return self.w2(x.to(dtype)) + + +class MoE(nn.Module): + """Top-k routed experts plus one shared expert every token goes through. Experts are split + across ranks, so `self.experts` is None for those another rank owns.""" + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + n_routed_experts, n_activated_experts = args.get_moe_config(layer_id) + self.layer_id = layer_id + self.dim = args.dim + assert n_routed_experts % world_size == 0, ( + f"Number of experts must be divisible by world size (world_size={world_size})" + ) + self.n_routed_experts = n_routed_experts + self.n_local_experts = n_routed_experts // world_size + self.n_activated_experts = n_activated_experts + self.experts_start_idx = rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.gate = Gate(layer_id, args) + expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None + self.experts = nn.ModuleList( + [ + Expert( + args.dim, + args.moe_inter_dim, + dtype=expert_dtype, + swiglu_limit=args.swiglu_limit, + ) + if self.experts_start_idx <= i < self.experts_end_idx + else None + for i in range(self.n_routed_experts) + ] + ) + assert args.n_shared_experts == 1 + self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit) + + def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> torch.Tensor: + shape = x.size() + x = x.view(-1, self.dim) + weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten()) + y = torch.zeros_like(x, dtype=torch.float32) + counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() + for i in range(self.experts_start_idx, self.experts_end_idx): + if counts[i] == 0: + continue + expert = self.experts[i] + idx, top = torch.where(indices == i) + y[idx] += expert(x[idx], weights[idx, top, None]) + if world_size > 1: + dist.all_reduce(y) + y += self.shared_experts(x) + return y.type_as(x).view(shape) + + +class Block(nn.Module): + """A block whose residual stream is `hc_mult` parallel copies (Hyper-Connections). + + Attention and FFN each sit between `hc_pre` (collapse the copies into one sublayer input) and + `hc_post` (expand back out, mixing the residual in through `comb`). `hc_mixes` derives all three + coefficient sets from the stream itself, `comb` made doubly stochastic by Sinkhorn. + + The coefficients a sublayer computes are used by the *next* one -- see `forward`. + """ + + attention_cls = Attention + + def __init__( + self, + layer_id: int, + args: ModelArgs, + engram_layout: EngramLayout | None = None, + ): + super().__init__() + self.layer_id = layer_id + self.norm_eps = args.norm_eps + self.attn = self.attention_cls(layer_id, args) + self.ffn = MoE(layer_id, args) + self.engram = None + if engram_layout is not None and layer_id in engram_layout.layer_ids: + self.engram = Engram(args, layer_id, engram_layout) + self.attn_norm = RMSNorm(args.dim, self.norm_eps) + self.ffn_norm = RMSNorm(args.dim, self.norm_eps) + self.hc_mult = hc_mult = args.hc_mult + self.hc_sinkhorn_iters = args.hc_sinkhorn_iters + self.hc_eps = args.hc_eps + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_attn_scale = nn.Parameter(torch.empty(3)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3)) + + def hc_mixes(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + """x: [b,s,hc,d], hc_fn: [mix_hc, hc*d], hc_scale: [3], hc_base: [mix_hc]. Returns the + pre / post / comb coefficients, split out of one projection of the flattened stream.""" + # normalized over the whole flattened hc*d stream, one statistic per token + x = x.flatten(2).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + return hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps) + + def hc_pre(self, x: torch.Tensor, pre_mix: torch.Tensor): + """Collapse the hc copies into one, weighted by pre_mix. [b,s,hc,d] x [b,s,hc] -> [b,s,d]""" + y = torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=2) + return y.to(x.dtype) + + def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor): + """Expand the sublayer output back to hc copies and mix the residual in through `comb`. + x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc] -> [b,s,hc,d]""" + y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) + return y.type_as(x) + + def forward( + self, + x: torch.Tensor, + start_pos: int, + pre_mix: torch.Tensor, + image_mask: torch.Tensor | None, + *attn_args, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`pre_mix` collapses the hc_mult copies down to one input for this block's attention. Each + sub-block's own `hc_mixes` produces the mix for the *next* one, so attention uses what the + previous layer's FFN produced and the FFN uses what this attention produced. + + image_mask: [b, s] bool, True inside image spans (selects the VL routing bias).""" + residual = x + attn_pre, attn_post, attn_comb = self.hc_mixes(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x = self.hc_pre(x, pre_mix) + x = self.attn_norm(x) + x = self.attn(x, start_pos, *attn_args) + x = self.hc_post(x, residual, attn_post, attn_comb) + + residual = x + ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + x = self.hc_pre(x, attn_pre) + x = self.ffn_norm(x) + x = self.ffn(x, image_mask) + x = self.hc_post(x, residual, ffn_post, ffn_comb) + return x, ffn_pre + + +class ParallelHead(nn.Module): + def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + self.norm_eps = norm_eps + self.hc_eps = hc_eps + self.part_vocab_size = vocab_size // world_size + # bf16 in the checkpoint, kept as fp32 here so the logits come out in fp32 directly + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor, full_logits=False): + """x: [b, s, d]. Generation only needs the last position, so that is the default.""" + if not full_logits: + x = x[:, -1] + logits = F.linear(x.float(), self.weight) + if world_size > 1: + all_logits = [torch.empty_like(logits) for _ in range(world_size)] + dist.all_gather(all_logits, logits) + logits = torch.cat(all_logits, dim=-1) + return logits + + +@lru_cache(1) +def get_dspark_topk_idxs(window_size: int, bsz: int, block_size: int, start_pos: int): + assert start_pos > 0 + matrix = torch.cat( + [ + torch.arange(min(window_size, start_pos + 1)), + window_size + torch.arange(block_size), + ] + ) + return matrix.int().view(1, 1, -1).expand(bsz, block_size, -1).contiguous() + + +class DSparkAttention(Attention): + def forward(self, x: torch.Tensor, start_pos: int, main_x: torch.Tensor): + assert self.compress_ratio == 0 + bsz, seqlen, _ = main_x.size() + win = self.window_size + rd = self.rope_head_dim + + main_freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + main_kv = self.kv_norm(self.wkv(main_x)) + apply_rotary_emb(main_kv[..., -rd:], main_freqs_cis) + act_quant(main_kv, fp8_block_size, scale_fmt, scale_dtype, True) + + if start_pos == 0: + if seqlen <= win: + self.window_kv_cache[:bsz, :seqlen] = main_kv + else: + cutoff = seqlen % win + self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = main_kv[:, -win:].split( + [win - cutoff, cutoff], dim=1 + ) + return x + + bsz, block_size, _ = x.size() + freqs_cis = self.freqs_cis[start_pos + seqlen : start_pos + seqlen + block_size] + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + + topk_idxs = get_dspark_topk_idxs(win, bsz, block_size, start_pos) + self.window_kv_cache[:bsz, start_pos % win] = main_kv.squeeze(1) + kv = torch.cat([self.window_kv_cache[:bsz], kv], dim=1) + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + o = o.view(bsz, block_size, self.n_local_groups, -1) + act_quant(o, fp8_block_size, scale_fmt, scale_dtype, True) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + +class DSparkMarkovHead(nn.Module): + def __init__(self, vocab_size: int, dspark_markov_rank: int): + super().__init__() + self.embed = ParallelEmbedding(vocab_size, dspark_markov_rank) + self.head = ParallelHead(vocab_size, dspark_markov_rank) + + def forward(self, token_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + embed = self.embed(token_ids) + logits = self.head(embed, full_logits=True) + return logits, embed + + +class DSparkConfidenceHead(nn.Module): + def __init__(self, input_dim: int): + super().__init__() + # proj in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for fp32 confidence score. + self.proj = Linear(input_dim, 1, dtype=torch.float32) + + def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor): + hidden = torch.cat([hidden, markov_embed], dim=-1) + return self.proj(hidden.float()).squeeze(-1) + + +class DSparkBlock(Block): + """DSpark stage stored under the mtp.* checkpoint namespace.""" + + attention_cls = DSparkAttention + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__(layer_id, args) + stage_id = layer_id - args.n_layers + self.block_size = args.dspark_block_size + self.noise_token_id = args.dspark_noise_token_id + self.temperature = args.temperature + if stage_id == 0: + assert len(args.dspark_target_layer_ids) > 0, "DSpark needs target layers" + self.main_proj = Linear(args.dim * len(args.dspark_target_layer_ids), args.dim) + self.main_norm = RMSNorm(args.dim, args.norm_eps) + if stage_id == args.n_mtp_layers - 1: + self.norm = RMSNorm(args.dim, args.norm_eps) + self.markov_head = DSparkMarkovHead(args.vocab_size, args.dspark_markov_rank) + self.confidence_head = DSparkConfidenceHead(args.dim + args.dspark_markov_rank) + self.embed: ParallelEmbedding | None = None + self.head: ParallelHead | None = None + + def forward(self, x: torch.Tensor, start_pos: int, pre_mix: torch.Tensor, main_x: torch.Tensor): + if start_pos == 0: + self.attn(x, start_pos, main_x) # prefill only seeds the window KV cache + return x, pre_mix + return super().forward(x, start_pos, pre_mix, None, main_x) # drafts are text: no VL bias + + def forward_embed(self, main_hidden: torch.Tensor, input_ids: torch.Tensor): + assert self.embed is not None + main_x = self.main_norm(self.main_proj(main_hidden)) + draft_input_ids = input_ids.new_full([input_ids.size(0), self.block_size], self.noise_token_id) + draft_input_ids[:, 0] = input_ids + x = self.embed(draft_input_ids) + x = x.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + return x, main_x + + def forward_head( + self, + x: torch.Tensor, + pre_mix: torch.Tensor, + input_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert self.head is not None + x = self.hc_pre(x, pre_mix) + logits = self.head(self.norm(x), full_logits=True) + output_ids = input_ids.new_empty(input_ids.size(0), self.block_size + 1) + output_ids[:, 0] = input_ids + markov_embeds = [] + for i in range(self.block_size): + logits_bias, markov_embed = self.markov_head(output_ids[:, i]) + logits[:, i].add_(logits_bias) + markov_embeds.append(markov_embed) + output_ids[:, i + 1] = sample(logits[:, i], self.temperature) + markov_embed = torch.stack(markov_embeds, dim=1) + confidence = self.confidence_head(x, markov_embed) + return output_ids, logits, confidence + + +def make_identity_pre_mix(x: torch.Tensor, hc_mult: int) -> torch.Tensor: + """initial one-hot mix""" + pre_mix = x.new_zeros(x.size(0), x.size(1), hc_mult, dtype=torch.float32) + pre_mix[:, :, 0] = 1.0 + return pre_mix + + +class SharedAttentionRuntime: + """What attention layers hand down the stack instead of recomputing. Layers run in order and + every source writes before its consumers read, so one slot each is enough and nothing needs + resetting between forwards. Sources: compress_kv and index_k from kv_source_layers, topk_idxs + from index_source_layers, candidates from candidate_source_layer.""" + + def __init__(self): + self.compress_kv: torch.Tensor | None = None + self.index_k: torch.Tensor | None = None + self.topk_idxs: torch.Tensor | None = None + self.candidates: torch.Tensor | None = None + + +# Only ever one model per process, same as world_size / rank / default_dtype above. +shared_attn = SharedAttentionRuntime() + + +class Transformer(nn.Module): + """DeepSeek-V4.1: embed -> expand to hc_mult copies -> blocks -> collapse -> logits. Building + this sets the globals at the top of the file. The tokenizer only feeds the engram token map.""" + + def __init__(self, args: ModelArgs, tokenizer=None): + global world_size, rank, default_dtype + world_size = dist.get_world_size() if dist.is_initialized() else 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16 + super().__init__() + self.max_seq_len = args.max_seq_len + self.temperature = args.temperature + self.norm_eps = args.norm_eps + self.hc_eps = args.hc_eps + self.engram_layout = EngramLayout.from_args(args) + self.engram_hash = ( + NgramHashState(args, self.engram_layout, tokenizer) if self.engram_layout is not None else None + ) + self.embed = ParallelEmbedding(args.vocab_size, args.dim) + self.layers = torch.nn.ModuleList() + for layer_id in range(args.n_layers): + self.layers.append(Block(layer_id, args, self.engram_layout)) + self.norm = RMSNorm(args.dim, self.norm_eps) + self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps) + self.mtp = torch.nn.ModuleList() + self.target_layer_ids = args.dspark_target_layer_ids + if args.dspark_block_size: + for layer_id in range(args.n_mtp_layers): + self.mtp.append(DSparkBlock(args.n_layers + layer_id, args)) + self.mtp[-1].embed = self.embed + self.mtp[-1].head = self.head + self.hc_mult = args.hc_mult + self.vision = None + if args.vision_enabled: + self.vision = ViT(args) + self.aligner = Aligner(args) + # learned embeddings for the image span delimiters + self.image_start = nn.Parameter(torch.empty(args.dim)) + self.image_end = nn.Parameter(torch.empty(args.dim)) + self.image_newline = nn.Parameter(torch.empty(args.dim)) + + @torch.inference_mode() + def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) + + def merge_image_embeddings(self, images, h: torch.Tensor): + """Overwrite each image's token span in h with its ViT/aligner features. The IMAGE slots take + the aligner rows in row-major order; the span delimiters take learned embeddings.""" + for i, sample in enumerate(images): + for img in sample or (): + types = img.types.to(h.device) + span = h[i, img.start : img.start + types.numel()] + span[types == IMAGE_START] = self.image_start.to(h.dtype) + span[types == IMAGE_END] = self.image_end.to(h.dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(h.dtype) + embeds = self.encode_image(img.patches.to(h.device), img.n_vit_h, img.n_vit_w) + span[types == IMAGE] = embeds.to(h.dtype) + + @torch.inference_mode() + def forward( + self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None + ): + """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it + has already filled, so the padding it uses internally never reaches here. token_types / + images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie + inside the first (start_pos 0) chunk.""" + image_mask = None if token_types is None else token_types >= 0 # TEXT is -1 + # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask + engram_mask = None if image_mask is None else ~image_mask + engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None + h = self.embed(input_ids) + if images is not None: + assert start_pos == 0, "image spans must be prefilled in a single chunk" + self.merge_image_embeddings(images, h) + # Expand to hc_mult copies for Hyper-Connections + h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + main_hiddens = [] + pre_mix = make_identity_pre_mix(h, self.hc_mult) + for i, layer in enumerate(self.layers): + if layer.engram is not None: + h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask) + # the MTP head reads the attention input of its target layers, not their output + if i in self.target_layer_ids: + main_hiddens.append(h.mean(dim=2)) + h, pre_mix = layer(h, start_pos, pre_mix, image_mask) + h = layer.hc_pre(h, pre_mix) + logits = self.head(self.norm(h)) + output_ids = sample(logits, self.temperature) + main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None + return output_ids, logits, main_hidden + + @torch.inference_mode() + def forward_spec(self, input_ids: torch.Tensor, main_hidden: torch.Tensor, start_pos: int = 0): + h, main_x = self.mtp[0].forward_embed(main_hidden, input_ids) + pre_mix = make_identity_pre_mix(h, self.hc_mult) + for layer in self.mtp: + h, pre_mix = layer(h, start_pos, pre_mix, main_x) + if start_pos == 0: + return None + return self.mtp[-1].forward_head(h, pre_mix, input_ids) + + +def sample(logits, temperature: float = 1.0): + """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU, + since it avoids the GPU-to-CPU sync in torch.multinomial.""" + if temperature == 0: + return logits.argmax(dim=-1) + logits = logits / max(temperature, 1e-5) + probs = torch.softmax(logits, dim=-1, dtype=torch.float32) + return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1) + + +if __name__ == "__main__": + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.manual_seed(0) + args = ModelArgs(dspark_block_size=6, dspark_target_layer_ids=(3, 4)) + x = torch.randint(0, args.vocab_size, (2, 150)) + model = Transformer(args) + + output_ids, logits, main_hidden = model(x[:, :128]) + model.forward_spec(output_ids, main_hidden) + for i in range(128, 150): + output_ids, logits, main_hidden = model(x[:, i : i + 1], i) + result = model.forward_spec(output_ids, main_hidden, i) + assert result is not None + output_ids, logits, confidence = result 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_e8_1bit_m1.py b/tests/test_e8_1bit_m1.py new file mode 100644 index 0000000..5b7380b --- /dev/null +++ b/tests/test_e8_1bit_m1.py @@ -0,0 +1,66 @@ +"""The 8-bit E8 residual kernels against a float64 decode of the same codes.""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") + +from keep.vq.e8 import e8_1bit_grid, e8_1bit_packed +from ramp.kernels import e8_1bit_m1 + +pytestmark = pytest.mark.skipif(not (getattr(mx, "metal", None) and mx.metal.is_available()), reason="needs a Metal GPU") +NONE = e8_1bit_m1.NO_RESIDUAL + + +def _layer(experts, out, inp, group_size, seed): + rng = np.random.default_rng(seed) + codes = rng.integers(0, 256, (experts, out, inp // 8), dtype=np.uint8) + scales = rng.uniform(0.01, 0.2, (experts, out, inp // group_size)).astype(np.float16) + dense = e8_1bit_grid().astype(np.float64)[codes].reshape(experts, out, inp) + dense *= np.repeat(scales.astype(np.float64), group_size, axis=-1) + return mx.array(codes), mx.array(scales), dense + + +def _close(got, want, tol): + got = np.array(got.astype(mx.float32)).astype(np.float64) + assert got.shape == want.shape, (got.shape, want.shape) + err = np.abs(got - want).max() / np.abs(want).max() + assert err <= tol, err + + +def test_pair_packed_keeps_every_value(): + pairs = e8_1bit_m1.pair_packed(e8_1bit_packed()) + nib = lambda p, n: ((p >> (4 * n)) & 0xF).astype(np.float64) + back = np.stack([(nib(pairs, k if d % 2 == 0 else k + 4) - 8) / 2 for d in range(8) for k in [d // 2]], axis=1) + np.testing.assert_array_equal(back, e8_1bit_grid()) + + +# group 32 -> 4 codewords per group (4-byte loads); 64 -> 8. inp=768 makes lanes take more than +# one step; out=93 leaves a ragged last tile. +@pytest.mark.parametrize("group_size,inp,out", [(32, 128, 96), (64, 768, 93), (256, 4096, 64)]) +@pytest.mark.parametrize("rows_per_simdgroup,simdgroups", [(1, 1), (8, 1), (4, 4)]) +@pytest.mark.parametrize("dtype,tol", [(mx.float32, 1e-5), (mx.bfloat16, 2e-2)]) +def test_both_layouts_match_the_dense_decode(group_size, inp, out, rows_per_simdgroup, simdgroups, dtype, tol): + codes, scales, dense = _layer(6, out, inp, group_size, seed=3) + book = mx.array(e8_1bit_m1.pair_packed(e8_1bit_packed())) + kw = {"output_dims": out, "rows_per_simdgroup": rows_per_simdgroup, "simdgroups": simdgroups} + rng = np.random.default_rng(4) + rhs = np.array([[5, 0, 3], [NONE, 2, 2], [1, NONE, 4]], dtype=np.uint32) + + def rows(n): # values as the kernel sees them, after rounding to dtype + x = mx.array(rng.normal(size=(n, inp)).astype(np.float32)).astype(dtype) + return x, np.array(x.astype(mx.float32)).astype(np.float64) + + def ref(x_row, e): + return np.zeros(out) if e == NONE else x_row @ dense[e].T + + x, xn = rows(3) + got = e8_1bit_m1.e8_1bit_m1_token_layout(x, codes, scales, book, mx.array(rhs), **kw) + assert got.dtype == dtype + _close(got, np.stack([[ref(xn[t], e) for e in rhs[t]] for t in range(3)]), tol) + + x, xn = rows(9) + got = e8_1bit_m1.e8_1bit_m1_per_route(x, codes, scales, book, mx.array(rhs.reshape(-1)), **kw) + _close(got, np.stack([ref(xn[r], e) for r, e in enumerate(rhs.reshape(-1))]), tol) diff --git a/tests/test_e8p_m1.py b/tests/test_e8p_m1.py new file mode 100644 index 0000000..cef5f38 --- /dev/null +++ b/tests/test_e8p_m1.py @@ -0,0 +1,195 @@ +"""Opt-in M=1 kernels for 16-bit E8P codes on GPUs without the NAX kernels.""" + +from __future__ import annotations + + +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}, True), # MTP verify window + ({"tokens": 32}, True), + ({"tokens": 33}, False), # prefill keeps the sorted-route kernels + ({"tokens": 0}, 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 + + +@pytest.mark.parametrize("metal,nax,expected", [(True, False, True), (True, True, False), (False, False, False)]) +def test_default_is_on_only_for_a_metal_gpu_without_nax(monkeypatch, metal, nax, expected): + from ramp.kernels import nax as nax_module + + monkeypatch.setattr(mx.metal, "is_available", lambda: metal) + monkeypatch.setattr(nax_module, "is_available", lambda: nax) + for value in (None, "", "auto"): + if value is None: + monkeypatch.delenv(e8p_m1.E8P_M1_ENV, raising=False) + else: + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, value) + assert e8p_m1.e8p_m1_enabled() is expected + + +def test_flag_overrides_the_default_and_is_validated(monkeypatch): + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "0") + 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, "1") + assert e8p_m1.should_use_e8p_m1(code_bits=16, tokens=1, sorted_indices=False, route_strategy="auto") is True + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "maybe") + with pytest.raises(ValueError, match=e8p_m1.E8P_M1_ENV): + e8p_m1.e8p_m1_enabled() + + +def _layer(code_bits=16, experts=4, out=32, inp=64, seed=0, group_size=32): + 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=group_size, 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("E8P M=1 kernel used with the flag off 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.setenv(e8p_m1.E8P_M1_ENV, "0") + _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") +# group 32 -> 4 codewords per group (4-codeword loads); 64 -> 8 (8-codeword +# loads). inp=768 makes lanes take more than one step; out=93 leaves a +# ragged last tile. +@pytest.mark.parametrize("group_size,inp,out", [(32, 128, 96), (64, 768, 93)]) +@pytest.mark.parametrize("rows_per_simdgroup,simdgroups", [(1, 1), (8, 1), (4, 4)]) +def test_token_and_per_route_kernels_match_the_dense_decode(group_size, inp, out, rows_per_simdgroup, simdgroups): + layer = _layer(experts=6, out=out, inp=inp, seed=3, group_size=group_size) + kw = {"output_dims": out, "rows_per_simdgroup": rows_per_simdgroup, "simdgroups": simdgroups} + dense = _dense(layer) + rng = np.random.default_rng(4) + rhs = mx.array([[5, 0, 3]], dtype=mx.uint32) + + x = rng.normal(size=(1, inp)).astype(np.float32) + got = e8p_m1.e8p_m1_token_layout(mx.array(x), layer.codes, layer.effective_scales(), layer.codebook, rhs, **kw) + 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, inp)).astype(np.float32) + got = e8p_m1.e8p_m1_per_route( + mx.array(xr), layer.codes, layer.effective_scales(), layer.codebook, rhs.reshape(-1), **kw + ) + 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) + + +@pytest.mark.skipif(not _metal_available(), reason="needs a Metal GPU") +@pytest.mark.parametrize("tokens", [2, 5]) +def test_multi_token_verify_layouts_match_the_dense_decode(tokens): + layer = _layer(experts=6, out=64, inp=128, seed=5) + dense = _dense(layer) + rng = np.random.default_rng(6) + rhs = np.stack([rng.choice(6, 3, replace=False) for _ in range(tokens)]).astype(np.uint32) + + x = rng.normal(size=(tokens, 128)).astype(np.float32) + got = e8p_m1.e8p_m1_token_layout( + mx.array(x), layer.codes, layer.effective_scales(), layer.codebook, mx.array(rhs), output_dims=64 + ) + want = np.stack([[x[t] @ dense[e].T for e in rhs[t]] for t in range(tokens)]) + np.testing.assert_allclose(np.array(got.astype(mx.float32)), want, rtol=1e-4, atol=1e-3) + + xr = rng.normal(size=(tokens * 3, 128)).astype(np.float32) + got = e8p_m1.e8p_m1_per_route( + mx.array(xr), layer.codes, layer.effective_scales(), layer.codebook, mx.array(rhs.reshape(-1)), output_dims=64 + ) + want = np.stack([xr[r] @ dense[e].T for r, e in enumerate(rhs.reshape(-1))]) + np.testing.assert_allclose(np.array(got.astype(mx.float32)), want, rtol=1e-4, atol=1e-3) + + +def test_glu_skips_the_sorted_prefill_path_for_few_token_e8p(monkeypatch): + from ramp.models.glm4_moe_adapter import QuantizedVQSwitchGLU + + glu = QuantizedVQSwitchGLU(gate_proj=_layer(out=32, inp=64), up_proj=_layer(out=32, inp=64), down_proj=_layer(out=64, inp=32)) + x = mx.zeros((1, 16, 64)) + indices = mx.zeros((1, 16, 6), dtype=mx.uint32) # 96 routes: past the sorted-prefill threshold + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "1") + assert glu._can_use_shared_sorted_route_prefill(x, indices) is False + monkeypatch.setenv(e8p_m1.E8P_M1_ENV, "0") + assert glu._can_use_shared_sorted_route_prefill(x, indices) is (glu.gate_proj.route_strategy != "direct") + + +def test_chunk_never_straddles_a_scale_group(): + assert e8p_m1._chunk(512, 8) == 8 # V4-Flash gate/up: 64 codewords per group + assert e8p_m1._chunk(16, 4) == 4 + with pytest.raises(ValueError, match="divisible by 4"): + e8p_m1._chunk(12, 2) + + +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, num_hash_layers=0, 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()) diff --git a/tests/test_runpod_v41_reap.py b/tests/test_runpod_v41_reap.py new file mode 100644 index 0000000..c1b52cb --- /dev/null +++ b/tests/test_runpod_v41_reap.py @@ -0,0 +1,356 @@ +"""CPU checks for the RunPod V4.1 REAP kit (scripts/runpod/v41_reap). + +The MoE below copies the reference ``inference/model.py`` (pinned revision) +Gate / Expert / MoE with world_size = 1 and nn.Linear in place of the FP8/FP4 +Linear. The patched forward must give the same output and the statistics must +match a direct float64 computation and pass KEEP's own validators. +""" + +from __future__ import annotations + +import importlib.util +import json +import struct +import types +from collections import Counter +from pathlib import Path + +import numpy as np +import pytest + +REPO = Path(__file__).resolve().parents[1] +KIT = REPO / "scripts" / "runpod" / "v41_reap" + + +def _load(name: str): + spec = importlib.util.spec_from_file_location(f"v41_{name}", KIT / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# --------------------------------------------------------------------------- collector + + +def _reference(torch): + nn, F = torch.nn, torch.nn.functional + ref = types.ModuleType("ref") + ref.world_size = 1 + + class Gate(nn.Module): + def __init__(self, n, dim, topk): + super().__init__() + self.topk, self.route_scale = topk, 1.5 + self.weight = nn.Parameter(torch.randn(n, dim) * 0.3) + self.bias = nn.Parameter(torch.randn(n) * 0.1) + + def forward(self, x, image_mask=None): + scores = F.softplus(F.linear(x.float(), self.weight.float())).sqrt() + indices = (scores + self.bias).topk(self.topk, dim=-1)[1] + weights = scores.gather(1, indices) + weights /= weights.sum(dim=-1, keepdim=True) + 1e-20 + weights *= self.route_scale + return weights, indices + + class Expert(nn.Module): + def __init__(self, dim, inter, swiglu_limit=7.0): + super().__init__() + self.w1 = nn.Linear(dim, inter, bias=False) + self.w2 = nn.Linear(inter, dim, bias=False) + self.w3 = nn.Linear(dim, inter, bias=False) + self.swiglu_limit = swiglu_limit + + def forward(self, x, weights=None): + dtype = x.dtype + gate = self.w1(x).float() + up = self.w3(x).float() + if self.swiglu_limit > 0: + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + gate = torch.clamp(gate, max=self.swiglu_limit) + x = F.silu(gate) * up + if weights is not None: + x = weights * x + return self.w2(x.to(dtype)) + + class MoE(nn.Module): + def __init__(self, layer_id, dim=16, inter=12, n=8, topk=2): + super().__init__() + self.layer_id, self.dim, self.n_routed_experts = layer_id, dim, n + self.experts_start_idx, self.experts_end_idx = 0, n + self.gate = Gate(n, dim, topk) + self.experts = nn.ModuleList([Expert(dim, inter) for _ in range(n)]) + self.shared_experts = Expert(dim, inter) + + def forward(self, x, image_mask=None): + shape = x.size() + x = x.view(-1, self.dim) + weights, indices = self.gate(x) + y = torch.zeros_like(x, dtype=torch.float32) + counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() + for i in range(self.experts_start_idx, self.experts_end_idx): + if counts[i] == 0: + continue + idx, top = torch.where(indices == i) + y[idx] += self.experts[i](x[idx], weights[idx, top, None]) + y += self.shared_experts(x) + return y.type_as(x).view(shape) + + ref.MoE = MoE + return ref + + +def test_patched_moe_matches_reference_and_direct_stats(): + torch = pytest.importorskip("torch") + F = torch.nn.functional + collect = _load("collect") + torch.manual_seed(0) + torch.set_grad_enabled(False) + ref = _reference(torch) + layers = [ref.MoE(i) for i in range(2)] + [ref.MoE(2)] # layer 2 plays an mtp block + xs = [torch.randn(1, 37, 16) for _ in layers] + expected = [m(x) for m, x in zip(layers, xs)] + + original = ref.MoE.forward + state = {"rank": 0} + stats = collect.SessionStats(torch=torch, n_layers=2, n_experts=8, dims={"hidden": 16, "down": 12}, device="cpu") + state["stats"] = stats + ref.MoE.forward = collect.make_moe_forward(ref, torch, F, None, state) + try: + got = [m(x) for m, x in zip(layers, xs)] + finally: + ref.MoE.forward = original + for a, b in zip(expected, got): + assert torch.equal(a, b) + + arrays = stats.arrays([0, 1]) + for layer, (m, x) in enumerate(zip(layers[:2], xs)): + flat = x.view(-1, 16) + w, idx = m.gate(flat) + count = np.zeros(8) + gnorm = np.zeros(8) + fmax = np.zeros(8) + hid = np.zeros((8, 16)) + down = np.zeros((8, 12)) + for row in range(flat.shape[0]): + for slot in range(idx.shape[1]): + e, g = int(idx[row, slot]), float(w[row, slot]) + ex = m.experts[e] + h = F.silu(ex.w1(flat[row]).clamp(max=7.0)) * ex.w3(flat[row]).clamp(-7.0, 7.0) + out = ex.w2(g * h).double() + count[e] += 1 + gnorm[e] += float(out.norm()) + fmax[e] = max(fmax[e], float(out.norm()) / g) + hid[e] += flat[row].double().square().numpy() + down[e] += h.double().square().detach().numpy() + np.testing.assert_array_equal(arrays["reap__route_count"][layer], count) + np.testing.assert_allclose(arrays["reap__gate_norm_sum"][layer], gnorm, rtol=1e-5) + np.testing.assert_allclose(arrays["reap__norm_max"][layer], fmax, rtol=1e-5) + np.testing.assert_allclose(arrays["importance_sum__hidden"][layer], hid, rtol=1e-5) + np.testing.assert_allclose(arrays["importance_sum__down"][layer], down, rtol=1e-4, atol=1e-6) + assert arrays["routes__expert_ids"].shape == (37, 2, 2) + assert arrays["routes__expert_ids"].dtype == np.int16 + + from keep.quality.dsv4_teacher_runner import validate_dsv4_calibration_capture + from keep.quality.reap import ReapStats + + validate_dsv4_calibration_capture(arrays, layers=[0, 1], num_experts=8) + assert ReapStats.from_arrays(arrays).routed_slots.tolist() == [74, 74] + torch.set_grad_enabled(True) + + +def test_head_capture_matches_full_logits_schema(): + torch = pytest.importorskip("torch") + F = torch.nn.functional + collect = _load("collect") + from keep.quality.dsv4_teacher_runner import validate_dsv4_logits_capture + + head = types.SimpleNamespace(weight=torch.randn(50, 8)) + head.forward = lambda x, full_logits=False: F.linear(x[:, -1].float(), head.weight) + positions = np.array([2, 3, 7]) + targets = torch.as_tensor([5, 9, 11]) + state = {"rank": 0, "positions": torch.as_tensor(positions), "targets": targets, "captured": {k: [] for k in collect.CAPTURE_KEYS}} + head.forward = collect.make_head_forward(head, torch, F, None, state, top_k=10, rows_per_slice=2) + x = torch.randn(1, 9, 8) + head.forward(x) + cap = {k: np.concatenate(v) for k, v in state["captured"].items() if v} + cap["positions"] = positions.astype(np.int32) + cap["target_token_ids"] = np.zeros(3, dtype=np.int32) + validate_dsv4_logits_capture(cap, top_k=10, supervised=3) + full = F.linear(x[0, positions], head.weight) + np.testing.assert_allclose(cap["logsumexp"], torch.logsumexp(full, -1).numpy(), rtol=1e-6) + assert (cap["topk_logit_ids"][:, 0] == full.argmax(-1).numpy()).all() + np.testing.assert_allclose(cap["target_logit"], full.gather(1, targets[:, None])[:, 0].numpy(), rtol=1e-6) + + # Student pass scored against that teacher: the same head gives KL 0 and full agreement. + lse = torch.as_tensor(cap["logsumexp"]) + state["captured"] = {k: [] for k in collect.CAPTURE_KEYS} + state["teacher"] = { + "ids": torch.as_tensor(cap["topk_logit_ids"].astype(np.int64)), + "logprob": torch.as_tensor(cap["topk_logit_values"].astype(np.float32)) - lse[:, None], + "targets": targets, + } + head.forward(x) + student = {k: np.concatenate(v) for k, v in state["captured"].items() if v} + assert set(student) == {"kl_on_teacher_topk", "student_nll", "top1_agree"} + np.testing.assert_allclose(student["kl_on_teacher_topk"], 0.0, atol=2e-3) + assert student["top1_agree"].all() + np.testing.assert_allclose(student["student_nll"], (lse - torch.as_tensor(cap["target_logit"])).numpy(), rtol=1e-5) + + +def test_pack_rules(tmp_path): + collect = _load("collect") + pack = tmp_path / "pack.jsonl" + rows = [ + {"id": "a", "session": "s1", "split": "calibration", "domain": "rust", "token_ids": [1, 2, 3], "target_spans": []}, + {"id": "b", "session": "e1", "split": "eval", "domain": "bash", "token_ids": [1, 2, 3, 4, 5], "target_spans": [[0, 2], [3, 5]]}, + ] + pack.write_text("".join(json.dumps(r) + "\n" for r in rows)) + sessions = collect.load_sessions(pack) + assert list(sessions) == ["s1", "e1"] + # targets 1 and 3, 4 (0 has no predecessor); logits are read one position earlier + assert collect.target_positions(rows[1]).tolist() == [0, 2, 3] + rows.append({**rows[1], "id": "c"}) + pack.write_text("".join(json.dumps(r) + "\n" for r in rows)) + with pytest.raises(ValueError, match="exactly one sample"): + collect.load_sessions(pack) + + +# --------------------------------------------------------------------------- pruner + + +def _write_safetensors(path: Path, tensors: dict[str, tuple[str, list[int], bytes]]) -> None: + header, offset, blobs = {}, 0, [] + for name, (dtype, shape, data) in tensors.items(): + header[name] = {"dtype": dtype, "shape": shape, "data_offsets": [offset, offset + len(data)]} + offset += len(data) + blobs.append(data) + raw = json.dumps(header).encode() + raw += b" " * (-len(raw) % 8) + path.write_bytes(struct.pack(" dict[str, tuple[str, list[int], bytes]]: + data = path.read_bytes() + (size,) = struct.unpack("= 3000 and report["tokens"]["d"]["eval"] >= 400 + sizes = Counter() + for r in cal: + sizes[r["session"]] += len(r["token_ids"]) + assert max(sizes.values()) <= 500 + assert len({r["session"] for r in ev}) == len(ev) and all(r["target_spans"] == [[1, len(r["token_ids"])]] for r in ev) + assert report["sources"]["fake"].get("non_english", 0) == 0 # the CJK row sits after the budgets are met + rows.insert(0, rows.pop()) + _, report = pack.build([source], render, eos=1, max_tokens=64, pack_tokens=0, session_tokens=500, + revisions={}, rows=lambda s, r: iter(rows)) # fmt: skip + assert report["sources"]["fake"]["non_english"] == 1 diff --git a/tests/test_runpod_v41_routed_precision.py b/tests/test_runpod_v41_routed_precision.py new file mode 100644 index 0000000..21a5147 --- /dev/null +++ b/tests/test_runpod_v41_routed_precision.py @@ -0,0 +1,135 @@ +"""Routed-precision pieces of the V4.1 REAP recovery kit: the residual cache, the reduced-rank +sidecar fit, the static precision allocator and the nested-code decoder.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import numpy as np +import pytest + +REPO = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO / "scripts" / "runpod" / "v41_reap")) +sys.path.insert(0, str(REPO / "src")) + +import local_eval # noqa: E402 +import routed_report # noqa: E402 +import precision_alloc # noqa: E402 + + +def test_residual_cache_hits_only_after_async_admission(): + ind = np.array([[0, 1], [0, 1], [2, 0]]) + wts = np.array([[0.5, 0.1], [0.5, 0.1], [0.5, 0.5]]) + counts = np.zeros((3, 5)) + hi = local_eval.residual_cache(ind, wts, k=1, tau=0.3, bounds=[3], warm=None, counts=counts) + # token 0: cold; expert 0 admitted. token 1: expert 0 resident, expert 1 never admitted (gate < tau). + # token 2: expert 2 not yet resident, expert 0 still is. Then 2 is admitted (evicting 0, k=1) and 0 + # is admitted again (evicting 2): one slot of cache thrashes, and each admission is a fetch. + assert hi.tolist() == [[False, False], [True, False], [False, True]] + assert counts[:, 4].tolist() == [2, 0, 1] + + +def test_residual_cache_resets_per_record_and_warm_start(): + ind = np.array([[0], [0], [0]]) + wts = np.ones((3, 1)) + counts = np.zeros((1, 5)) + hi = local_eval.residual_cache(ind, wts, k=1, tau=0.3, bounds=[2, 1], warm=None, counts=counts) + assert hi[:, 0].tolist() == [False, True, False] # the second record starts cold again + assert counts[0, 4] == 2 + hi = local_eval.residual_cache(ind, wts, k=1, tau=0.3, bounds=[3], warm=np.array([0]), counts=np.zeros((1, 5))) + assert hi.all() + + +def test_precision_allocator(): + precision_alloc._check() + + +def test_sidecar_reduced_rank(): + pytest.importorskip("torch") + import sidecar_fit + + sidecar_fit._check() + + +def test_nested_decoder_matches_local_eval(): + import keep_quant + from keep.vq.e8 import e8_1bit_grid + + rng = np.random.default_rng(0) + z = {"code_bits": np.asarray(8), "w1_codes": rng.integers(0, 256, (2, 4, 64), dtype=np.uint16), + "w1_scales": rng.random((2, 4, 1)).astype(np.float16)} + got = keep_quant.decode_codes(z, "w1", 1) + grid = e8_1bit_grid().astype(np.float32) + want = (grid[z["w1_codes"][1].astype(np.int64)] * z["w1_scales"][1].astype(np.float32)[..., None]).reshape(4, -1) + assert np.allclose(got, want) + + +def test_residual_cache_oracle_and_trace(): + ind = np.array([[0, 1], [2, 1], [2, 0], [3, 3]]) + wts = np.array([[0.9, 0.1], [0.5, 0.1], [0.5, 0.2], [0.3, 0.3]]) + counts = np.zeros((4, 5)) + # record 0 = tokens 0-2: expert 0 has mass 1.1, expert 2 has 1.0, expert 1 has 0.2 -> hot {0, 2} + hi = local_eval.residual_cache(ind, wts, k=2, tau=0.3, bounds=[3, 1], warm=None, counts=counts, policy="oracle") + assert hi.tolist() == [[True, False], [True, False], [True, True], [True, True]] + assert counts[:, 4].sum() == 0 # the oracle never fetches + trace: list = [] + local_eval.residual_cache(ind, wts, k=2, tau=0.3, bounds=[3, 1], warm=None, counts=counts, trace=trace) + assert trace == [(0, 0, 0), (0, 1, 2), (1, 0, 3)] # (record, token in record, expert); gate 0.2 never admits + + +def test_compare_bootstrap(tmp_path): + import json + + rng = np.random.default_rng(1) + sessions = {f"s{i}": {"domain": f"d{i % 2}"} for i in range(4)} + for name, shift in (("a", 0.1), ("b", 0.0)): + (tmp_path / name).mkdir() + for s in sessions: + np.savez(tmp_path / name / f"{s}.npz", kl=np.abs(np.random.default_rng(hash(s) % 97).normal(1, 0.2, 500)) + shift) + (tmp_path / "summary.json").write_text(json.dumps({"variants": {"a": sessions, "b": sessions}})) + out = routed_report.compare(f"{tmp_path}:a", f"{tmp_path}:b", reps=200) + assert abs(out["kl_diff"] - 0.1) < 1e-9 and out["ci95"][0] <= 0.1 <= out["ci95"][1] + + +def test_residual_cache_reports_the_resident_set(): + ind = np.array([[0, 1], [2, 1], [2, 0]]) + wts = np.array([[0.9, 0.1], [0.5, 0.1], [0.5, 0.2]]) + resident = np.zeros((3, 4), dtype=bool) + local_eval.residual_cache(ind, wts, k=2, tau=0.3, bounds=[3], warm=np.array([3, 1]), counts=np.zeros((4, 5)), + resident=resident) + # token 0 sees the warm set {3, 1}; expert 0 is admitted and evicts 3; then 2 evicts 1 + assert resident.tolist() == [[False, True, False, True], [True, True, False, False], [True, False, True, False]] + resident[:] = False + local_eval.residual_cache(ind, wts, k=2, tau=0.3, bounds=[3], warm=None, counts=np.zeros((4, 5)), + policy="oracle", resident=resident) + assert resident.all(axis=0).tolist() == [True, False, True, False] # the record's top two by mass: 0 and 2 + + +def test_gate_masked_routes_among_allowed_experts(): + torch = pytest.importorskip("torch") + import types + + torch.manual_seed(0) + gate = types.SimpleNamespace(weight=torch.randn(8, 5), bias=torch.zeros(8), gate_temp=1.0, score_func="sqrtsoftplus", + topk=3, norm_topk_prob=True, route_scale=1.5) + x = torch.randn(4, 5) + scores = torch.nn.functional.softplus(x @ gate.weight.T).sqrt() + w, i = local_eval.gate_masked(gate, x, torch.ones(4, 8, dtype=torch.bool), torch.nn.functional.linear) + assert torch.equal(i, scores.topk(3, dim=-1)[1]) # nothing masked: the reference router + assert torch.allclose(w.sum(-1), torch.full((4,), 1.5)) + allowed = torch.ones(4, 8, dtype=torch.bool) + allowed[:, i[0, 0]] = False + w2, i2 = local_eval.gate_masked(gate, x, allowed, torch.nn.functional.linear) + assert not (i2 == i[0, 0]).any() and torch.allclose(w2.sum(-1), torch.full((4,), 1.5)) # re-routed, renormalized + few = torch.zeros(4, 8, dtype=torch.bool) + few[:, :2] = True # fewer allowed than top-k: the surplus slot carries no weight + w3, i3 = local_eval.gate_masked(gate, x, few, torch.nn.functional.linear) + assert ((w3 > 0).sum(-1) == 2).all() and (i3[w3 > 0] < 2).all() + + + +def test_hot_set_takes_the_hottest_experts_the_checkpoint_holds(): + assert local_eval.hot_set(np.array([7, 2, 9, 4, 1]), 2, [1, 2, 4]) == {2, 4} + assert local_eval.hot_set(np.array([7, 2]), 3, [2, 7, 8]) == {2, 7} + assert routed_report.FP4_BYTES == 18_800_640 diff --git a/tests/test_runpod_v41_stream.py b/tests/test_runpod_v41_stream.py new file mode 100644 index 0000000..d015e3b --- /dev/null +++ b/tests/test_runpod_v41_stream.py @@ -0,0 +1,290 @@ +"""The layer-streamed driver against DeepSeek's own Transformer.forward, on CPU at toy size. + +The real ``model.py`` and ``engram.py`` (tests/fixtures/v41_reference) run with +plain-torch kernels. The toy model keeps every feature that makes streaming +subtle: + +- hyper-connections; +- compressed KV shared from source layers; +- an indexer shared across layers; +- a candidate-block source and a layer consuming it; +- an Engram layer; +- a MoE. +""" + +from __future__ import annotations + +import json +import sys +import types +from pathlib import Path + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +REPO = Path(__file__).resolve().parents[1] +REF_DIR = REPO / "tests" / "fixtures" / "v41_reference" +KIT = REPO / "scripts" / "runpod" / "v41_reap" +sys.path.insert(0, str(KIT)) + +VOCAB = 96 + + +def _install_stubs() -> None: + import torch_kernel + + sys.modules["kernel"] = torch_kernel + image = types.ModuleType("image_processor") + image.IMAGE = image.IMAGE_END = image.IMAGE_NEW_LINE = image.IMAGE_START = None + vision = types.ModuleType("vision") + vision.Aligner = vision.ViT = None + sys.modules.setdefault("image_processor", image) + sys.modules.setdefault("vision", vision) + try: + import sympy + except ImportError: # engram only needs isprime + sympy = types.ModuleType("sympy") + sympy.isprime = lambda n: n > 1 and all(n % p for p in range(2, int(n**0.5) + 1)) + sys.modules["sympy"] = sympy + + +@pytest.fixture(scope="module") +def toy(tmp_path_factory): + _install_stubs() + import stream + + ref = stream.import_reference(REF_DIR) + import engram + + engram.build_compressed_token_map = lambda tokenizer: (list(range(VOCAB)), VOCAB) + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(0) + + base = dict( + max_batch_size=1, max_seq_len=64, dtype="bf16", expert_dtype=None, vocab_size=VOCAB, + dim=64, moe_inter_dim=32, n_layers=6, n_mtp_layers=0, n_heads=4, n_routed_experts=8, + n_activated_experts=2, route_scale=1.5, swiglu_limit=10.0, q_lora_rank=32, head_dim=32, + rope_head_dim=8, norm_eps=1e-6, o_groups=2, o_lora_rank=16, window_size=8, + compress_ratios=(0, 2, 2, 1, 1, 0), kv_source_layers=(1, 3), index_source_layers=(1, 3, 4), + index_n_heads=2, index_head_dim=16, index_topk=4, candidate_source_layer=3, + candidate_topk_blocks=2, candidate_block_size=2, original_seq_len=64, + engram_layer_ids=(2,), engram_max_ngram_size=3, engram_vocab_size=40, engram_n_heads=2, + engram_head_dim=32, engram_pad_id=2, engram_compressed_vocab_size=VOCAB, + ) + layout = ref.EngramLayout.from_args(types.SimpleNamespace(**{**base, "engram_num_embeddings": (0,)})) + rows = int(sum(p for per_ngram in layout.primes[0] for p in per_ngram)) + args = ref.ModelArgs(**base, engram_num_embeddings=(rows,)) + model = ref.Transformer(args, tokenizer=None) + + gen = torch.Generator().manual_seed(1) + with torch.no_grad(): + for name, p in model.named_parameters(): + shape = p.shape + if p.dtype == torch.float8_e4m3fn: + p.copy_((torch.randn(shape, generator=gen) * 0.5).to(p.dtype)) + elif p.dtype == torch.float8_e8m0fnu: + p.copy_(torch.full(shape, 0.5).to(p.dtype)) + elif "norm" in name or name.endswith(("q_weight", "k_weight")): + p.copy_(1 + 0.1 * torch.randn(shape, generator=gen)) + elif "hc_" in name and "scale" in name: + p.copy_(0.5 + 0.1 * torch.randn(shape, generator=gen)) + else: + fan_in = shape[-1] if len(shape) > 1 else 1 + p.copy_(torch.randn(shape, generator=gen) / max(fan_in, 1) ** 0.5) + + from safetensors.torch import save_file + + ckpt = tmp_path_factory.mktemp("hf") + sd = {k: v.contiguous() for k, v in model.state_dict().items()} + names = sorted(sd) + half = len(names) // 2 + shards = {"model-00001.safetensors": names[:half], "model-00002.safetensors": names[half:]} + weight_map = {} + for shard, keys in shards.items(): + save_file({k: sd[k] for k in keys}, str(ckpt / shard)) + weight_map.update({k: shard for k in keys}) + (ckpt / stream.INDEX).write_text(json.dumps({"weight_map": weight_map})) + + records = [torch.randint(0, VOCAB, (1, n), generator=gen) for n in (37, 23, 57)] + return types.SimpleNamespace(ref=ref, stream=stream, args=args, model=model, ckpt=ckpt, records=records) + + +def _reference_hidden(toy, tokens): + captured = {} + handle = toy.model.norm.register_forward_hook(lambda m, i, o: captured.setdefault("h", o)) + try: + with torch.inference_mode(): + toy.model.forward(tokens, 0) + finally: + handle.remove() + return captured["h"] + + +def _streamer(toy): + stream = toy.stream + ckpt = stream.HFCheckpoint(toy.ckpt) + rows = {layer: stream.EngramRows.from_checkpoint(ckpt, layer, toy.ref.fp8_block_size) for layer in toy.args.engram_layer_ids} + return stream.LayerStreamer(toy.ref, toy.args, ckpt.load, rows, toy.model.engram_hash, "cpu") + + +@pytest.mark.parametrize("grouping", ["one-moe-batch", "moe-per-record"]) +def test_streamed_forward_matches_reference(toy, grouping): + reference = [_reference_hidden(toy, t) for t in toy.records] + streamer = _streamer(toy) + items = [ + toy.stream.Item(tokens=t, group=0 if grouping == "one-moe-batch" else i) for i, t in enumerate(toy.records) + ] + with torch.inference_mode(): + streamer.forward(items) + streamed = [streamer.final_hidden(item) for item in items] + for want, got in zip(reference, streamed): + assert got.shape == want.shape + torch.testing.assert_close(got.float(), want.float(), rtol=0, atol=0) + + +def test_record_alone_equals_record_in_chunk(toy): + """The shared_attn save/restore: a record's output does not depend on its chunk-mates.""" + + streamer = _streamer(toy) + alone = [toy.stream.Item(tokens=toy.records[2], group=0)] + crowd = [toy.stream.Item(tokens=t, group=i) for i, t in enumerate(toy.records)] + with torch.inference_mode(): + streamer.forward(alone) + streamer.forward(crowd) + torch.testing.assert_close(streamer.final_hidden(crowd[2]), streamer.final_hidden(alone[0]), rtol=0, atol=0) + + +def test_pack_bits_round_trip(toy): + mask = torch.rand(3, 5, 13) > 0.5 + assert torch.equal(toy.stream.unpack_bits(toy.stream.pack_bits(mask)), mask) + + +# --------------------------------------------------------------------------- the two sweeps (collect.py) + +LAYERS = list(range(6)) +DIMS = {"hidden": 64, "down": 32} + + +def _record(tokens, name, split): + ids = tokens[0].tolist() + return {"id": f"{name}-{len(ids)}", "session": name, "split": split, "domain": "code", "token_ids": ids, "target_spans": [[3, len(ids)]]} + + +def _reference_stats(toy, collect, records, dims): + stats = collect.SessionStats(torch=torch, n_layers=6, n_experts=8, dims=dims, device="cpu") + original = toy.ref.MoE.forward + toy.ref.MoE.forward = collect.make_moe_forward(toy.ref, torch, torch.nn.functional, None, {"rank": 0, "stats": stats}) + try: + hidden = [_reference_hidden(toy, torch.as_tensor([r["token_ids"]])) for r in records] + finally: + toy.ref.MoE.forward = original + return stats.arrays(LAYERS), hidden + + +def _sweep(toy, collect, monkeypatch, chunk, split, students=None, teacher_mask=None): + streamer = _streamer(toy) + monkeypatch.setattr(toy.ref.MoE, "forward", toy.ref.MoE.forward) # install_hooks patches it; undo at teardown + state = collect.install_hooks(streamer, top_k=16, rows_per_slice=7) + return collect.run_chunk(streamer, state, chunk, split=split, students=students or {}, dims=DIMS, device="cpu", + teacher_mask=teacher_mask), streamer + + +def _assert_stats_match(got, want): + assert set(got) == set(want) + for key, value in want.items(): + if np.issubdtype(value.dtype, np.integer): + np.testing.assert_array_equal(got[key], value, err_msg=key) + else: + np.testing.assert_allclose(got[key], value, rtol=1e-6, atol=1e-12, err_msg=key) + + +def test_calibration_sweep_matches_reference(toy, monkeypatch): + import collect + + r0, r1, r2 = toy.records + chunk = [("cal-0", [_record(r0, "cal-0", "calibration"), _record(r1, "cal-0", "calibration")]), + ("cal-1", [_record(r2, "cal-1", "calibration")])] + want = {name: _reference_stats(toy, collect, records, DIMS)[0] for name, records in chunk} + results, _ = _sweep(toy, collect, monkeypatch, chunk, "calibration") + assert set(results) == {("teacher", "cal-0"), ("teacher", "cal-1")} + for name, _ in chunk: + _assert_stats_match(results[("teacher", name)], want[name]) + assert results[("teacher", "cal-0")]["routes__expert_ids"].shape == (60, 6, 2) + + +def test_masked_calibration_measures_the_pruned_models_routing(toy, monkeypatch): + import collect + + chunk = [("cal-0", [_record(toy.records[0], "cal-0", "calibration")])] + gen = torch.Generator().manual_seed(3) + kept = {str(layer): sorted(torch.randperm(8, generator=gen)[:5].tolist()) for layer in LAYERS} + mask = collect.survivor_masks({"survivors": kept}, 8) + with torch.no_grad(): + full = {layer: toy.model.layers[layer].ffn.gate.bias.clone() for layer in LAYERS} + for layer, m in mask.items(): + toy.model.layers[layer].ffn.gate.bias[m] = float("-inf") + want = _reference_stats(toy, collect, chunk[0][1], DIMS)[0] + for layer, bias in full.items(): + toy.model.layers[layer].ffn.gate.bias.copy_(bias) + results, _ = _sweep(toy, collect, monkeypatch, chunk, "calibration", teacher_mask=mask) + got = results[("teacher", "cal-0")] + _assert_stats_match(got, want) + for layer in LAYERS: # nothing routes to a masked expert, and the unmasked run differs + assert got["reap__route_count"][layer][mask[layer].numpy()].sum() == 0 + plain, _ = _sweep(toy, collect, monkeypatch, chunk, "calibration") + assert not np.array_equal(plain[("teacher", "cal-0")]["reap__route_count"], got["reap__route_count"]) + + +def test_eval_sweep_teacher_and_students_match_reference(toy, monkeypatch): + import collect + + F = torch.nn.functional + chunk = [(f"eval-{i}", [_record(t, f"eval-{i}", "eval")]) for i, t in enumerate(toy.records[:2])] + gen = torch.Generator().manual_seed(7) + half = {"survivors": {str(layer): sorted(torch.randperm(8, generator=gen)[:4].tolist()) for layer in LAYERS}} + students = {"8": collect.survivor_masks({"survivors": {str(layer): list(range(8)) for layer in LAYERS}}, 8), + "4": collect.survivor_masks(half, 8)} + + want = {name: _reference_stats(toy, collect, records, {}) for name, records in chunk} + with torch.no_grad(): + full = {layer: toy.model.layers[layer].ffn.gate.bias.clone() for layer in LAYERS} + for layer, mask in students["4"].items(): + toy.model.layers[layer].ffn.gate.bias[mask] = float("-inf") + masked = {name: _reference_hidden(toy, torch.as_tensor([records[0]["token_ids"]])) for name, records in chunk} + for layer, bias in full.items(): + toy.model.layers[layer].ffn.gate.bias.copy_(bias) + + results, _ = _sweep(toy, collect, monkeypatch, chunk, "eval", students) + assert set(results) == {(s, name) for s in ("teacher", "8", "4") for name, _ in chunk} + weight = toy.model.head.weight.detach() + for name, records in chunk: + stats, (hidden,) = want[name] + teacher = results[("teacher", name)] + _assert_stats_match({k: v for k, v in teacher.items() if k in stats}, stats) + assert not any(k.startswith("importance_sum") for k in teacher) + + positions = collect.target_positions(records[0]) + targets = torch.as_tensor(records[0]["token_ids"])[positions + 1] + logits = F.linear(hidden[0, positions].float(), weight) + np.testing.assert_allclose(teacher["logsumexp"], torch.logsumexp(logits, -1).numpy(), rtol=1e-6) + np.testing.assert_allclose(teacher["target_logit"], logits.gather(1, targets[:, None])[:, 0].numpy(), rtol=1e-6) + np.testing.assert_array_equal(teacher["positions"], positions) + + everyone = results[("8", name)] + np.testing.assert_allclose(everyone["kl_on_teacher_topk"], 0.0, atol=2e-3) + assert everyone["top1_agree"].all() + + pruned = F.linear(masked[name][0, positions].float(), weight) + nll = torch.logsumexp(pruned, -1) - pruned.gather(1, targets[:, None])[:, 0] + np.testing.assert_allclose(results[("4", name)]["student_nll"], nll.numpy(), rtol=1e-5) + assert not np.allclose(results[("4", name)]["student_nll"], everyone["student_nll"]) + + +def test_chunk_sessions_keeps_sessions_whole(): + import collect + + sessions = [(n, [{"token_ids": [0] * size}]) for n, size in (("a", 5), ("b", 5), ("c", 20), ("d", 1))] + chunks = collect.chunk_sessions(sessions, budget=24, streams=2) + assert [[n for n, _ in c] for c in chunks] == [["a", "b"], ["c"], ["d"]] diff --git a/tests/test_runpod_v41_torch_kernel.py b/tests/test_runpod_v41_torch_kernel.py new file mode 100644 index 0000000..c16375b --- /dev/null +++ b/tests/test_runpod_v41_torch_kernel.py @@ -0,0 +1,96 @@ +"""The FP8/FP4 stand-ins in torch_kernel.py against block-loop references written like the kernels.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +REPO = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO / "scripts" / "runpod" / "v41_reap")) +sys.path.insert(0, str(REPO / "src")) + +import torch_kernel as k # noqa: E402 + + +@pytest.fixture(autouse=True) +def _keep_default_dtype(): + """The GEMMs return the default dtype; set fp32 here and leave other tests' default alone.""" + before = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + yield + torch.set_default_dtype(before) + + +def _e8m0(shape, lo=-6, hi=2): + return (2.0 ** torch.randint(lo, hi, shape).float()).to(torch.float8_e8m0fnu) + + +def test_decode_e2m1_matches_numpy_decoder(): + from keep.convert.fp4_expert import decode_fp4_e2m1 + + codes = torch.arange(256, dtype=torch.uint8).view(16, 16) + want = decode_fp4_e2m1(codes.numpy().view(np.int8)) + got = k.decode_e2m1(codes.view(torch.float4_e2m1fn_x2)) + np.testing.assert_array_equal(got.numpy(), np.asarray(want, dtype=np.float32)) + + +def test_act_quant_scales_are_pow2_and_inplace_is_quant_dequant(): + torch.manual_seed(0) + x = torch.randn(5, 96, dtype=torch.bfloat16) * 7 + y, s = k.act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu) + sf = s.float() + assert torch.equal(sf, 2.0 ** torch.log2(sf).round()) # powers of two + amax = x.float().unflatten(-1, (-1, 32)).abs().amax(-1) + assert (amax / sf <= 448).all() and (amax / sf > 224).all() # the smallest pow2 scale that fits + k.exact_quant = True + try: + z = x.clone() + k.act_quant(z, 32, "ue8m0", torch.float8_e8m0fnu, inplace=True) + finally: + k.exact_quant = False + deq = (y.float().unflatten(-1, (-1, 32)) * sf[..., None]).flatten(-2).to(torch.bfloat16) + assert torch.equal(z, deq) + + +def _block_loop(A, sa, B, sb_of, K, kb): + """sum over K blocks of (A_k @ B_k^T) * sa[:, k] * sb[:, k], as the tilelang kernels accumulate.""" + out = torch.zeros(A.shape[0], B.shape[0]) + for i in range(K // kb): + part = A[:, i * kb : (i + 1) * kb] @ B[:, i * kb : (i + 1) * kb].T + out += part * sa[:, i : i + 1] * sb_of(i)[None, :] + return out + + +def test_fp8_gemm_matches_block_loop(): + torch.manual_seed(1) + M, N, K, bs = 7, 70, 96, 32 + a = (torch.randn(M, K) * 50).clamp(-448, 448).to(torch.float8_e4m3fn) + b = (torch.randn(N, K) * 50).clamp(-448, 448).to(torch.float8_e4m3fn) + a_s, b_s = _e8m0((M, K // bs)), _e8m0(((N + bs - 1) // bs, K // bs)) + got = k.fp8_gemm(a, a_s, b, b_s, torch.float8_e8m0fnu, bs) + rows = torch.arange(N) // bs + want = _block_loop(a.float(), a_s.float(), b.float(), lambda i: b_s.float()[rows, i], K, bs) + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-3) + + +def test_fp4_gemm_matches_block_loop(): + torch.manual_seed(2) + M, N, K, act_bs = 6, 40, 128, 32 + a = (torch.randn(M, K) * 50).clamp(-448, 448).to(torch.float8_e4m3fn) + b = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8).view(torch.float4_e2m1fn_x2) + a_s, b_s = _e8m0((M, K // act_bs)), _e8m0((N, K // 32)) + got = k.fp4_gemm(a, a_s, b, b_s, torch.float8_e8m0fnu, act_bs) + B = k.decode_e2m1(b) + want = _block_loop(a.float(), a_s.float().repeat_interleave(act_bs // 32, 1), B, lambda i: b_s.float()[:, i], K, 32) + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-3) + + +def test_e2m1_round_ties_to_even(): + v = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, -0.25, -2.5, 6.0, 0.3]) + want = [0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0, -0.0, -2.0, 6.0, 0.5] + assert k._e2m1_round(v).tolist() == want diff --git a/tests/test_switch_routing.py b/tests/test_switch_routing.py index 97f6631..0269247 100644 --- a/tests/test_switch_routing.py +++ b/tests/test_switch_routing.py @@ -414,6 +414,9 @@ def recording_helper(x_arg, codes, scales, codebook, sorted_rhs, sorted_lhs, **k monkeypatch.setattr(glm4_moe_adapter.nax, "is_available", lambda: False) monkeypatch.setattr(glm4_moe_adapter, "gather_vqmm_sorted_routes", recording_helper) + # Without NAX, 16-bit E8P at <= 32 tokens now takes ramp.kernels.e8p_m1 by + # default; this pins the sorted-route fallback behind it. + monkeypatch.setenv("RAMP_E8P_M1", "0") with pytest.warns( RuntimeWarning,