Let's optimize Nemotron 3.5 Lightning - a worklog

I came across NVIDIA’s Nemotron 3.5 Lightning model recently. I rented a RTX PRO 6000 box, thanks to Huggingface and Modal for their GPU credits! I had Claude Fable. So yeah, spent my weekend on squeezing as much as possible to serve this model. What you’ll be reading is a worklog of how I approached this optimization setup. Please feel free to reach out to me if you have any questions / feedback.

TL;DR worklog numbers:

Note:

• “lossless” means in speculative decoding every draft token is verified against the full model, so the output distribution is unchanged and verified.

Grey is the official BF16 baseline. Blue is our optimization efforts. Red is the naive quantize-everything run, which failed the quality gate.

This post is about a single request: one prompt, one RTX PRO 6000, lowest TTFT and TPOT I can get. Concurrency is a different problem (this checkpoint already does ~4,578 tokens/sec in aggregate at concurrency 256) and that’s Part 2.

Hardware: one RTX PRO 6000 Blackwell Server Edition (sm120, 188 SMs, 1.60 TB/s memory bandwidth).

Software: vLLM 0.28.0rc2, plus a few custom changes.

Three terms: TPOT (time per output token) is decode speed. TTFT (time to first token) is prefill speed. Throughput is tokens/sec: per query in this post, aggregate across all concurrent requests in Part 2. TPOT and TTFT are bounded by different parts of the hardware, which is most of what this post is about.

Step 0: Serving BF16 checkpoint

Stock vLLM:

vllm serve nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 \
  --max-model-len 16384 \
  --max-num-seqs 256 \
  --tensor-parallel-size 1

--max-num-seqs 256 is not tuning. This is a 23-layer Mamba hybrid, every decode sequence holds a Mamba state block, and vLLM’s default of 1024 doesn’t fit and kills CUDA graph capture. That cost me my first hour!

TTFT: 91.5 ms, TPOT: 5.58 ms

One BF16 decode step. The kernels are closely packed with no gaps, so there is no idle backlogs to win. The single widest kernel is the BF16 lm_head, the gemvx one (the bottom one): 479 µs per call. lm_head turns the last hidden state into one token’s logits.

179 tokens/sec for a single request. Not bad. Where is the ceiling?

30B params × 2 bytes is 60 GB, but this is a MoE: only 6 of 128 experts fire per token, so the active set is ~3B params. The dense parts (Mamba projections, shared experts, attention, the output head) are read on every step regardless, which brings the bytes actually touched per token to about 3.5B params, or ~7 GB in BF16. If decode did nothing but stream those weights at peak bandwidth:

But we’ve achieved 179 tok/s, so

of the card’s memory bandwidth, before any optimization (I mean vLLM is already optimized). Now what? Is there a way to bring down the 7 GB? Yes you’re right, it’s quantization! (Batching would amortize the reads across requests, but that’s not the goal here.)

One more number to keep: TTFT is 91.5 ms, and it is about to get worse.

Step 1: W4A16 - quantize the weights

The obvious first move: 4-bit weights, BF16 activations. The format is NVIDIA’s NVFP4: 4 bits per weight, with group_size = 16. I built the checkpoint with NVIDIA’s ModelOpt repo. Decode now reads 4× fewer weight bytes per token.

TPOT: 3.07 ms, 1.8× over BF16, about what the bytes math predicts. Weight-only quantization is a good deal and most people should stop here.

TTFT went the other way: 91.5 → 153.3 ms, prefill 67% slower. Prefill pushes the whole prompt (~1k tokens) through the model at once, so the op is compute-bound. The Marlin kernel stores 4-bit weights but dequantizes them to 16-bit inside the kernel and does the math in BF16. At decode this dequant can be overlapped with memory read making the arithmetic essentially free. At prefill it is extra work on the hot path! And we’re not using the native FP4 cores, that’s bad.

One MoE layer at prefill. The two Marlin GEMMs: 396.8 and 338.5 µs to push 1000 tokens through. The same two calls take ~20 µs each at decode. This is the entire TTFT regression.

That asymmetry is the rest of the post. Blackwell has native FP4 tensor cores, up to 4× BF16 throughput, but they need 4-bit activations too. That’s W4A4.

