NVIDIA Releases Nemotron-Labs-3-Puzzle-75B-A9B: A Compressed Hybrid MoE LLM Delivering 2.03x Server Throughput at Matched User Throughput


Large hybrid MoE models like Nemotron-3-Super are accurate but expensive to serve. Their active parameters, KV cache, and Mamba state cap how many users a node can hold at a given per-user token rate. NVIDIA AI team has released Nemotron-Labs-3-Puzzle-75B-A9B, a compressed variant of Nemotron-3-Super. The parent model has 120.7B total and 12.8B active parameters. The compressed model has 75.3B total and 9.3B active parameters.

The deployment target was fixed before the architecture search began. Target one was 2x server throughput at 100 tokens per second per user. Target two was 8 concurrent 1M-token requests on a single H100. Three checkpoints on Hugging Face: BF16, FP8, and NVFP4.

TL;DR

  • 120.7B/12.8B active compresses to 75.3B/9.3B active, with the 88-block hybrid layout preserved.
  • 8xB200 total throughput rises 1.60x to 2.14x over Super at matched NVFP4 and matched user throughput.
  • Single-H100 1M-token concurrency goes 1 to 8, driven by a 70 GB to 44.5 GB weight drop.
  • Iterative Puzzle beats single-step Puzzle by 0.57 average points at the same compression target.
  • Arena-Hard-V2 (-4.2) and SWE-Bench (-2.6) are the real costs; RULER and AA-LCR barely move.

Nemotron-Labs-3-Puzzle-75B-A9B

Nemotron-3-Super is a hybrid Mamba-Transformer MoE model. Puzzle-75B-A9B preserves the parent’s block layout exactly. It has 88 blocks: 40 Mamba, 40 MoE, and 8 attention blocks.

What changed is capacity inside those blocks:

QuantitySuperPuzzle-75B-A9BRatio
Total parameters120.7B75.3B62.4%
Active parameters12.8B9.3B73.1%
Mamba SSM state size1289675%
MoE routed expert intermediate size26881280-2688Mean 59.9%
Activated routed experts per token224-18Mean 50%
Active routed expert capacity (relative)100%8.7%-62.3%Mean 30.9%

The number of routed experts, the shared expert size, and the MoE latent size are unchanged. Attention layers were left untouched. The proposed research’s stated reason is that Nemotron-3-Super is already very KV-cache efficient. Mamba layers were pruned uniformly, because inference frameworks do not support a different SSM state size per layer.

https://arxiv.org/pdf/2607.04371

The result is not a uniformly scaled-down teacher. The above figure shows the allocation across depth. Puzzle preserved capacity in selected middle and late layers, and cut hard elsewhere.

Benchmark and Performance

The below table reports Pareto-optimal total throughput on a single 8xB200 node, with single-step decoding.

Scenario (in/out)UT floorSuper (tok/s)Puzzle-75B-A9B (tok/s)Boost
50K / 2K>= 1005,1288,2101.60x
50K / 2K>= 1253,7846,4121.69x
50K / 2K>= 1502,5324,5231.79x
8K / 64K>= 10020,93942,6012.03x
8K / 64K>= 12513,07427,9182.14x
8K / 64K>= 1508,52218,0472.12x

Both models were served at matched NVFP4 weights, FP8 KV cache, and FP16 Mamba state. The gap therefore reflects compression, not a change in numeric format. The prefill-heavy 50K/2K regime gains least. The decode-heavy 8K/64K regime gains most.

On a single 8xH100 node at UT = 100, the gains are smaller. They are 1.91x on 50K/2K and 1.82x on 8K/64K. Both models there use FP8 weights, FP8 KV cache, and FP32 Mamba state.

On a single H100 at 1M context, the binding constraint flips from compute to memory. Super’s NVFP4 weights occupy about 70 GB of the 80 GB HBM budget. Each 1M-token request adds about 4 GB of KV cache. Effective concurrency is therefore 1.

