AttnQuantLinear

Reference for the drop-in attnquant layers (AttnQuantLinear, AttnQuantStackedLinear, AttnQuantMLP), including what they replace, how the shared block W0 is injected with set_w0, and the auto-selected forward paths.

core/attnquant_linear.py and core/attnquant_mlp.py define the layers a attnquant model substitutes for MAX’s built-in linear stack. Each layer rebuilds its projection weight from a theta bundle instead of loading a stored dense matrix:

y = s * (W0 x) + U (Vᵀ x)

Three classes cover the substitution:

• AttnQuantLinear: replaces max.nn.Linear, the per-projection layer.
• AttnQuantStackedLinear: replaces the fused StackedLinear used for attention q_proj, k_proj, v_proj.
• AttnQuantMLP: replaces max.nn.MLP, the gate, up, down block.

For the reconstruction math, see reconstruction. For how a model wires its projections onto these layers, see extending attnquant.

AttnQuantLinear

AttnQuantLinear accepts the same constructor signature as max.nn.Linear (in_dim, out_dim, dtype, device, has_bias, quant_config, name, and so on) plus a lora_rank, and it absorbs any extra keyword arguments. That makes functools.partial(AttnQuantLinear, lora_rank=...) a drop-in for the linear_cls parameter every standard MAX architecture threads through construction. The lora_rank argument is r, the rank of the U, V delta. The constructor default is 4; the shipped models pass the config’s lora_rank (default 32) through the partial.

Instead of computing x @ weight.T against a stored dense matrix, __call__ rebuilds the projection weight from the theta bundle at inference time. There are two math-identical paths, auto-selected by accelerator and lora_rank:

• Custom op: ops.custom("generated_linear", ...) reconstructs the weight inside a Mojo kernel and never materializes the full [out_dim, in_dim] tensor. Metal-oriented, selected at low rank; supports rank up to 64 (the kernel’s shared-memory limit) and raises above that.
• GEMM: stock max.graph matmuls that express the same output through MAX’s device GEMM kernels. Selected on NVIDIA and other non-Metal accelerators.

Set ATTNQUANT_GEMM to force a path, or leave it unset to auto-select. The selection rules live in core/gemm_formulation.py. See serving configuration for the full environment-variable reference.

Owned vs injected weights

The per-layer theta weights are Weight objects owned by the layer, so their state-dict names come out as <name>.s, <name>.U, <name>.V:

The layer doesn’t own the shared block W0. The model declares it once (one per family per band) and injects it before the forward pass:

layer.set_w0(w0)   # must be called before __call__

Calling the layer without an injected W0 raises RuntimeError. The outside-graph precompute path isn’t an exception inside the layer: it never builds AttnQuantLinear at all. The shipped models substitute vanilla Linear/MLP there, so set_w0() is never called (see Fast path). The model’s _inject_w0_recursive walks the module tree and calls set_w0() on every AttnQuantLinear child, matching each family by its attribute name.

Fast path: precompute outside the graph

Set ATTNQUANT_PRECOMPUTE_OUTSIDE=1 (recommended) to materialize the effective weight once at load time:

W_eff = s[:, None] * W0 + U @ Vᵀ

precompute_weff_outside_graph computes W_eff in numpy from the decomposed theta weights, casts it to the model dtype, stores it under each projection’s standard .weight key, and removes the theta keys from the state dict. The model then builds with vanilla Linear/StackedLinear/MLP classes, so each projection runs a single x @ W_eff.T matmul with zero runtime reconstruction overhead (the same launch count as MAX’s built-in dense path), and Q/K/V and gate-up fusion come from the vanilla stacked layers concatenating those weights. At the passthrough operating point (ATTNQUANT_BLOCKS set to the layer count), output is native-equivalent to the built-in path under the logit-parity checks. With ATTNQUANT_QUANT=nvfp4, each W_eff is additionally quantized to 4-bit NVFP4 at load (see configuration).