Step 2: W4A4 - Quantize Everything

Expectation: the model card says Nemotron was pretrained with NVFP4, the expert layers ran 4-bit weights and activations during training, so 4-bit activations at inference should be gentle. I naively quantized everything quantizable to W4A4 with NVIDIA’s “four-over-six” weight-scaling recipe and 1000 calibration samples.

Result: perplexity 3.296. didn’t expect this…generations fall into repetition loops, “====” and speed barely improved.

I gave up the project here for a while. My notes literally say “W4A16 is good enough, just ship it.” That conclusion was wrong. I hadn’t yet worked out which layers can take 4-bit activations. Did I just expect a recurrent Mamba state to play well with quant error!?

Step 3: W4A4 checkpoint with the right coverage

Digging with Claude Fable turned up that Nemotron Super 120B ships with an actual W4A4 recipe from NVIDIA, and it DOES NOT quantize everything:

• Routed experts → 4-bit weights and activations. Sparse (6 of 128 fire per token) and FP4-native from pretraining.
• Shared experts → FP8. They run on every token, so error there is paid on every token.
• Mamba projections → FP8. Mamba layers carry recurrent state and quantization error compounds across the sequence. I had forgotten about Mamba entirely.
• Attention and router → BF16.

This is my quant config yaml:

quant_cfg:
  # ROUTED MoE experts -> W4A4
  - quantizer_name: '*mlp.experts*weight_quantizer'
    cfg: {$import: nvfp4_four_over_six}
  - quantizer_name: '*mlp.experts*input_quantizer'
    cfg: {$import: nvfp4_dynamic}

  # SHARED experts -> FP8 per-tensor W8A8
  - quantizer_name: '*shared_experts*'
    cfg: {num_bits: e4m3}

  # Mamba mixer projections -> FP8 per-tensor W8A8
  - quantizer_name: '*mixer.in_proj*'
    cfg: {num_bits: e4m3}
  - quantizer_name: '*mixer.out_proj*'
    cfg: {num_bits: e4m3}

  # Attention, router, norms, embeddings -> BF16, untouched

I ported this coverage to Lightning, plus one addition: act_headroom, a newer activation-scale calibrator that leaves room for outlier activations instead of clipping at the observed max.

Result: perplexity 2.595, down from 3.296. Coverage alone recovered 0.70 of perplexity. This is sorted now.

TTFT: ~120 ms, 22% faster than our weight-only W4A16. FP4 tensor cores doing FP4 math at prefill, as Step 1 predicted. Can we still do better?

Same prefill MoE layer. Top (W4A16) runs Marlin: 396.8 + 338.5 µs. Bottom (our W4A4) runs CUTLASS FP4: 177.9 + 136.2 µs, plus a few small FP4 setup kernels.

TPOT is 3.47 ms, slower than W4A16’s 3.07. More quantization, slower decode. I profile that properly in the interlude.

Step 4: LM Head Quantization

One dense layer was still BF16: the lm_head. 131,072-token vocabulary × 2,688 hidden = 352M parameters, read from memory on every step to produce one token’s logits.

704 MB in BF16 → 176 MB in NVFP4 saves 528 MB of reads per step, ~330 µs at 1.60 TB/s.

The recipe change is two lines:

  # LM HEAD -> W4A16 four-over-six.
  - quantizer_name: 'lm_head.weight_quantizer'
    enable: true
    cfg: {$import: nvfp4_four_over_six}

Result: TPOT 3.47 → 3.16 ms. Predicted 330 µs, saved 330 µs. Perplexity 2.551, within 1.2% of the weight-only checkpoint.

The quantized lm_head selected in Perfetto: 139 µs per step to read the 176 MB vocab matrix and produce one token’s logits.

Okay, now prefill is much faster, decode still 3% slower than W4A16.

Interlude: why didn’t decode get faster?

W4A4 made prefill 22% faster than W4A16, but decode is slower. I traced one decode step of each checkpoint before optimizing anything else.

At batch size 1, a GEMM is a memory read. The whole weight matrix streams in from HBM to multiply against one row of activations; the arithmetic is nearly free and the traffic is everything. Both checkpoints store the same 4-bit weights, so they read the same bytes. 4-bit activations cannot speed up decode because decode is not compute-bound. They only paid off at prefill.

