AttnQuant inference for MAX

Quantize every projection in Gemma 3 and Gemma 4 to NVFP4 with no calibration data, giving 3.18x smaller weights and faster decode as a drop-in MAX architecture.

NVIDIA’s NVFP4 checkpoint for Gemma 4 31B, nvidia/Gemma-4-31B-IT-NVFP4, quantizes only the MLP projections to 4-bit and keeps all 230 attention projections in 16-bit bfloat16. AttnQuant quantizes all 410 projections to NVFP4, needs no calibration data, and installs as a drop-in MAX custom architecture: the same max serve and max generate commands you already use.

Batch-1 decode is memory-bound: every token reads every weight once, so the bytes read per token set the pace. Attention coverage is the difference between reading 19.3 GB per token and the NVIDIA checkpoint’s 31.5 GB. Against the unquantized bfloat16 model’s 61.4 GB, that’s 3.18x fewer bytes.

Configuration	MAX built-in	AttnQuant
Gemma 4 31B decode, NVIDIA NVFP4 checkpoint	92.0 tok/s	100.6 tok/s	+9.4%
Gemma 4 31B decode + MTP γ=3, NVIDIA NVFP4 checkpoint	247.5 tok/s	290.1 tok/s	+17.2%
Bytes read per token, NVIDIA NVFP4 checkpoint	31.5 GB	19.3 GB	1.63x fewer
Whole-model weight bytes, bfloat16 baseline	61.4 GB	19.3 GB	3.18x smaller

One NVIDIA B200, concurrency 1 on /v1/chat/completions, reproduced across three machines that agree to within 2%.

Results

Full coverage, no calibration: ATTNQUANT_QUANT=nvfp4 quantizes every projection at load time using round-to-nearest. There’s no calibration set to collect, no dataset to license, and nothing to tune per model.
A real accuracy trade, priced: Quantizing attention costs about 0.41 pp of MMLU-Pro on top of the 0.44 pp that MLP-only NVFP4 already costs (84.03 against bf16’s 84.88). Set ATTNQUANT_QUANT_IGNORE=self_attn to match NVIDIA’s coverage instead; you get the accuracy back and give up the speed.
Speculative decoding composes: With a Gemma 4 assistant drafter at γ=3, attnquant serves 290.1 tok/s, 2.9x non-speculative, and the coverage advantage holds at +17.2%.
Low time-to-first-token at concurrency: run_max.py enables in-flight batching, which lets prefill ride along with the decode batch. Median TTFT at concurrency 64 is 220.5 ms against 434.5 ms without it.
Nothing to retrain: AttnQuant decomposes the checkpoint analytically at load and rebuilds each weight exactly, so at the shipped operating point output matches the original model under logit-parity verification.

Where the speed comes from

AttnQuant reads fewer bytes per token because it quantizes more of the model. That’s the whole mechanism, and it’s measurable directly: match NVIDIA’s coverage with ATTNQUANT_QUANT_IGNORE=self_attn and throughput returns to baseline (92.6 against 92.0), while full coverage leads by 9.4%.

At ATTNQUANT_BLOCKS=60, the shipped setting, the decomposition reconstructs the original weights bit-for-bit, so it costs nothing at runtime and changes no output. Device graph capture adds about 3% on either path. Q/K/V fusion cuts the graph from 21 kernel launches per layer to 8.

Throughput has the full comparison, including against vLLM, and the harness behind every number.

How it works

AttnQuant represents each projection weight as a theta bundle (a shared block W0, a per-row scale s, and a low-rank correction U, V) computed from the trained checkpoint in one analytic pass, no training. At load time it rebuilds each effective weight once (W_eff = s·W0 + U·Vᵀ), optionally quantizes it to NVFP4, and hands MAX a graph identical to the base model’s plus fusion. At the production operating point every rebuilt weight equals the original, which is why the output is native-equivalent.

Try it

Serving Gemma 3 1B takes two environment variables and one flag:

ATTNQUANT_BLOCKS=26 ATTNQUANT_PRECOMPUTE_OUTSIDE=1 \
  pixi run python run_max.py serve \
    --model google/gemma-3-1b-it \
    --custom-architectures attnquant \
    --port 3100

The quickstart walks through install, generate, and serve from a clean checkout.

Interactive throughput comparison

Select a model and quantization to compare attnquant against the built-in path. The default view is the flagship result: Gemma 4 31B with NVFP4.

Higher is better · accent = attnquant · one NVIDIA B200, max serve + max benchmark at concurrency 1 on /v1/chat/completions, capture pinned explicitly · the +11.6% is the capture-matched comparison and its mechanism is still under investigation · the +10.2% badge is against vLLM, not the MAX built-in

Decode throughput	100.6 tok/s (serve, c=1, chat endpoint)
With MTP drafter	290.1 tok/s at gamma=3; +17.2% over built-in, +11.1% over vLLM
vs external, no speculation	MAX 92-93 vs vLLM 110.5 and TRT-LLM 110.9 — MAX ~19% slower; traced to a dead sync in MAX's sampler, whose removal reaches 120.9
TRT-LLM + MTP	unsupported for Gemma 4 in 1.3.0rc22 — no speculative comparison exists
Mechanism	NVFP4 coverage, ~92% of it — match the coverage and the gap is +0.7%
Speculative efficiency	MAX realises 97.2% of its acceptance ceiling vs vLLM's 84.2%
Measured accept length	2.75-2.78 at gamma=3, chat; validated to ~2% against vLLM's direct counter
Weight bytes per token	19.295 GB attnquant vs 31.482 GB built-in NVFP4 (1.63x fewer)
Accuracy cost of that coverage	-0.41 pp MMLU-Pro (n=12032); -0.85 pp total vs bf16
Endpoint dependence	under MTP the chat endpoint reverses the sign vs /v1/completions
NVFP4 kernel efficiency	3.76 TB/s achieved = 47% of peak (bf16 families reach 57-65%)
Harness	max serve + max benchmark, c=1, /v1/chat/completions, MDE ~5%
GPU	NVIDIA B200

What’s supported today

Gemma 4 31B (Gemma4ForConditionalGeneration, 60 layers): +9.4% over built-in NVFP4 and +17.2% under MTP, with ~92% of each gain coming from NVFP4 coverage rather than the decomposition; 3.18x whole-model weight compression, for −0.41 pp of MMLU-Pro.
Gemma 3 1B (Gemma3ForCausalLM, 26 layers): device graph capture on a path the built-in architecture can’t capture at all, worth +30.1% decode in bf16.
Any decoder you port: the core is model-agnostic; the extending guide shows the six-file pattern that wires a new architecture onto it.

Both models serve through max generate and max serve on Linux NVIDIA GPUs and Apple silicon, with the faster of two math-identical execution paths auto-selected per accelerator. Two limits before you deploy: the bf16 fast path measures about 1.9x the built-in path’s peak GPU memory (NVFP4 recovers most of that), and each model runs on a single device because tensor parallelism isn’t supported yet.

Next steps

Quickstart: from clean checkout to served endpoint.
Serving configuration: every ATTNQUANT_* variable and config.json field.
How attnquant works: the theta bundle and the reconstruction math.
Metrics: every published number with the command that reproduces it.