The inside-graph variant ATTNQUANT_PRECOMPUTE=1 recomputes U @ Vᵀ on every forward pass inside the graph. Prefer ATTNQUANT_PRECOMPUTE_OUTSIDE for serving.

Compatibility surface

To satisfy framework code that inspects a linear layer, AttnQuantLinear exposes a placeholder .weight with the correct [out_dim, in_dim] shape (its value is never used for compute), plus .bias, .weight_scale, .input_scale, and .weight_scale_2 (all None on the unquantized attnquant path). It implements the Shardable protocol so sharding_strategy reads and writes work correctly, but shard() raises NotImplementedError for multi-device: tensor parallelism doesn’t yet compose with in-kernel reconstruction (multi_gpu_supported=False).

AttnQuantStackedLinear

The base StackedLinear concatenates its children’s .weight tensors into one stacked weight and calls linear() once. That path materializes the full weight and never invokes the custom op, so it’s incompatible with attnquant children.

AttnQuantStackedLinear overrides __call__ to run each child’s own reconstruction and return a list of per-child outputs. This keeps grouped-query attention (GQA) correct: when q_proj, k_proj, and v_proj have different output dimensions, each output keeps its own out_dim rather than being obscured by concatenation. Child weights appear at self_attn.q_proj.* (not qkv_proj.q_proj.*), matching the checkpoint layout without a per-architecture weight mapping.

When ATTNQUANT_QKV_FUSE=1 (the default) and the active path allows it, the three Q/K/V projections share their x matmul across a single fused dispatch. With outside-graph precompute the model uses the vanilla StackedLinear, which concatenates the precomputed .weight tensors and runs one matmul, then splits by out_dim. This fusion (one matmul instead of three) cuts the launch count; it isn’t where the measured decode win comes from. That isolates to device graph capture.

AttnQuantMLP

The base max.nn.MLP has a fused path that concatenates gate_proj.weight and up_proj.weight into one tensor and calls linear() once. That path materializes the stacked weight and never invokes the custom op, so it’s incompatible with attnquant children.

AttnQuantMLP always runs each projection through its own reconstruction, then applies the configured activation with SwiGLU gating (activation(gate) * up, using the model’s hidden_activation, for example gelu_tanh for Gemma 3), and dispatches down_proj. With outside-graph precompute (ATTNQUANT_PRECOMPUTE_OUTSIDE=1) the shipped models replace it with the vanilla max.nn.MLP, whose fused path concatenates the precomputed gate and up weights into one matmul. On the reconstruction paths, __call__ checks these in order:

• Fused MLP custom op (ATTNQUANT_FUSED_MLP=1, custom-op path only): a single generated_mlp custom op fuses gate, up, activation, and down into one dispatch.
• GEMM gate-up fusion (ATTNQUANT_QKV_FUSE=1, GEMM path): batches gate and up into one fused GEMM, then activation and down_proj separately.
• Separate dispatches (fallback): calls gate_proj, up_proj, and down_proj each through its own __call__.

Like the other layers, AttnQuantMLP.shard() returns [self] for a single device and raises NotImplementedError for multiple devices, because tensor parallelism isn’t supported on the attnquant path.

Serve a attnquant model

These layers register through the custom architecture, so you serve with the standard MAX flow through the run_max.py wrapper. Set ATTNQUANT_BLOCKS to the layer count (passthrough) and ATTNQUANT_PRECOMPUTE_OUTSIDE=1 for the fast path. To generate from a prompt:

ATTNQUANT_BLOCKS=26 ATTNQUANT_PRECOMPUTE_OUTSIDE=1 pixi run python run_max.py generate \
  --model google/gemma-3-1b-it \
  --custom-architectures attnquant \
  --prompt "..." \
  --max-new-tokens 128

To start an OpenAI-compatible endpoint:

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 \
  --max-length 4096 \
  --device-memory-utilization 0.9

See serving configuration for every environment variable and config.json field.