One MoE layer at decode. The two expert GEMMs are slower on our W4A4 (bottom, CUTLASS: 25.41 + 19.58 µs) than on W4A16 (top, Marlin: 21.70 + 17.63 µs), plus a few small FP4 setup kernels.

So decode is stuck at batch size 1, not stuck in general. If one token can’t keep the hardware busy, the fix is more tokens per step, not a better kernel.

Step 5: MTP speculative decoding - the model drafts for itself

More tokens per step is speculative decoding: a small drafter proposes k tokens, the big model verifies all of them in one pass, and rejection sampling keeps the output distribution identical. This model has a built-in MTP head in the checkpoint, trained to predict the token after next. I turned it on at k=3. Acceptance on real text is nearly flat per position (0.65 / 0.60 / 0.56), about 2.8 accepted tokens per step.

Result: TPOT 2.31 ms, ~404 tokens/sec per query, TTFT ~47 ms. Every step now verifies four tokens at once, which is exactly the batch-size-1 problem from the interlude going away. Fastest config so far, and I could have stopped here…haha, two more steps please!

Step 6: quantize the drafter

The MTP head is a full transformer block with its own attention and its own 128-expert MoE, and it ships in BF16. So on a 4-bit checkpoint the drafter was still running 16-bit GEMMs: ~805 µs per step.

I quantized the drafter’s experts to the same NVFP4 by hand, because the exporter never touches the MTP block.

Two issues:

• The quantization metadata lives in three places and vLLM only reads the copy in config.json.
• Even with 4-bit weights, vLLM builds the drafter unquantized by default and bit-casts the packed tensors into BF16. Output still correct (rejection sampling), just no speedup but the acceptance was 1.03. Fix: "quantization": "modelopt_mixed" in the speculative config.

Acceptance held at 2.80, same as the BF16 drafter. Result: TPOT 2.31 → 2.21 ms, ~418 tokens/sec per query, TTFT ~48 ms.

I was going to call this the final config. Then I traced a decode step one more time. Last one, I promise!

Step 7: prep stall in drafter

The trace showed the GPU idle for about a quarter of every step, as a stall after the drafter’s third pass while Python prepared the next step. Roughly 6 ms of host work per ~5.3 ms of GPU work: the GPU finishes and waits.

Fix: Record the drafter’s 2nd and 3rd passes as one CUDA graph. A small vLLM plugin to capture and launch one full CUDAGraph. The trace confirmed 1.27 ms less host time per step.

With prefix caching on, this Mamba hybrid waits for the accepted-token count on the CPU before it can prepare the next step, a synchronize() every step. Prefix caching does nothing for a single stream, so I turned it off.

One decode step. Step 6 (top) at 7 ms 425 µs, Step 7 (bottom) at 5 ms 483 µs. Same GPU work (same three execute_context graphs in both kernel tables). The long empty stretch after the drafter’s kernels at the top-right strip is the GPU waiting for Python.

Result: TPOT 2.21 → 1.97 ms, ~465 tokens/sec per query, TTFT ~49 ms. 12% faster than Step 6, and within ~5% of the GPU floor (1.88 ms/token, the step’s GPU work alone). vLLM’s newer opt-in model runner fixes the same problem its own way, fused draft passes and no CPU sync.

Final: 179 → ~465 tokens/sec per query, 91.5 → 49 ms to first token on one RTX PRO 6000. 2.6× the per-query throughput of the BF16 baseline and first-token latency nearly halved, from a checkpoint and a serving recipe built along the way.

Coming in Part 2: filling the GPU

This post optimized one request at a time. How many tokens per second one GPU can push when packed with concurrent requests needs a different optimization setup: KV-cache admission, the Mamba-state paging that caps it, and speculative decoding starting to hurt once the machine is full. Same GPU, same checkpoint, ~4,578 tokens/sec in aggregate at concurrency 256. How that number was reached, and the benchmarking traps on the way, is the next post.

The checkpoint is on Hugging Face, serving flags on the card: naveenrajk/nemotron35-w4a4-v13.