Puzzle-75B-A9B’s NVFP4 weights occupy about 44.5 GB. Attention layout is unchanged, so per-request KV cost is unchanged. Concurrency at 1M rises to 8. Aggregate decode throughput at that concurrency is roughly 4x Super’s single-request throughput. Prefill of a 990K-token prompt is about 1.2x faster.

How Iterative Puzzle Works

Puzzle is a decomposed neural architecture search framework, implemented here as Puzzletron. It defines a discrete search space of alternative layer implementations. Each alternative gets a quality score. A mixed-integer program then selects one alternative per layer under a deployment constraint.

Three pruning techniques form the search space:

  • Intermediate channel pruning: Channels inside each routed expert are ranked by contribution to the expert’s output. All experts within one MoE layer are pruned to a uniform size, for kernel compatibility.
  • Top-k reduction: The number of experts a token is routed to varies per layer, up to the parent’s k=22.
  • Mamba SSM pruning: The SSM state size drops from 128 to 96 channels.

The SSM result is measured. Dropping 128 channels to 96 speeds the SSM kernel 1.2x to 1.3x during decode. This holds at batch sizes between 8 and 512. Channels were ranked by estimated contribution to the Mamba layer output. The estimate averaged over 67M tokens of validation data. Appendix A shows this beats random channel selection under aggressive pruning.

The original formulation assumes replacement quality impacts are approximately additive. Each candidate block is scored inside the unmodified parent. That ignores higher-order interactions between replacements.

Iterative Puzzle alternates bounded compression with short knowledge distillation recovery. It builds a sequence M0, M1, … MR instead of jumping to the target. Scores are recomputed against the current compressed model, not the original parent.

Three stages were used:

  1. MoE weights to 75% of teacher capacity, Mamba SSM state to 75%. Healed for 24B tokens.
  2. MoE weights to 60% of teacher capacity. Healed for 43.2B tokens.
  3. Activated routed-expert budget to 50%, allocated heterogeneously. Healed for 52.8B tokens.
https://arxiv.org/pdf/2607.04371

The above table compares this against a single-step Puzzle baseline at the same target. The three-step procedure averages 69.05 across ten benchmarks, against 68.48. Gains appear on MMLU-Pro, GPQA, HLE, AA-LCR, LiveCodeBench, SciCode, and RULER-256K. IFBench-Instruction fell 0.2 points and IFBench-Prompt fell 0.5.

Recovery: Distillation, RL, and Verbosity

Knowledge distillation ran on 30% pretraining data and 70% SFT data from Nemotron-3-Nano. During the Puzzle phase, KD used a 32K sequence length. Recovery then trained at 128K, and scaled to 512K. The budget was up to 100B tokens, with a 16M-token global batch, in Megatron-LM.

RL post-training adopted Stage 2 of the Nemotron-3-Super RL pipeline, focused on software engineering. Phase 2.1 did single-step tool-use comparison. Phase 2.2 moved to end-to-end sandbox RL, where agents run up to 200 turns. Both phases used a KL penalty of 0. The team swept learning rates, then averaged the resulting weights.

https://arxiv.org/pdf/2607.04371

The above Figure 4 shows what each stage contributed. Short-context KD recovers most categories to over 97% of Nemotron-3-Super. Long-context KD then lifts long-input and long-generation benchmarks specifically. The research team states that RL’s impact in these experiments was small.

Verbosity is the quiet detail. After the last Puzzle iteration, the model generated 132% of Super’s token count. That fell to 99% after the full recovery pipeline.

Deployment: Quantization and Multi-Token Prediction

Two post-training quantization recipes were produced: FP8 W8A8 targets Hopper and NVFP4 W4A4 targets Blackwell.

