Jet-Nemotron

PostNAS pipeline — Begins with a pre-trained full-attention model and freezes MLPs, then proceeds in four steps: learn optimal placement or removal of full-attention layers, select a linear-attention block, design a new attention block, and run a hardware-aware hyperparameter search.

Learning where full attention actually matters — A once-for-all super-network plus beam search identifies only a few layers as critical, and the important layers differ by task.

JetBlock: linear attention with dynamic convolution — The new block adds a kernel generator that produces input-conditioned causal convolutions applied to V tokens and removes static convolutions on Q/K. They report higher math and retrieval accuracy vs. prior linear blocks at similar training and inference speed.

Hardware-aware design insight — Through grid search at fixed KV cache size, they show generation speed tracks KV cache more than parameter count. The work shares head/dimension settings that hold throughput roughly constant while improving accuracy. This leads to comparable tokens/s with more parameters and better scores.

Results at a glance — Jet-Nemotron-2B outperforms or matches small full-attention models on MMLU, MMLU-Pro, BBH, math, commonsense, retrieval, coding, and long-context tasks, while delivering up to 47x decoding throughput at 64K and as high as 53.6x decoding and 6.14x prefilling speedup at 256K on H100.