GPU acceleration

The mechanisms behind fast attnquant decode (CUDA graph capture, which carries the whole measured delta on both models, plus Q/K/V fusion and NVFP4 quantization) and the profiling data behind each.

AttnQuant applies three mechanisms to GPU decode: CUDA graph capture eliminates kernel launch overhead, Q/K/V fusion issues one matmul where the built-in path issues three, and NVFP4 quantization cuts whole-model weight bytes 3.18x. Only the first carries a measured throughput delta on either supported model; the other two are structural savings that stay below the harness floor at batch 1. This page explains each mechanism and the data behind it. For the full benchmark tables, see throughput.

Where GPU time goes

Weight-reading matmuls dominate the GPU timeline. An nsys profile of both models on an NVIDIA B200 shows the kernel breakdown:

On the 31B model, matmul kernels dominate at 83.5% of GPU time. These are GEMVs reading weights from memory, and the GPU reports 98.3% utilization.

That doesn’t make batch-1 decode bandwidth-bound. A 1.90x cut in bytes moved per step buys only 1.19x on the serve path, and an nsys decomposition of a 12.29 ms built-in NVFP4 step measures where the rest goes: 3.94 ms (32%) is irreducible byte movement at 8 TB/s, ~3.42 ms (28%) is weight-GEMM time above that floor, ~2.45 ms (20%) is byte-invariant GPU work (attention, norms, elementwise, activation quantization), and 3.32 ms (27%) is host gap. Host gap plus byte-invariant GPU work is 5.77 ms: 47% of the step is untouched by moving fewer bytes. An earlier T = bytes/BW + C fit put that term at 78% (serve) / 66% (generate); it was labeled inferred and it overestimated, because bf16 and NVFP4 kernels don’t achieve the same effective bandwidth.

The kernels themselves are the other half of the answer. Achieved bandwidth by weight family on the same profile: attention projections (bf16, unquantized) 4.59 TB/s or 57% of peak, MLP (NVFP4, block-scaled) 3.76 TB/s or 47%, and lm_head / tied embedding (bf16) 5.16 TB/s or 65%. The NVFP4 kernels are the least bandwidth-efficient family on the step, which is why cutting MLP bytes 3.56x doesn’t convert. It isn’t a dequantization pass: there’s no dequant kernel at all (0 launches, 0.00% of GPU time), and the block-scaled GEMM consumes fp4 weights directly.

See throughput for the full tables.

CUDA graph capture

At batch-1 decode, a transformer issues hundreds of kernel dispatches per token. CUDA graph capture replays the entire decode step as a single graph instead of issuing them one at a time. The conventional 5–10 µs-per-launch estimate would make that worth 1.5–3 ms per token, but that figure is inferred and it doesn’t describe every path: on the built-in NVFP4 path, forcing capture on put 99.1% of kernels into graphs and moved the measured host gap from 3316 µs to 3318 µs, no change, which matches the ~0% throughput result below. Capture is worth what it’s measured to be, per path.

Capture is enabled on both supported models: supports_device_graph_capture=True in models/gemma3/arch.py:87 and models/gemma4/arch.py:89.

On Gemma 3 1B it’s the entire measured attnquant advantage, and the built-in path can’t match it. On the serve harness at concurrency 1, pooled over 4 independent server sessions per capture-off cell:

Capture on versus off is +30.1% (Mann-Whitney z=−4.28), stable in every session. Forcing capture on the built-in path raises NotImplementedError: Gemma3Inputs does not define model ABI buffers (pipeline_model.py:205): MAX’s Gemma 3 architecture never implements the buffer ABI that serve-side capture requires. AttnQuant supplies it as AttnQuantGemma3Inputs.buffers, so on this model capture is a capability attnquant adds rather than a flag the built-in path leaves unset. With capture matched off, the two paths show no reproducible difference; this model’s within-cell spread runs 30-57% and is session-level, putting its minimum detectable effect near 20%.

On Gemma 4 31B with NVFP4, capture isn’t the mechanism. Measured with capture pinned explicitly on every cell:

Capture is worth about +3% on each path, with overlapping ranges, so neither is resolved. The architecture declaration (supports_device_graph_capture) only selects the default; an explicit --device-graph-capture overrides it on either path, so it’s not the source of the advantage.

The advantage is NVFP4 coverage, and 92% of it. AttnQuant quantizes all 410 projections and reads 19.295 GB per token, while NVIDIA’s shipped checkpoint quantizes 180 and reads 31.482 GB. (Its ignore list also names lm_head, but that entry is moot: tie_word_embeddings: true makes the tied head bf16 on both paths.) Match the coverage with ATTNQUANT_QUANT_IGNORE=self_attn and the gap collapses from +9.4% to +0.7%, ranges overlapping. That also means the comparison isn’t precision-matched, and the coverage costs about 0.41 pp of MMLU-Pro.

Q/K/V fusion

The precompute-outside path fuses the query, key, and value projections into a single matmul where the built-in path issues three. This cuts both the dispatch count and HBM traffic: one weight read instead of three.

Fusion has no measured throughput attribution on either model. No measurement here isolates it, so it’s not claimed as the source of any delta:

Fusion is a real structural saving (it cuts the precompute-outside path from 21 kernel launches per layer to 8), and on the 1B any effect sits below a ~20% measurement floor. On the 31B NVFP4 the advantage is attributed: a coverage-matched control accounts for about 92% of it, leaving no room for a material fusion effect. Treat fusion as a launch-count property of the graph, not a throughput claim.

NVFP4 quantization

NVFP4 (4-bit weights with per-block float8 scales) reduces the bytes each projection reads. Set ATTNQUANT_QUANT=nvfp4 and the precompute-outside path quantizes each effective weight at load time, storing packed uint8 weights plus float8 block scales instead of bf16.

On Gemma 4 31B, projection weights go from 61.4 GB to 16.5 GB (3.56x), and whole-model weight bytes from 61.4 GB to 19.3 GB (3.18x) once the 2.82 GB of bf16 embeddings and lm_head are counted. The ratio isn’t 4x because the fp8 block scales cost 0.5 bit/param. An earlier figure of “RSS 246 GB to 76 GB” described host-RSS load-time peaks, not weight bytes.

The two NVFP4 paths don’t load the same bytes. The built-in path loads nvidia/Gemma-4-31B-IT-NVFP4, whose quantization_config.ignore list leaves every self_attn layer in bf16 (the tied lm_head is bf16 either way); attnquant loads google/gemma-4-31b-it and quantizes every projection at load. They’re both “NVFP4,” but they’re not numerically equivalent: attnquant moves ~19.3 GB of weights per token against ~31.5 GB, 1.63x fewer, and the two score the same: +0.10 ± 0.50 pp on MMLU-Pro head-to-head (p=0.924), with attnquant 0.86 pp below bf16 where MLP-only quantization is 0.95 ± 0.45 pp below it. See quality for the paired measurements.

The quantization win itself is smaller than the byte cut. Measured at batch/concurrency 1, a 1.90x reduction in bytes moved per decode step buys 1.31x on the model-only max generate path and 1.19x on the serve path. See throughput for the profiled decomposition of that gap.

Quantization runs once at load (measured at ~51 minutes for the 31B model’s 410 projections on the first run), and a disk cache of the quantized state dict cut the next load to ~22 minutes.

How the pieces fit

Three execution choices make the mechanisms above land on MAX’s fastest kernels:

• Precompute-outside (ATTNQUANT_PRECOMPUTE_OUTSIDE=1): builds the effective weight once at load, so each projection is a single vanilla matmul and the compiled graph is identical to the base model’s. MAX’s own optimized B200 GEMM kernels apply directly. This is the recommended decode path.
• GEMM formulation (auto-selected on NVIDIA): expresses reconstruction with stock ops.matmul, which MAX lowers to tensor-core GEMM kernels.
• Custom op (generated_linear.mojo): fuses the shared-block dot product and low-rank delta in one kernel with no intermediate materialization. It wins on Metal at low rank; on NVIDIA the GEMM path wins at every rank.

See reconstruction for the math behind each path and throughput for the complete benchmark method.