ComponentBF16 baselineFP8 checkpointNVFP4 checkpoint
Sparse and shared MoE GEMMsBF16FP8NVFP4
Mamba GEMMsBF16FP8FP8
Mamba SSM cacheFP32FP32FP16+SR
KV cacheFP8FP8FP8
RouterFP32FP32FP32
Attention QKV/output, MoE latent projections, LM headBF16BF16BF16

Both recipes calibrated on 256 post-training SFT samples. NVFP4 used max calibration, not the AutoQuantize sensitivity search used for Super. The resulting checkpoint is slightly more aggressively quantized, and performed similarly.

NVFP4 is not natively supported on Hopper. It is still used for the 1M-context H100 target, because HBM capacity binds there.

Puzzle-75B-A9B inherits a shared MTP head from Super. Parameters are shared across MTP steps, so one head applies recursively at inference. Transferring Super’s trained head directly gave similar acceptance lengths.

The research team then identifies a training-inference mismatch. Teacher-forced MTP training feeds the full shifted hidden-state sequence. Autoregressive drafting instead feeds a mixture of target-model and MTP-generated hidden states. Acceptance rates fall at deeper draft positions.

Continued training on the transferred head addresses this. On SPEED-Bench at draft length 7, average acceptance length rose from 3.45 to 4.34. That is roughly 25% to 30%, concentrated at later draft positions. Unlike Super, the NVFP4 checkpoint barely degrades: 4.31 against 4.34.

Where Compression Helps and Where It Hurts

Benchmark (BF16)SuperPuzzle-75B-A9BDelta
MMLU-Pro83.882.4-1.4
AIME25 (no tools)92.289.7-2.5
GPQA (no tools)80.578.6-1.9
LiveCodeBench82.181.1-1.0
SciCode (subtask)42.340.6-1.7
SWE-Bench (OpenHands)59.556.9-2.6
Arena-Hard-V272.868.6-4.2
AA-LCR56.856.9+0.1
RULER 1M93.992.2-1.7
MMLU-ProX79.577.5-2.0

The research paper’s own summary is that instruction-following and agentic evaluations lose most. Arena-Hard-V2 is the worst case, at -4.2 points. RULER stays within roughly 1 to 2 points at 256K, 512K, and 1M.

Three BF16 results do not regress. AA-LCR gains 0.1, Scale AI Multi-Challenge ties at 56.6, and TauBench Telecom gains 0.4.

NVFP4 costs little on top of compression. On RULER 1M the NVFP4 checkpoint scores 93.2, above BF16’s 92.2. HLE is the clearest NVFP4 cost, dropping from 16.5 to 15.7. FP8 results sit in Appendix E, and track BF16 closely. SWE-Bench is not reported for the FP8 checkpoint.

Use Cases

  • Ultra-long-context RAG on one GPU: A document analysis service at 1M context goes from 1 concurrent request to 8. Aggregate decode throughput at that concurrency is roughly 4x.
  • Interactive coding assistants: At UT >= 100 tok/s in the 8K/64K regime, one node serves 2.03x the tokens. Adjusted for verbosity, that is 2.16x the completed requests per minute.
  • Prefill-heavy document pipelines: The 50K/2K regime gains only 1.60x. Compression helps less when prompt processing dominates compute.
  • Agentic SWE loops: Check the 2.6-point SWE-Bench gap against your task mix. RL recovery targeted this capability, and only partly restored it.

Deployment Explorer





Source link

  • Related Posts

    Datalab Lift vs the Field: How a 9B Schema-First Extractor Compares with NuExtract3, LlamaExtract, Marker, and Docling

    Datalab’s Lift is a focused document extraction tool with a specific promise: give it a PDF or image plus a JSON Schema, and it returns schema-shaped JSON directly. Instead of…

    Google AI Studio Adds Import from GitHub to Build a Deployable App

    Google AI Studio is rolling out an ‘import from GitHub’ feature inside its Build mode. It takes a repo and transforms it into a runtime-compatible…

    Leave a Reply

    Your email address will not be published. Required fields are marked *