Nemotron Puzzle / LatentMoE on Metal (e.g. Bolt-74B-A9B) #962

Activity

1. trickyfalcon commented on Oct 16, 2026

   

   trickyfalcon

   on Oct 16, 2026

   Author

   More actions

   Confirmed on an M3 Max (128 GB) with TensorFold 1.0.12 (Homebrew ashhart/tensorfold/tensorfold): serving Bolt's nemotron_h_puzzle MLX checkpoint refuses immediately at config —

   the native engine has no backend for nemotron_h_puzzle checkpoints yet; the Python engine 0.6.6 may serve them
   

   So the gap is at the backend level (not a quant/shape mismatch), and it fails fast — before any weights load. capabilities --json on Metal names nemotron but not puzzle/omni.

   For contrast, the same checkpoint does serve on Metal via the 0.6.x MLX path at ~45 tok/s single / ~51 at 4 streams (write-up linked in the issue). Happy to help validate a native Metal nemotron_h_puzzle backend on Apple Silicon if one gets started.
2. ashhart commented on Oct 17, 2026

   

   ashhart

   on Oct 17, 2026

   Owner

   More actions

   Thank you, @trickyfalcon, the 1.0.12 refusal and your M3 Max proof distinguish the missing backend from a malformed checkpoint.

   We checked the released Metal registry at 833bbfc: it has Lightning, but no puzzle/omni backend. The current Bolt consolidation is CUDA work, so it does not fill that Metal gap.

   Please share the exact checkpoint revision and pinned PoC changes as design/reference evidence for a native Metal proposal. We will keep it with the shared latent-MoE and per-layer expert work, with drafts and owned images included where the checkpoint supports them. The Python serving line remains frozen; your measured behavior is still useful for the native port.

   Posted by Ash's AI Assistant Synthena
3. trickyfalcon commented on Oct 17, 2026

   

   trickyfalcon

   on Oct 17, 2026

   Author

   More actions

   Thanks @ashhart — understood, the Python line stays frozen; posting this purely as design/reference for a native Metal nemotron_h_puzzle backend.

   Checkpoint (exact revision). We published a first-party MLX build to serve as a concrete, pullable reference:

   • bytebunker/NVIDIA-Nemotron-3.5-Bolt-74B-A9B-MLX-4bit @ 660c6bb50620682f9c79fc75187ea5c2077ae767
   • MLX affine 4-bit, group_size 64; embeddings 8-bit; router/gate, fc1_latent_proj/fc2_latent_proj, lm_head, norms and conv1d kept full-precision; ~4.674 bpw.
   • Quantized by us from nvidia/NVIDIA-Nemotron-3.5-Bolt-74B-A9B-BF16; numerically equivalent to the community rapid-mlx build (identical structure and zero-points, per-group scales within 1e-3, dequantized weights within a single quant level), so either serves as a reference.

   The load-bearing point: this layout is affine, not NVFP4/ModelOpt. Metal already reads affine — so nemotron_h_puzzle-on-Metal is a backend/dataflow problem, not a new-quant problem.

   PoC changes (pinned, reference). The 0.6.x Metal proof is two additive files + one guarded fused-path patch, all here:
   → https://github.com/bytebunkerlabs/dgx-spark-serve/tree/main/contrib/tensorfold-nemotron-h-puzzle

   • mlx_lm/models/nemotron_h_puzzle.py — builds each block from its own per-layer args and stacks experts at that layer's count (the heterogeneity the uniform loader trips on); top-k is uniform (9), surfaced globally for the fused kernel.
   • tensorfold/families/nemotron_h_puzzle/ — registers the model_type, reuses the nemotron_h loader/kernels.
   • fused-latentmoe-{kernels,model}.patch — teaches the fused MoE the three latent steps: down-project to the 1024 latent before the experts, up-project after the weighted sum, shared expert stays on hidden. Guarded on moe_latent_size, so the Lightning path is byte-for-byte unchanged.

   Design shape (lines up with your latent-MoE + per-layer expert work). The latent dataflow mirrors cuda_latent.zig run() — that's the blueprint for the Metal path. Two pieces are new to the Metal side: (1) routed experts run in the 1024 latent (fc1/fc2 around them, shared expert on hidden), and (2) per-layer heterogeneous expert counts (240/304/368/432/496) with top-k uniform at 9 — the grouped kernels stay count-agnostic, but router/route/add_norm bake E/TOPK, so those need per-shape variants. Text backbone only on our side (vision tower dropped); glad to extend where the omni checkpoint supports it.

   Happy to shape this however you prefer — an issue-level design doc first (gate agreed here), then I fork from the release branch and build to the bit-exactness bar. What format works best for you?
4. ashhart commented on Oct 17, 2026

   

   ashhart

   on Oct 17, 2026

   Owner

   More actions

   Thank you @trickyfalcon! A pinned affine checkpoint settles the format question: Metal already reads affine 4-bit, so this is backend work, as you say. The native shape we'd take: per-layer expert counts in the Nemotron loader first, then the latent down- and up-projection around the shared MoE, with the Lightning path byte-for-byte unchanged. Both build on the shared latent-MoE work in the Bolt consolidation (#897 with #938). We'll read your proof of concept for the design and write the Zig natively, or review a native PR from you. The gate: drafted == plain, resumed == fresh and concurrent == solo on an M-series Mac, against your published build.

   Posted by Ash's AI Assistant Synthena
5. trickyfalcon commented on Oct 18, 2026

   

   trickyfalcon

   on Oct 18, 2026

   Author

   More actions

   Thanks @ashhart — that's exactly the split I was hoping for. I'll take the native PR path: write the Zig/Metal nemotron_h_puzzle backend and bring it to the gate you set.

   Confirming the shape so we're aligned before the kernel time goes in:

   1. Per-layer expert counts in the Nemotron loader first — parse the Puzzle block_configs into per-layer n_routed_experts (240/304/368/432/496) and moe_intermediate_size, following the existing per-layer arrays; num_experts_per_tok stays uniform at 9.
   2. Latent down/up-projection around the shared MoE — fc1_latent_proj (4096→1024) after route, fc2_latent_proj (1024→4096) after the weighted combine; routed experts run in the 1024 latent, the shared expert stays in hidden. I'll mirror the CUDA dataflow from the shared latent-MoE work in Nemotron 3.5 Bolt 74B-A9B on CUDA: puzzle layers, LatentMoE, ggml IQ2_XXS/Q8_0 weights, image input, host prompt states (on #853) #897/nemotron: serve NVIDIA Nemotron 3.5 Bolt (nemotron_h_puzzle) on CUDA #938.
   3. Lightning path byte-for-byte unchanged — everything latent-guarded on moe_latent_size, new kernels registered with their M1–M5 contract limits, test-golden untouched.

   The gate I'll hit, on an M-series Mac, against bytebunker/NVIDIA-Nemotron-3.5-Bolt-74B-A9B-MLX-4bit:

   • drafted == plain token-for-token, resumed == fresh, concurrent (2/4) == solo;
   • per-op bit-equality and per-kernel row-invariance at the real heterogeneous shapes;
   • teacher-forced KL vs mlx_lm within your bar;
   • the kernel-contract SUMMARY, gates 1–4, with fp32 accumulation in the latent.

   I'll branch from the current release (tell me if you'd rather I base it elsewhere), keep it text-backbone-scoped to start, and defer the MTP draft head until the main path is bit-exact. I'll open a draft PR early so you can steer the design before the kernels harden.