Execution paths

How attnquant runs the reconstruction through precompute-outside (the production path), the GEMM formulation, or the generated_linear custom op, and how the fastest path is auto-selected.

Every execution path in attnquant computes the same reconstruction identity, y = s·(W0 x) + U(Vᵀ x). What differs is when and where the rebuild happens, and that’s what decides speed. There are three paths:

1. Precompute-outside (ATTNQUANT_PRECOMPUTE_OUTSIDE=1): rebuild each weight once at load time; serve a graph identical to the base model’s. The production path.
2. GEMM formulation: express the rebuild as stock max.graph matmuls that MAX lowers to its optimized device GEMM kernels. Default for in-graph reconstruction on NVIDIA.
3. Custom op (generated_linear): rebuild per element inside a Mojo kernel, never materializing the weight. Wins on Metal at low rank.

Precompute-outside: the production path

Set ATTNQUANT_PRECOMPUTE_OUTSIDE=1 and attnquant materializes each layer’s effective weight once at load time, outside the graph:

W_eff = s · W0 + U · Vᵀ

W_eff is stored under the projection’s standard .weight key and the theta keys are removed, so the model builds with vanilla MAX Linear/StackedLinear/MLP classes and the compiled graph is identical to the base model’s, with the same launch count and kernel selection. Two things come for free on top:

• Q/K/V and gate-up fusion: the vanilla stacked layers concatenate the precomputed weights, so the three attention projections run as one matmul and gate/up as another. This takes the path from 21 kernel launches per layer to 8. It’s a launch-count property, not a measured throughput win: on both supported models the measured delta isolates to device graph capture, and with capture matched off neither model shows a reproducible fusion advantage (see throughput).
• NVFP4 quantization: with ATTNQUANT_QUANT=nvfp4, each W_eff is quantized to 4-bit at load (see configuration).

If a layer’s theta keys can’t be matched to a shape- and band-compatible W0 at load, the precompute raises immediately with the offending layers listed. A mismatch here almost always means ATTNQUANT_BLOCKS differs between decomposition and serving.

Don’t confuse this flag with ATTNQUANT_PRECOMPUTE, which folds W_eff inside the graph and recomputes the delta on every forward pass.

Precompute-outside is the fast path by a measured margin, not by convention. With the rebuild left in the graph, Gemma 3 1B bf16 runs at 0.657x vanilla dense at matched capture (150.4 vs 229.1 tok/s, CI [0.453, 0.871]), below a pre-registered 0.80 gate even with attnquant_qkv_gemm cutting the path to 8 GEMM launches per layer. See throughput.

The GEMM formulation

When the rebuild stays in the graph, attnquant_gemm expresses it with stock max.graph ops, which MAX lowers to its tuned device GEMM kernels:

y = s[None, :] · (x @ W0.T) + (x @ V) @ U.T

To keep the dispatch count low, the path fuses W0.T and V along the input axis so one x @ [W0.T, V] matmul produces both the shared-block term and the low-rank down-projection, then a second matmul applies the up-projection. The matmuls run in the input dtype (bf16 on GPU) and accumulate in float32 via Tensor Cores; the final elementwise combine also runs in the input dtype, matching vanilla dense precision.

On the GEMM path, the three attention projections batch into two dispatches (one fused GEMM over the concatenated W0 blocks and down-projections, one block-diagonal up-projection GEMM) via attnquant_qkv_gemm (ATTNQUANT_QKV_FUSE=1, on by default). The MLP fuses gate and up into one GEMM the same way.

The custom op

The generated_linear Mojo custom op reconstructs each weight inside the kernel, one dispatch per projection, never materializing the [out_dim, in_dim] tensor. It has CPU, Metal, and NVIDIA (B200/SM100) implementations, all accumulating in float32. On the B200 it uses vectorized 128-bit loads, warp-shuffle reductions, one warp per output row, and a tensor-core (MMA) variant for large output dimensions; the CUDA-only features compile in only on machines with an NVIDIA accelerator, so the same kernel file serves Metal and CPU.

This path wins at low rank on Metal (for example ATTNQUANT_RANK=4), where per-element reconstruction is cheaper than routing through large matmuls. On NVIDIA the GEMM path is faster at every rank measured, so the auto-select routes there. The kernel’s shared-memory layout supports rank up to 64; attnquant refuses higher ranks on this path and points you to the GEMM path.

On the custom-op side, the MLP can also fuse gate, up, activation, and down into a single generated_mlp dispatch (ATTNQUANT_FUSED_MLP=1, on by default).

How the path is selected

use_gemm_formulation resolves ATTNQUANT_GEMM in this order:

AttnQuant detects the accelerator once via max.driver.accelerator_api and logs the decision once at startup. For a normal serve, don’t set any of this by hand: ATTNQUANT_PRECOMPUTE_OUTSIDE=1 with ATTNQUANT_GEMM unset gives the fastest configuration on NVIDIA.

See runtime configuration for the full list of environment variables and their defaults.