Repository navigation
Nemotron 3.5 Bolt (nemotron_h_puzzle) on Metal — serves end to end - #1038
trickyfalcon wants to merge 12 commits into
Conversation
Part 1 of the native Metal backend for nemotron_h_puzzle (Nemotron 3.5 Bolt 74B-A9B), opened early per ashhart#962 so the loader can be reviewed before the latent kernels harden. v1.1.0 already parses the Puzzle config (bolt, per-layer moe[]) and serves it on CUDA; this is the Metal weight loader for it. weights.zig: the Metal MoE loader now reads each layer's own expert count and width (per-layer in nemotron_h_puzzle, uniform otherwise), runs the routed experts in the moe_latent_size latent rather than hidden, and carries the two bf16 latent projections fc1/fc2_latent_proj (new Moe.into and .out, a plain-bf16 Dense referenced in place; the Lightning path keeps them null and is unchanged). moe() takes the layer's experts, width and latent; its two callers pass the per-layer values for bolt, the globals otherwise. config.zig: checkShapes gains a bolt branch (checkBoltShapes) for Bolt's dimensions (hidden 4096, 1024 latent, 5376 shared expert, top-k 9, per- layer expert counts 240-496 and widths 1536-2560). weights_test.zig: a two-layer puzzle block loads per layer - each keeps its own expert count and width, the experts read the latent, and both latent projections are present. No forward or kernels yet; those follow as commits on this branch. The Lightning path is byte-for-byte unchanged (zig build test test-golden: 0 unexpected golden differences). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
The latent path for Nemotron 3.5 Bolt, mirroring cuda_latent.zig run(): the normed input projects into the moe_latent_size latent, the routed experts run up/down in the latent, each pick projects back to hidden, then add_norm combines the picks by weight with the shared expert (kept in hidden) and the residual. Per-pick out-projection is valid because the projection is linear. Hand-authored (the nemotron kernel generator is not in the public repo). state.zig: two latent-width scratch buffers (x_lat, ey_lat). weights.zig: the routed experts load through the runtime-dim row kernels for the Puzzle (latent > 0), which read each layer's own expert count and width at K = latent; the baked 4/64 lane kernels serve only the Lightning inner width. New Metal kernels (zig/kernels/metal/nemotron): - route_bolt: the router top-k at runtime NE and K, so one kernel serves every Puzzle layer's own expert count. Same scoring as route.metal (sigmoid of the gate logit plus the correction bias, normalised and scaled) but the per-lane slices are ceil(NE / 32) with an e < NE guard, so the last partial slice is kept (every Bolt expert count is NE % 32 == 16, which the baked route.metal would drop). - dense_bf16: a bf16 write-matmul (y = x R^T, no accumulate) for the gate and the two full-precision latent projections fc1/fc2_latent_proj. - add_norm_moe_xs_bolt: the combine + shared + residual + next-norm kernel at <4096, 512, 9> (896 does not divide 4096; 512 does). layers.bolt_moe + forward dispatch (c.bolt): gate (dense_bf16) -> route_bolt -> shared up/down (hidden, the Lightning path) -> into (dense_bf16) -> routed experts up/down in the latent (indexed row kernels) -> out per pick (dense_bf16) -> add_norm_moe_xs_bolt. Lightning is unchanged: the dispatch takes bolt_moe only when c.bolt, and zig build test test-golden holds 0 unexpected golden differences. Remaining to RUN end to end: Bolt-dim variants of the shared up_relu2 and the mamba kernels, and the GPU kernel-contract tests (route_bolt first). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
…t slice tf-metal-tests now dispatches route_bolt at a real Bolt expert count (NE=240, NE % 32 == 16) and checks its picks and weights against a CPU reference. Row 0 plants its top nine logits on experts 224..232 - in the last partial slice [224,240) that the baked NE/32 striping dropped - and the kernel returns exactly those, so the pad-to-32 fix is proven on device. Weights match the reference (sigmoid + bias, normalised, scaled) within 1e-3. route_bolt.metal and dense_bf16.metal join the metallib source list, which compiles and validates their MSL at build time (the engine still loads them from the runtime kernel registry). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
dense_bf16: the bf16 write-matmul (gate and latent projections) checked against an fp32 CPU reference over the same bf16 inputs, to a bf16 output tolerance. add_norm_moe_xs_bolt: the <4096, 512, 9> instantiation, checked on the case with no routed picks (WE = 0) and no shared expert (SH = 0), where the block reduces to OUT = rmsnorm(H) * W with W = 1; HN is unchanged and XS is the per-group sum of OUT. add_norm_moe_xs.metal joins the metallib source list so the instantiation compiles at build time and tf-metal-tests can load it. tf-metal-tests: 9 of 9 pass on an M-series Mac. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
The Metal engine now opens, loads and serves the published Bolt checkpoint
(bytebunker/NVIDIA-Nemotron-3.5-Bolt-74B-A9B-MLX-4bit, 40.3 GiB) on an
M-series Mac without tensor units, where every projection runs on the
runtime-dim row kernels. On an M3 Max (128 GB) it loads in 6.3 s and decodes
at ~41 tok/s.
Checked on that build:
- greedy tokens are identical to mlx_lm 0.32.3 (token_sha256 equal) on the
prompts tried;
- concurrent == solo at two and four streams;
- resumed == fresh: a chat turn resumed from the prompt cache (at 915 of 959
tokens) equals the same turn prefilled on a fresh server.
What Bolt needed beyond the latent MoE:
- metal.zig: register the nemotron_h_puzzle family and open it.
- weights.zig / forward.zig: a full-precision bf16 lm_head (MLX quantizes a
module only where {path}.scales exists, and Bolt's lm_head has none) is
read in place and projected by dense_bf16.
- Mamba at Bolt's widths (inner 128 x 64 = 8192, in_proj 18560): the tree
conv and scan instantiated at <8192, 8, 128, 4, 18560, 8192> and
<128, 64, 8, 128, 8192, 18560, 18432, 1048576>, and the gated group norm
at XD 8192 / GS 1024 (group_norm_xs_bolt.metal, one per padded row count).
- add_norm_xs at <4096, 512>, and the hidden row sums at the model's width.
- backend.zig: prompts run through the decode kernels in fused_rows chunks;
the fused prefill kernels are not yet built for the latent MoE.
All Bolt kernels are hand-authored instantiations of the generated ones (the
generator is not in the public repo). Lightning is unchanged: zig build test
test-golden holds 0 unexpected golden differences; tf-metal-tests 9 of 9.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
…does bolt_moe now follows mlx_lm's NemotronHMoE order and roundings: the picks are summed by weight in the latent ((y * scores).sum(-2) in fp32, one bf16 rounding), the sum is projected back to hidden once by fc2_latent_proj, then the shared expert and the residual are added (bf16(y + shared), then bf16(h + y)). It replaces the per-pick out-projection, which was equal in exact arithmetic but rounded differently, and costs nine times the fc2_latent_proj work. - combine_latent.metal: the latent combine, with add_norm_moe_xs's routed-sum order (a sequential fma over the picks). - add_norm_moe_xs_bolt is now <4096, 512, 1>: one pick of weight 1 (exact under the fma) carries the projected MoE output into add_norm; Scratch gains a ones buffer for it. - tf-metal-tests: combine_latent bit-exact against a host fma reference; 10 of 10 pass. On the published Bolt checkpoint (M3 Max), greedy output is token-exact with mlx_lm 0.32.3 on 7 of 8 comparison prompts (was 5 of 8); the eighth diverges at an exact tie in mlx_lm's own logits. concurrent == solo (x2, x4) and resumed == fresh still hold; decode ~42-43 tok/s. Lightning unchanged (test-golden 0 unexpected). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
…orward) Nemotron had no fidelity tool; this adds one on the qwen35/lfm2 pattern. - tf-nemotron-forward MODEL TOKENS.npy OUT.bin: native teacher-forced bf16 logits for every position, the ids run through the decode kernels fused_rows rows a forward (Lightning or Bolt). - tools/zig/nemotron_fidelity.py: teacher-forced NLL, top-1 agreement and logit RMS of the native logits against mlx-lm run one row at a time, in 32-row chunks and whole, with mlx-lm's own chunked-vs-serial variation as the floor. Every native pass runs before mlx-lm loads, so a 74B checkpoint is never held twice. The long fixture is the running Python's json.decoder and argparse sources (public, deterministic). Bolt on an M3 Max (mlx 0.32.3, mlx-lm 0.32.0), native vs mlx-lm serial, with mlx-lm's own chunk32-vs-serial floor in brackets: fixture tokens top-1 logit RMS NLL native / mlx fibonacci-raw 15 0.933 (1.000) 0.181 (0.444) 2.890 / 2.841 gpu-chat-no-think 35 0.971 (0.971) 0.179 (0.191) 5.715 / 5.724 contribution-guide-512 512 0.922 (0.947) 0.322 (0.262) 3.293 / 3.306 python-stdlib-2048 2048 0.988 (0.987) 0.494 (0.484) 0.362 / 0.364 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
… 600 lines) The branch added 39 lean findings: multi-line comments and docstrings in the new kernels, tests and tools, and layers.zig past 600 lines. Each comment is now one line of at most 120 columns; bolt_moe moves to its own bolt.zig (with upRelu2 made pub for it) and the group-norm pick becomes an inline loop, so layers.zig is 580 lines. No behavior change: the served Bolt output is the same token for token (7 of 8 prompts token-exact with mlx_lm, as before). lean_new: 0 new (503 on the base and on this head); zig fmt --check clean; test-kernel-contract, test, test-golden (0 unexpected) and tf-metal-tests (10 of 10) pass. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
|
Thanks @trickyfalcon. The common lane host and latent combination order are now in place, and your M3 Max receipts help. Please keep it as a draft while you reject top-k above any layer's expert count, add the latent-tensor negatives, and refuse unsupported tensor-unit chips before loading. The M5 path can reach unloaded row pipelines and Lightning-sized kernels. Reuse the shared GPU core for reusable operations and retain RADIO images as a shipping gate. The fidelity table exceeds its quoted RMS floor on two fixtures; please correct that claim and attach pinned KL, per-operation bit/row, Lightning and session-equivalence receipts. Drafting is correctly marked inapplicable without MTP weights. Posted by Ash's AI Assistant Synthena |
…ding - config: a puzzle layer whose top-k exceeds its own expert count is refused (the global check compared against the largest layer only) - weights: a latent projection that is missing, quantized, transposed or not bf16 is refused with its name; tests cover each, and experts stored over hidden - model: Bolt is refused on tensor-unit chips (M5) before any weight is read, since the Builder would tile its projections for Lightning-sized lane kernels - the Bolt checks move to config_bolt.zig (config.zig stays under 600 lines) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
The hand-written route_bolt, dense_bf16 and combine_latent kernels are gone; Bolt now runs the ops/ kernels the prefill set already compiles with MLX's options: - gemv: ops/gemv.metal gains MLX's bm4/bm8 bn1 gemv beside its bn8 one; Bolt's gate, latent projections and bf16 lm_head take MLX's own tiles for each (K, N), a row a z - route: tf_route_topk (mlx_lm group_expert_select) at every Bolt expert count - combine: tf_moe_combine_rows follows col_reduce_small's lane order past 8 picks (8 lanes, pick 8 onto pick 0); at 8 or fewer picks its sums are unchanged tools/zig/bolt_op_cases.py checks each against MLX 0.32.3 on the Bolt checkpoint, every row of a 1-16 row pass equal to its own one-row op and MLX's metallib replay equal too (check_mlx_ops.py --ops bolt_gemv,bolt_route,bolt_combine). The fidelity tool now reports KL(mlx || native) a position and pins the checkpoint and commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
…past 8 picks Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
|
Thanks. All of that is in, and it stays a draft.
Receipts, on the M3 Max with the pinned checkpoint:
I corrected the fidelity claim. Native RMS is above mlx-lm's chunk spread on two fixtures. Switching to the exact ops didn't move that, because the rest of the M1–M4 decode path isn't MLX-exact, and the model spreads one-ULP differences to that size. So the per-op table is the bit claim, and the fidelity table bounds the whole-model effect. Details are in the updated description. |
|
Thanks @trickyfalcon. The per-layer top-k bound, latent-tensor negatives, early tensor-unit refusal and shared-op reuse address our earlier source asks. Please also make Model.load release the Metal device and queue it creates before that refusal; its current error path releases neither. Keep this draft for RADIO images and the full numerical/session gate. Your M3 results are useful; please attach the updated pinned KL/fidelity table and raw operator, row-invariance, Lightning and session receipts. We'll retain the unchanged numerical bars and run the integrated M3 acceptance and config-only M5 refusal/cleanup cells before shipping this backend. Posted by Ash's AI Assistant Synthena |
A config refused after the device opens (Bolt on a tensor-unit chip, a bad shape, a time-step clamp) returned without releasing either. A test counts the device's references across a refused load. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ERqbTvabLDMhhBM4yG3Ydq
|
Thanks. It stays a draft for RADIO and the full gate. The raw receipts are below, all from the M3 Max on the pinned checkpoint. The fix only touches the load error path, so these runs still stand. Per operator and row invariance: check_mlx_ops.py --ops bolt_gemv,bolt_route,bolt_combine (236/236)Pinned KL / fidelity: nemotron_fidelity.py --hash-weights (the model path is shortened to its folder name){
"mlx": "0.32.3",
"mlx_lm": "0.32.0",
"pins": {
"model": "NVIDIA-Nemotron-3.5-Bolt-74B-A9B-MLX-4bit",
"config_sha256": "7279c754382619fd2d849835d78c3137ce27a972efb1eadb86c61585af405985",
"index_sha256": "3b6189b91b55b18edd3be8f512d6b901686b2eb0370d1685266e91e79d01eba0",
"shards": {
"model-00001-of-00009.safetensors": "7e50d79de713d0ff7bc3be41ef52c8cd3c3e35cc3331f03a91b0916d30088639",
"model-00002-of-00009.safetensors": "999d37fa98e8926af495b59250808c4dcb01e7272ade97387aeb5a6ebfd9d78a",
"model-00003-of-00009.safetensors": "c75f5fe7261db2374b8e23cdd2cb1827140799ac13fa536cccbdb9d242be61da",
"model-00004-of-00009.safetensors": "4da61085cc4f0505569412ee72d2c853b819c0f2f2c8ce85470149d8437b2864",
"model-00005-of-00009.safetensors": "6f03752282c67af6b2b2e8e3c66d63a13cbd4cd2c19f8e622f40a0a7c9299caf",
"model-00006-of-00009.safetensors": "6aafcfc9d99618addbecf93b6ef6f1adc8114941e9ca4a0c1ba3cbf54fbb4d9d",
"model-00007-of-00009.safetensors": "87238841e0aad4279e95908268d34e3ed1b5da5de0d04ce99fb4f8c45af67c35",
"model-00008-of-00009.safetensors": "3e94ea58a946e1d54521ef18baa6c73e61066f68588fb2b46726ede3aac24dba",
"model-00009-of-00009.safetensors": "c19f22b58b4c42a739374585623d16addf2de352022b9d8a440ba131ba5ff720"
},
"engine_commit": "ce28820e47c16988f56288ad64ad464e0738184a",
"engine_dirty": true
},
"fixtures": [
{
"fixture": "fibonacci-raw",
"tokens": 15,
"token_sha256": "b7769551532c0ad8b9a1b8124e3bb9e3edcf8cca5d64125c020edb80d824adbb",
"native_nll": 2.8919690871110246,
"mlx_serial_nll": 2.8406177388455425,
"mlx_chunk32_nll": 2.8742868031743534,
"mlx_full_nll": 2.8742868031743534,
"native_vs_serial": {
"top1_agreement": 0.9333333333333333,
"max_abs": 1.5,
"rms": 0.2695379667229362,
"finite": true,
"kl_mean": 0.003982284659141044,
"kl_p99": 0.009924088716917365,
"kl_max": 0.010204187780618004
},
"native_vs_chunk32": {
"top1_agreement": 0.9333333333333333,
"max_abs": 4.125,
"rms": 0.6525186847388184,
"finite": true,
"kl_mean": 0.006040667485176172,
"kl_p99": 0.01711581182153379,
"kl_max": 0.01742250315447551
},
"mlx_chunk32_vs_serial": {
"top1_agreement": 1.0,
"max_abs": 3.125,
"rms": 0.4436320970441915,
"finite": true,
"kl_mean": 0.006796612407219586,
"kl_p99": 0.019070950104862386,
"kl_max": 0.01995933719745038
},
"mlx_full_vs_serial": {
"top1_agreement": 1.0,
"max_abs": 3.125,
"rms": 0.4436320970441915,
"finite": true,
"kl_mean": 0.006796612407219586,
"kl_p99": 0.019070950104862386,
"kl_max": 0.01995933719745038
}
},
{
"fixture": "gpu-chat-no-think",
"tokens": 35,
"token_sha256": "34c41fc50825b5fd887436ed1f9f2ae69d710777a221f94b1d5cbb5717737fe4",
"native_nll": 5.716686995554169,
"mlx_serial_nll": 5.724390312130104,
"mlx_chunk32_nll": 5.765810753377128,
"mlx_full_nll": 5.717642529948139,
"native_vs_serial": {
"top1_agreement": 0.9428571428571428,
"max_abs": 1.25,
"rms": 0.1893557441630768,
"finite": true,
"kl_mean": 0.010642203432166349,
"kl_p99": 0.08333138989714994,
"kl_max": 0.10881155142696898
},
"native_vs_chunk32": {
"top1_agreement": 0.9714285714285714,
"max_abs": 1.375,
"rms": 0.21460021831939105,
"finite": true,
"kl_mean": 0.00901477917118718,
"kl_p99": 0.04355610256285816,
"kl_max": 0.047731224040055864
},
"mlx_chunk32_vs_serial": {
"top1_agreement": 0.9714285714285714,
"max_abs": 1.75,
"rms": 0.19092943749106223,
"finite": true,
"kl_mean": 0.007755403077574087,
"kl_p99": 0.030876112146067048,
"kl_max": 0.03177423913459955
},
"mlx_full_vs_serial": {
"top1_agreement": 0.9714285714285714,
"max_abs": 1.75,
"rms": 0.23128675328521206,
"finite": true,
"kl_mean": 0.0072931219236618315,
"kl_p99": 0.031277037836049904,
"kl_max": 0.03227898497237537
}
},
{
"fixture": "public-contribution-guide-512",
"tokens": 512,
"token_sha256": "b0206242758037798cce991293f0f600a7c5113d6b652705b43ada0bd2b192ce",
"native_nll": 3.2981265319963677,
"mlx_serial_nll": 3.306361377972774,
"mlx_chunk32_nll": 3.2939761279482083,
"mlx_full_nll": 3.3011254621412025,
"native_vs_serial": {
"top1_agreement": 0.91796875,
"max_abs": 8.25,
"rms": 0.4013358564352756,
"finite": true,
"kl_mean": 0.010945506139855322,
"kl_p99": 0.04042149809126025,
"kl_max": 0.09108781051450346
},
"native_vs_chunk32": {
"top1_agreement": 0.921875,
"max_abs": 9.4375,
"rms": 0.39862855150561066,
"finite": true,
"kl_mean": 0.010753584871592985,
"kl_p99": 0.047534459930563815,
"kl_max": 0.1400134702895854
},
"mlx_chunk32_vs_serial": {
"top1_agreement": 0.947265625,
"max_abs": 4.25,
"rms": 0.26151797632237317,
"finite": true,
"kl_mean": 0.009838347827565704,
"kl_p99": 0.03959047257870294,
"kl_max": 0.05834875115422212
},
"mlx_full_vs_serial": {
"top1_agreement": 0.9453125,
"max_abs": 8.8125,
"rms": 0.4033499529019144,
"finite": true,
"kl_mean": 0.010674325096614975,
"kl_p99": 0.042492825393599755,
"kl_max": 0.06645138175069834
}
},
{
"fixture": "python-stdlib-2048",
"tokens": 2048,
"token_sha256": "5a42facb3a83f6c531383a19c57a666bf245328afc4138b97029ad2be1afadab",
"native_nll": 0.36548583462730677,
"mlx_serial_nll": 0.36374177519935724,
"mlx_chunk32_nll": 0.36367116087373574,
"mlx_full_nll": 0.36616538895453143,
"native_vs_serial": {
"top1_agreement": 0.986328125,
"max_abs": 11.0625,
"rms": 0.557423203819268,
"finite": true,
"kl_mean": 0.005274513941993687,
"kl_p99": 0.06450120232168344,
"kl_max": 0.875234221777027
},
"native_vs_chunk32": {
"top1_agreement": 0.98974609375,
"max_abs": 10.5,
"rms": 0.5284347706154673,
"finite": true,
"kl_mean": 0.004072006140364548,
"kl_p99": 0.0500922697477258,
"kl_max": 0.11082624491880542
},
"mlx_chunk32_vs_serial": {
"top1_agreement": 0.9873046875,
"max_abs": 11.875,
"rms": 0.4840688789466054,
"finite": true,
"kl_mean": 0.0051351756370919535,
"kl_p99": 0.06548356509049968,
"kl_max": 0.6026136509687315
},
"mlx_full_vs_serial": {
"top1_agreement": 0.98828125,
"max_abs": 11.1875,
"rms": 0.5309294850846997,
"finite": true,
"kl_mean": 0.005368460118554758,
"kl_p99": 0.06725235402622899,
"kl_max": 0.5935126875951473
}
}
]
}concurrent == solo: xpu_concurrent_receipt.py, 2 and 4 streamsresumed == fresh: xpu_resume_receipt.py, cached server then --prompt-cache-gib 0# cached
{"sampling": "greedy", "turn1": {"sha": "8351ea926665", "prompt": 1508, "cached": 0, "prefill_s": 14.557032}, "turn2": {"sha": "e54437d714b4", "prompt": 1580, "cached": 1502, "prefill_s": 0.863641417}}
{"sampling": "seeded", "turn1": {"sha": "cc1f8a2a3a78", "prompt": 1508, "cached": 1502, "prefill_s": 0.086221167}, "turn2": {"sha": "235fa147bcc0", "prompt": 1606, "cached": 1502, "prefill_s": 1.138095875}}
# fresh
{"sampling": "greedy", "turn1": {"sha": "8351ea926665", "prompt": 1508, "cached": 0, "prefill_s": 15.623639125}, "turn2": {"sha": "e54437d714b4", "prompt": 1580, "cached": 0, "prefill_s": 13.90681975}}
{"sampling": "seeded", "turn1": {"sha": "cc1f8a2a3a78", "prompt": 1508, "cached": 0, "prefill_s": 12.926058917}, "turn2": {"sha": "235fa147bcc0", "prompt": 1606, "cached": 0, "prefill_s": 13.742770125}}Lightning 30B-A3B (d9d758fb): main vs this branch, token_sha per request (32/32 equal)# main
{
"p0-greedy-chat": "5ae8a877bf7a",
"p0-greedy-raw": "434cb488a3e9",
"p0-seeded-chat": "ca84ce37feaf",
"p0-seeded-raw": "154bb8b237f1",
"p1-greedy-chat": "b18a343f6d01",
"p1-greedy-raw": "2ca4bcfac5c6",
"p1-seeded-chat": "75ad7305d004",
"p1-seeded-raw": "994dfa915cf0",
"p2-greedy-chat": "2f212f7460d7",
"p2-greedy-raw": "6ce81e25bb38",
"p2-seeded-chat": "4aac7cb43ab7",
"p2-seeded-raw": "11ccbd49c519",
"p3-greedy-chat": "ebd63876f10b",
"p3-greedy-raw": "866008e96b3a",
"p3-seeded-chat": "bb7b0f3e6107",
"p3-seeded-raw": "ae3c33a23e79",
"p4-greedy-chat": "99f683d98b2a",
"p4-greedy-raw": "eb0c21a3d01c",
"p4-seeded-chat": "80d51dfae249",
"p4-seeded-raw": "341615eefe95",
"p5-greedy-chat": "15ffafbac276",
"p5-greedy-raw": "878d1d439cda",
"p5-seeded-chat": "6a59107e2c89",
"p5-seeded-raw": "bbf7d21d4f9e",
"p6-greedy-chat": "689eaf56d4b3",
"p6-greedy-raw": "b269199162ce",
"p6-seeded-chat": "92fbbb26fe47",
"p6-seeded-raw": "4fcf8b6c20b4",
"p7-greedy-chat": "104a1428c2f1",
"p7-greedy-raw": "fbf60b9de08e",
"p7-seeded-chat": "e8ce18600195",
"p7-seeded-raw": "3f7c0c8cec66"
}
# this branch
{
"p0-greedy-chat": "5ae8a877bf7a",
"p0-greedy-raw": "434cb488a3e9",
"p0-seeded-chat": "ca84ce37feaf",
"p0-seeded-raw": "154bb8b237f1",
"p1-greedy-chat": "b18a343f6d01",
"p1-greedy-raw": "2ca4bcfac5c6",
"p1-seeded-chat": "75ad7305d004",
"p1-seeded-raw": "994dfa915cf0",
"p2-greedy-chat": "2f212f7460d7",
"p2-greedy-raw": "6ce81e25bb38",
"p2-seeded-chat": "4aac7cb43ab7",
"p2-seeded-raw": "11ccbd49c519",
"p3-greedy-chat": "ebd63876f10b",
"p3-greedy-raw": "866008e96b3a",
"p3-seeded-chat": "bb7b0f3e6107",
"p3-seeded-raw": "ae3c33a23e79",
"p4-greedy-chat": "99f683d98b2a",
"p4-greedy-raw": "eb0c21a3d01c",
"p4-seeded-chat": "80d51dfae249",
"p4-seeded-raw": "341615eefe95",
"p5-greedy-chat": "15ffafbac276",
"p5-greedy-raw": "878d1d439cda",
"p5-seeded-chat": "6a59107e2c89",
"p5-seeded-raw": "bbf7d21d4f9e",
"p6-greedy-chat": "689eaf56d4b3",
"p6-greedy-raw": "b269199162ce",
"p6-seeded-chat": "92fbbb26fe47",
"p6-seeded-raw": "4fcf8b6c20b4",
"p7-greedy-chat": "104a1428c2f1",
"p7-greedy-raw": "fbf60b9de08e",
"p7-seeded-chat": "e8ce18600195",
"p7-seeded-raw": "3f7c0c8cec66"
} |
Draft, opened early per #962 so the design can be steered before the kernels harden. The native Metal backend for
nemotron_h_puzzle(Nemotron 3.5 Bolt 74B-A9B), on the current release (v1.1.0), which already parses the Puzzle config and serves it on CUDA; this adds the Metal side. It serves the published checkpoint end to end (bytebunker/NVIDIA-Nemotron-3.5-Bolt-74B-A9B-MLX-4bit): on an M3 Max (128 GB) it loads in ~6 s and decodes at ~40 tok/s.Receipts (M3 Max, applegpu_g15, the published MLX-4bit build)
Pinned checkpoint:
config.jsonsha2567279c7543826…, index3b6189b91b55…, all nine shards hashed in the fidelity report (--hash-weights). Reference: MLX 0.32.3, mlx-lm 0.32.0. Engine: this branch.check_mlx_ops.py --ops bolt_gemv,bolt_route,bolt_combine)group_expert_selectat k 9, and the latent combine; MLX's own metallib kernel replayed at each shape agrees tooxpu_concurrent_receipt.py)xpu_resume_receipt.py, CONTRIBUTING.md, two turns)--prompt-cache-gib 0, greedy and seedednum_nextn_predict_layers: 0)test-golden0 unexpected; and Lightning 30B-A3B MLX-4bit (d9d758fb) served bymainand by this branch on the M3 Max: 32/32 requests the same token_sha (8 prompts × greedy/seeded × chat/raw)zig build test,test-bolt-host,test-kernel-contract,tf-metal-tests, leanTeacher-forced fidelity vs mlx-lm one token at a time (
tools/zig/nemotron_fidelity.py, KL(mlx ‖ native) per position in nats). mlx-lm's own 32-chunk run against its serial run is given as a reference spread, not a bound:A correction to the earlier version of this table: native RMS is above mlx-lm's chunk spread on two fixtures (contribution-guide and python-stdlib), so "inside the floor" did not hold. Moving the gate, latent projections, head, route and combine onto MLX-exact kernels did not lower these numbers. The previous build and this one differ from each other by RMS 0.17–0.51, the same size as either differs from mlx-lm. The rest of the decode path on M1–M4 (the row kernels for the 4-bit projections, Mamba and attention) is not MLX-exact, and the model turns any one-ULP difference into this spread. The per-op table is the bit-level claim; the fidelity table bounds the whole-model effect.
What's in it
fc1/fc2_latent_projand the unquantized bf16lm_headread in place.bolt.zig), in mlx_lm's NemotronHMoE order: gate gemv →tf_route_topk→fc1_latent_proj→ routed experts →tf_moe_combine_rows→fc2_latent_proj→+ shared→+ residual. The gemvs, route and combine come fromops/and the prefill set compiles them with MLX's options; the hand-written route, dense and combine kernels are gone.ops/gemv.metalgains MLX's bm4/bm8 bn1 gemv (tf_gemv_bf16_bm{4,8}_tm4) beside its bn8 one.tf_moe_combinenow followscol_reduce_small's lane order past 8 picks (8 lanes, pick 8 onto pick 0); at 8 picks or fewer its sums are unchanged.bolt_op_cases.py(the per-op cases),tf-nemotron-forwardandnemotron_fidelity.py(teacher-forced logits, KL, pins).Not yet / next
prefill_mixers) for the latent MoE; prompts run through the decode kernels 16 rows at a time.Deferred, with their gates kept