Phase 0 toward V4.1 on a 128 GB Mac: lossless mxfp8 residents and kill-test tooling

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")

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]);
+    }
+}

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))

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"
+}

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"
+}

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"
+}