Repository navigation
Phase 0 toward V4.1 on a 128 GB Mac: lossless mxfp8 residents and kill-test tooling - #3
Merged
Merged
Conversation
…fp8) bind_deepseek_v4_flash_non_vq_weights(resident_format="mxfp8") binds every FP8 resident that lands on a plain nn.Linear, or on wo_a's MultiLinear, as its shipped e4m3 bytes and ue8m0 exponents in a weights-only mxfp8 module (nn.QuantizedLinear, or mlx-lm's QuantizedMultiLinear for wo_a) instead of decoding it to bf16. Every weight dequantizes bit-identically to the bf16 decode, at 8.25 instead of 16 bits per weight. Activations are not quantized. The default stays "bf16", so nothing changes until the Mac kill test passes. load_authenticated_dsv4_composite passes the option through, and the report lists the repacked tensors and the .scales they created, without counting the scales as logical tensors. keep.convert.fp8_block.repack_fp8_block_as_mxfp8 does the byte copy and refuses anything that would need rounding: scales that are not exact powers of two, a NaN scale byte, or a 32-wide group that would straddle two source blocks. Building the same weights through mx.quantize(mode= "mxfp8") is not exact (it changed 608 of 98,304 values in a check), which is why the copy goes through this function. benchmarks/bench_dsv4_mxfp8_residents.py is the kill test: bit-exactness against today's decode, then M=1 latency against bf16 and affine-8 g64, with wo_a timed in its grouped form. Kill if any bit differs, if mxfp8 is under 1.5x bf16, or if it is over 1.1x affine-8. On CPU MLX 0.31.2 the exactness half passes at all eight V4-Flash resident shapes; the speed half has to run on the Mac. Estimate from the roadmap's byte table: on V4-Flash this would shrink residents from about 15.7 GB to about 9.9 GB in memory and cut decode reads by a similar amount.
The roadmap found the quiet gate one-sided: a row could fault gigabytes of
its own working set back in and still stamp itself clean on zero pageouts
and swapouts. One post-fix row passed with 245,255 page-ins (3.83 GB).
collect_metric_snapshot now also reports pageins, swapins and compressor
decompressions (totals and deltas). New headline rows are schema_version 2
and carry an inbound_gate. A v2 row is clean only if its outbound deltas
are zero, as before, and it stayed under 4,096 page-ins, 4,096
decompressions (64 MiB each at 16 KiB pages) and zero swap-ins. vm_stat is
system-wide, so zero page-ins is not attainable: the cleanest Task 4 row
saw 144. The rows that faulted their working set back saw 7,296 to
304,047. A row may not declare a laxer gate than the harness's. Version 1
rows, including the committed Task 4 evidence, replay under their original
rule; the quiet record's keys are unchanged.
Also:
- --resident-format {bf16,mxfp8} on bind-proof, run-row and series binds
residents through the new opt-in mxfp8 path. Rows and the bind proof
record the format, and a summary refuses rows that mix formats.
- --rewarm-after-quiet (opt-in) re-runs the untimed warmup after the quiet
window, so the timed window measures steady state instead of pages macOS
compressed while the harness sat idle.
…ration pass keep.quality.reap scores each routed expert by REAP's criterion, the mean over the slots routed to it of gate weight times the L2 norm of its output before gating. The inputs come from what KEEP's teacher runner already computes: the model's own gate (bias-corrected selection, sqrtsoftplus weights, renormalisation) and the per-route expert outputs. Upstream REAP's code instead takes softmax of raw logits, picks top-k without the correction bias, and runs every expert on every token. Here the cost is a norm and five scatter-adds per layer, folded into float64 once per layer like the imatrix. - layer_major_prefill takes an optional ReapAccumulator; the forward stays bit-identical (tested alone and alongside the imatrix). - Calibration sessions now also write reap__* arrays. Capture validation checks them and requires their route counts to match the imatrix's. The imatrix finalizer ignores them, so existing runs still finalize. - finalize_dsv4_reap_stats sums them, and refuses if any session lacks them rather than silently under-scoring what that session routed to. - produce_dsv4_teacher_cache.py finalize-reap writes reap-stats.npz. reap-select picks a uniform survivor set and reports per-layer Jaccard of the pruned set against route frequency. That is the plan's kill test: Jaccard >= 0.9 means REAP adds nothing over the counts KEEP already had. - expert_saliency offers reap, frequency, weighted_frequency and ean_mean, plus optional shrinkage of rare experts toward the layer mean. - super_expert_candidates applies REAP's outlier rule to max |f|, so negative outliers are not invisible as they are upstream. - prune_checkpoint_plan maps checkpoint names to a uniform prune: it renumbers kept experts and slices gate weight, bias and bias_vl (V4.1's image-token bias, which must not be left at the old width). It refuses hash-routed blocks (tid2eid), which need a remapping policy. A test checks that the sliced gate picks the same experts with the same renormalised weights as the full gate with pruned experts masked out of selection.
…ier kill test) Keeping every routed expert of a model too large for RAM means a hot set in memory and the rest read from SSD on a miss, on the decode critical path. keep.quality.expert_residency answers whether that is fast enough before an offload engine is built: - RouteRecorder keeps the first 2,048 tokens of routing per layer. In a causal model those prefill routes are also the decode routes. DSV4 calibration sessions now record them (routes__expert_ids int16, routes__gate_weights float16, about 1 MB per session). - simulate_residency replays them through a per-layer hot set under three policies: static (a fixed REAP-ordered set), LRU, and decayed use counts with a half-life (oMLX's V4.1 policy). It can warm-start from REAP saliency. fetch_threshold models a hybrid that drops light misses as a pruned model would and reports the gate mass it drops. simulate_sessions resets residency at each session boundary and pools raw counts. - residency_report applies the plan's rule per hot-set size: KILL if the best policy that fetches every miss still misses more than 5% of routed slots. decode_tokens_per_second turns misses into a tok/s estimate for a given SSD rate. - benchmarks/replay_dsv4_expert_residency.py runs it over a calibration run's sessions plus reap-stats.npz. layer_major_prefill's new hook is generalised to route_observers, so the REAP accumulator and the route recorder ride the same pass. The forward is still bit-identical.
Every hyper-connection turns hc_mult x hc_mult mixing logits into a doubly-stochastic matrix through a row softmax and 19 more rounds of row/column normalisation. The model's compiled path still launches a few kernels per round, because each normalisation is a reduction, and it runs 86 times per decoded V4-Flash token. The roadmap measured 30-70 ms of sinkhorn per token. ramp.kernels.hc_sinkhorn.fused_hc_sinkhorn does the whole sinkhorn in one Metal dispatch, one thread per matrix, with the values in registers, following the compiled path's operation order. exp and summation order are not guaranteed bit-identical, so it is opt-in: RAMP_HC_SINKHORN=metal selects it in HyperConnection (also used by the GLM-5 next adapter), and the default path is untouched. _hc_split_logits_ops computes pre, post and the comb logits exactly as the compiled path does. hc_sinkhorn_reference states the algorithm in float32 NumPy; a test pins it to the compiled path on CPU. Metal-only tests compare the fused kernel with the compiled path and HyperConnection's outputs with the flag on and off; they skip without a Metal GPU. benchmarks/bench_hc_sinkhorn.py is the kill test: KILL if the fused path is under 3x faster at decode shape or differs by more than 2e-6 relative.
…t NAX The shipped DeepSeek-V4-Flash artifact uses 16-bit E8P codes. Every M=1 fast path in QuantizedVQSwitchLinear is gated on 8-bit codes, and the E8P tensor-unit kernels are M5-only (RAMP refuses them on an M4 Max, where 36 of 94 kernel tests are wrong). So on M1-M4 every decoded token's routed projections fall through to the scalar kernel, which spends a whole threadgroup per output element. gather_vqmm_m1.metal already had an E8P branch that nothing dispatched to. With RAMP_E8P_M1=1 (opt-in), M=1 calls with 16-bit codes, unsorted routes and the auto/direct strategy use: - e8p_m1_token_layout for gate/up: the existing M=1 kernel with CODE_BITS=16, one activation row against top-k experts; - e8p_m1_per_route for down: the same kernel source with the activation read offset by route. It is built as a separate kernel from a two-line text substitution, so the production 8-bit kernel is byte-identical. A test proves undoing the substitution restores the file. The flag is off by default. Tests pin the decision table and the dispatch wiring, and Metal-only tests compare both kernels with a dense decode of the same codes. benchmarks/bench_e8p_m1.py is the M4 kill test. It times one routed FFN at decode shape against today's dispatch, sweeps row packing and threadgroup codebook staging, and checks the outputs agree. Verdict: PASS at >= 200 GB/s effective, TUNE at 100-200 (next step: decode each codeword once instead of once per dimension), KILL under 100 (serve on the M5). Its synthetic projections reproduce the roadmap's 2,129,920 bytes per projection per expert.
scripts/phase0_kill_tests.sh runs every Phase 0 kill test that needs no model download (mxfp8 residents, fused sinkhorn, single-token E8P, and the Metal parity tests). It refuses to start with custom wired limits set, writes JSON evidence under artifacts/benchmarks/phase0-<host>-<date>/ and prints one verdict per test. If KEEP_RESIDENT_DIR and KEEP_VQ_DIR are set, it also runs the real-weight variants. docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.md lists each opt-in switch, its kill criterion and the decision a pass unlocks. It also gives the calibration steps the REAP and residency tests need, which must write into a new output directory because the runner skips existing sessions.
Records the ranked plan for running DeepSeek-V4.1-Flash on one 128 GB Mac, what Phase 0 built and verified on this branch, the next actions in order (Mac kill tests, REAP calibration rerun, CI fix), and the open decisions.
mxfp8 residents: exact at all eight shapes, but KILL on speed. At M=1 it runs 0.88x bf16 and 1.31x affine-8 (three runs: 0.88-0.90x bf16). E8P M=1: outputs agree, KILL at 52.8 GB/s best (2.38x production, rows_per_threadgroup 32). Sinkhorn: PASS at decode (3.69x, 5.7e-7), but 0.50x at 1024 rows. Also fix test_hyper_connection_output_with_the_flag, which skipped on Linux and failed on Metal: one-layer args need num_hash_layers=0.
A decode-once M=1 kernel (popcount parity, one abs-grid lookup per codeword, wide code loads) ran at 113 GB/s against 154 GB/s for the existing per-dimension branch, measured kernel-only with 256 calls per sync and best of 20 interleaved runs. The compiler already lifts the per-codeword work out of the 8-dimension loop. The kernel and both scripts are kept under the Phase 0 evidence directory; the source tree is unchanged. The kill test's KILL comes mostly from syncing after every FFN. With 43 dependent FFNs per sync, as in one decode step, RAMP_E8P_M1=1 reaches 118 GB/s (TUNE) against 50 GB/s for production.
v1 synced after every FFN, so its KILL at 52.8 GB/s mostly measured CPU-GPU round trips. v2 chains --ffns-per-sync dependent FFNs (default 43, one per V4-Flash layer) before one mx.eval and reports the median time per FFN. On the M4 Max, two runs both give TUNE: 122-126 GB/s best with 8-16 rows per threadgroup, 2.4-2.5x production's 50.5 GB/s, outputs within 6.4e-4. This corrects the 118 GB/s in the previous commit, whose chain script fed a [1, 6, 4096] activation forward and so sent gate and up through the per-route kernel.
v1 synced after every sub-millisecond matmul, so it mostly measured CPU-GPU round trips and called KILL (mxfp8 0.88x bf16). v2 chains --calls-per-sync dependent calls (default 43, one per V4-Flash layer) before one mx.eval; each call's activation depends on the previous output so the GPU cannot overlap them. On the M4 Max both runs PASS: mxfp8 is 1.66x bf16 and takes 1.01-1.05x affine-8's time, still bit-exact at all eight shapes. The large tensors read at 350-450 GB/s.
The machine has the 40-core GPU (546 GB/s theoretical). Reductions over 1 GiB buffers read at 474-477 GB/s and a bf16 matvec at 458 GB/s. The plan's 409.4 GB/s estimates are therefore conservative.
RAMP_E8P_M1=1 now dispatches to e8p_m1.metal instead of the generic M=1 kernel's E8P branch, for both gate/up (token layout) and down (per route). Each simdgroup owns several output rows, lanes load 4 or 8 codewords at once (whichever divides a scale group), and each threadgroup decodes the abs grid once into half2 pairs; signs are applied by XOR on the half2 bits. The decode is bit-identical to the reference on all 65,536 codewords. On the M4 Max the v2 kill test (43 chained FFNs per sync) now PASSES narrowly: 201.9-205.5 GB/s best over two runs, 3.97x production, up from 122-126 GB/s. A loads-only stub reaches ~410 GB/s, so decode ALU still takes half the time. The flag stays opt-in. The text-substituted per-route copy of gather_vqmm_m1.metal is gone; that file and the 8-bit M=1 path are unchanged. Tuning scripts and kernel variants are kept under the Phase 0 evidence directory.
With RAMP_E8P_M1 unset (or auto), 16-bit E8P codes at M=1 now use e8p_m1.metal on any Metal GPU where ramp.kernels.nax.is_available() is false, which means M1-M4. NAX GPUs and CPU-only runs are unchanged. RAMP_E8P_M1=0 turns it off and =1 forces it on. Across the 29 test files that touch these modules on the M4 Max, the only difference is one flaky speculative-runtime test, which failed only with the default off. The kill test passed by 1-3% on a loaded machine, so it should be rerun on an idle one.
…non-NAX GPUs The tuned E8P kernel now takes any number of routes: route r multiplies activation row r / X_DIV (top_k for gate/up, 1 for down) by expert rhs_indices[r]. should_use_e8p_m1 accepts 1-32 token rows, and the GLU's shared sorted-route prefill path declines those calls, so verify windows no longer fall to the scalar or sorted-route kernels on M1-M4. Synthetic V4-Flash FFN on the M4 Max (256 experts, top-6, 43 chained FFNs per sync): 3.9x the previous path at 1 token, 4.7-5.5x at 2-8, 2.6-3.3x at 11-21 and 9x at 32, outputs within 5.7e-4. Prefill above 32 tokens is unchanged. default_enabled() is no longer cached (nax.is_available already is), and the no-native fallback test now pins RAMP_E8P_M1=0.
…act pruner) Adds scripts/runpod/v41_reap/, a kit that runs on a rented 8xH200 pod and produces new DeepSeek-V4.1-Flash-1017 artifacts, so the Mac downloads a pruned checkpoint instead of the 510 GB source. - collect.py: runs DeepSeek's reference torch code under torchrun (model.py sha256 is pinned) with MoE.forward patched to record per-session REAP stats, imatrix sums (hidden/down), route traces and top-2048 eval logits. It writes KEEP's session schema, so finalize-reap, reap-select and finalize consume the output unchanged. A masked-router student pass scores each candidate expert count against the teacher (KL, NLL change, top-1 agreement) from a single converted checkpoint. - prune.py: copies safetensors bytes (FP8/FP4 are never decoded, no torch), renumbers the kept experts, slices gate weight/bias/bias_vl, leaves mtp.* untouched, updates both configs, and writes the index and SHA256SUMS. - build_pack.py: builds an English-only calibration and held-out pack from public datasets in V4.1's own chat/DSML format, with group-disjoint splits. - score.py and pod.sh (fetch/pack/gpu/prune stages; the GPU stage stops the pod at MAX_GPU_HOURS). - Runbook, plus the calibration-data research (including vetting of the five suggested datasets). All of this is verified on CPU only (tests/test_runpod_v41_reap.py and a local end-to-end through the real finalizers). Nothing has run on a GPU yet.
8xH200 was out of capacity in every data center that offers network volumes, so 'pod.sh all' now runs the whole job on one pod: - fetch and pack in parallel, then convert, the teacher pass and scoring (N = 256/192/128/96); - delete the TP8 copy, then prune N=192 (REAP 50%) and N=96 (REAP 75%). The pod stays up for the Mac to pull the results; the watchdog stops it at MAX_GPU_HOURS, and any failed stage stops it at once.
The reference Block modules now run one layer at a time over chunks of whole records, loading each layer straight from the HF shards. No convert.py, no TP8 copy, and any 80 GB+ GPU works. - stream.py: HFCheckpoint (per-layer loads with convert.py's fixups), EngramRows (host memmap gather of the FP8 tables), LayerStreamer (per-record save and restore of shared_attn, candidates bit-packed). - collect.py: --sweep calibration|eval, --shard I/N, --chunk-tokens. The eval sweep runs the teacher and every masked-router student side by side. The npz schemas are unchanged. - pod.sh: one collect.py per GPU per sweep; the convert step and the TP8 directory are gone. - tests: the pinned model.py and engram.py run on CPU at toy size with plain-torch kernels. The streamed forward matches Transformer.forward bit for bit. Calibration stats, teacher logits and student scores match the reference.
… smoke stage, prune 192)
…ed it after a complete pack)
The hidden state carried between layers costs about 63 KB per resident token on the GPU. With 524288-token chunks it held about 65 GB, and the indexer's 16 GiB allocation for a 16K record ran out of memory. CHUNK_TOKENS is now a knob with a default of 262144.
…sidecars, routed tiers Local tooling for recovering the published DeepSeek-V4.1-Flash REAP checkpoints on the M4 Max: - local_eval.py scores router-masked variants against the pod's teacher logits, with KEEP code overrides, imatrix capture, residual sidecars, and per-token precision tiers (--dynamic, --tau) under a no-stall LRU residual cache with async SSD admission (--cache-k, --cache-warm). - keep_quant.py fits E8P 2-bit and E8 1-bit codes with pod or pruned-model importance, and nested residual codes on top of another run (--residual-of). - precision_alloc.py: static per-expert rung allocation under a unit budget. - sidecar_fit.py: sequential reduced-rank ridge sidecars fit against the teacher MoE, plus sidecar_pod.sh for a one-GPU pod run. - Recovery log with steps 1-2d: selection on pruned-model routing explains the static allocation gain; second-pass pruning and pruned-model importance add nothing.
…er-track plan in the log
… 1-bit codes as uint8
…equal memory (C1 holds in the pool of 192)
…sliced router decodes at 30 tok/s
… the switch, in-place slot writes)
…hape; benchmark script
…X (144 residents, 26-28 tok/s on the M4 Max)
…he 2-bit codes or the cache
…ttice cache, correct script, 20.7 tok/s)
…ob for the 192 removed experts
…d as one safetensors stream
… (two code directories, full router file)
…s, 23-27 tok/s, 63 GiB)
…84_pod.sh, the tau sweep over the pool of 384
…checkpoint from codes alone
…t 0.763 static; tau 0.3 as default)
…apping; Metal pages in the whole buffer)
…t's FP4 experts; K5 result in the build plan
…ard's settings, tool-use KL note corrected
…th static 144), FP4 on the hottest experts (gate not met below 200 residents)
…il: 0 and 2 of 8)
… by 0.072), 24-sample script check, temperature and top_p rescue checks
…cripts, no spelling failures, 87 GiB peak)
…us a 1-bit residual; B2 result (64 hot at 3 bits level with 32 at FP4 at 92 GB)
…uto); docs: 180 residents with the drafter (21 of 24 sampled scripts, median 32 tok/s, 88 GiB)
… bundle (366 GB to 226 GB checkpoint, relative pool paths); verified from a clean oMLX directory
…sted tier, tested against a float64 decode
…ok/s with 40 GB held, 36.7 before); run_v41.sh sets tau 0.3; rubric item 6 to 10
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Phase 0 of the plan to run DeepSeek-V4.1-Flash on one 128 GB Mac (target: the M4 Max). This is the work that needs no V4.1 download and can be proven on the V4-Flash artifact. Every runtime change is opt-in; defaults stay identical to
mainuntil the matching kill test passes on the Mac.Start here:
docs/deepseek-v4-flash/2026-11-05-v41-plan-and-phase0-status.mdhas the ranked plan, what this branch did, the next actions and the open decisions.What landed
3c6d9f1wo_a) kept as their source bytes in weights-only mxfp8, instead of decoded to bf16. 8.25 instead of 16 bits per weight, every weight bit-identicalresident_format="mxfp8"benchmarks/bench_dsv4_mxfp8_residents.py450c925--resident-formatand--rewarm-after-quietbd7fd54finalize-reap,reap-select(REAP vs. route-frequency Jaccard) and a router/expert slicing plan that includesbias_vland refuses hash-routed layersreap-select33a5681benchmarks/replay_dsv4_expert_residency.py5baed30RAMP_HC_SINKHORN=metalbenchmarks/bench_hc_sinkhorn.py00580c4RAMP_E8P_M1=1benchmarks/bench_e8p_m1.py7deda77scripts/phase0_kill_tests.sh(one command on the Mac) and a runbook,docs/deepseek-v4-flash/2026-11-05-phase0-kill-tests.mdb392eb3Verified on Linux (CPU MLX 0.31.2)
mx.quantize(mode="mxfp8")is not exact (608 of 98,304 values changed).main's (308 failures that need Metal or Mac-local files), plus 22 new passing tests.CI:
validate-recipesfails onmaintooerror: recipe 'glm52_reap_504b_family_evidence_v2_20260710' requires strict input hashes and does not allow --no-hash.mainhas failed this workflow with the same error on every one of its runs, the latest being run 37649131452 ate04a363. This PR touches no recipe and no build code.--no-hash, but two recipes declarestrict_inputs: true, andkeep/build/runner.pydeliberately refuses unhashed plans for those.mainor here.Not verified: needs the Mac
Every Metal kernel's parity and every speed number. Run:
The REAP and residency tests also need a fresh V4-Flash calibration run into a new output directory. This stays a draft until those results are in.
Known gaps
nax_e8p_m32n64(M5-only), so MTP rows cannot validate on an M4 yet.