nemotron-3.5-lightning-mtp-head-mlx-4bit

Nemotron-3.5-Lightning-30B-A3B — 4-bit MLX multi-token-prediction head

This repo holds a single 772 MB weights file: the multi-token-prediction (MTP) head of NVIDIA-Nemotron-3.5-Lightning-30B-A3B, converted to MLX and quantised to 4-bit. It exists because mlx-community/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-4bit does not carry the MTP head. Pairing that 4-bit target with this file gives you a complete MTP speculative-decoding setup on Apple Silicon; nothing else is required.

What is inside

Extracted from the base release (mtp.layers.0 and mtp.layers.1) and stored here:

• the draft-input transforms — enorm, hnorm, eh_proj
• one attention block and one MoE block
• the head's own final_norm

The target model's embedding table and lm_head are not included; the head borrows both from the checkpoint it drafts for.

How it drafts

Given token t, the head predicts token t+2 from the pair (next token, hidden state):

x = eh_proj(concat(enorm(embed(tok_{t+1})), hnorm(h_t)))
x = attention(x) -> MoE(x) -> final_norm(x)
logits = lm_head(x)        # lm_head and embedding belong to the target model

h_t is the target's residual stream taken just before its final norm. Since the target model verifies every draft the head proposes, quantisation error in the head cannot corrupt the output — it only moves the acceptance rate up or down.

Using it

Run MTP speculative decoding (two draft tokens per step) against the mlx-community 4-bit target with sevren-ai/mlx-heimr-sparse-down:

uv run heimr-sparse-down --mtp 2

Numbers

On an Apple M5 Pro, sparse decoding with the head attached (--mtp 2) reaches 150–164 tok/s on a code prompt, against 105 tok/s for stock mlx_lm. Full measurements live in the mlx-heimr-sparse-down repo's README.

Tensor inventory

Provenance and quantisation details

Source weights: mtp.* tensors (bf16) from nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4. Quantisation: affine 4-bit, 64-wide groups, applied with mlx.nn.quantize in round-to-nearest mode with no calibration data. config.json in this repo mirrors those settings.