Repository navigation
Phase 0 toward V4.1 on a 128 GB Mac: lossless mxfp8 residents and kill-test tooling - #3
Draft
jackmaxwil wants to merge 8 commits into
Draft
jackmaxwil wants to merge 8 commits into
jackmaxwil wants to merge 8 commits into
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.
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.