Serve and generate with attnquant

Run a attnquant Gemma 3 or Gemma 4 model on MAX through run_max.py, register the custom architecture with --custom-architectures, and turn on the fast path.

A attnquant model registers as a MAX custom architecture, so you serve and generate with the same max serve and max generate commands as any other MAX model. Point MAX at a Hugging Face checkpoint, add --custom-architectures attnquant, and run the command through the run_max.py wrapper.

This page covers the run commands, how the custom architecture loads over MAX’s built-in, and the two settings that give you the fastest decode on NVIDIA. For the full flag reference, see configuration.

Run through run_max.py

Invoke every attnquant command through run_max.py, not the bare max CLI. The wrapper clears stale MODULAR_* environment variables that a previous install can leave behind, adds the repo root to sys.path so import attnquant resolves, then calls the MAX entrypoint. The shape of every command is:

pixi run python run_max.py <serve|generate> --custom-architectures attnquant ...

Generate a completion

To run a one-off completion against the Gemma 3 1B checkpoint, pass a prompt and a token budget:

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

Serve an endpoint

To bring up an OpenAI-compatible endpoint, pass a port. Set ATTNQUANT_BLOCKS to the model’s layer count for the passthrough operating point, and on NVIDIA set ATTNQUANT_PRECOMPUTE_OUTSIDE=1 (see Set the fast path on NVIDIA):

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 8192 \
    --device-memory-utilization 0.9

The Gemma 4 31B checkpoint (60 layers) serves through the same command with a different --model. Add ATTNQUANT_QUANT=nvfp4 for 3.18x smaller weights and the fastest measured decode (100.6 tok/s on one B200, serve harness at concurrency 1 on /v1/chat/completions):

ATTNQUANT_BLOCKS=60 ATTNQUANT_QUANT=nvfp4 ATTNQUANT_PRECOMPUTE_OUTSIDE=1 \
  pixi run python run_max.py serve \
    --model google/gemma-4-31B-it \
    --custom-architectures attnquant \
    --port 3100

Once the server is up, call it like any OpenAI-compatible endpoint:

curl http://localhost:3100/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "google/gemma-3-1b-it",
    "messages": [{"role": "user", "content": "What is the capital of France?"}]
  }'

How the custom architecture loads

When you pass --custom-architectures attnquant, MAX imports the package and reads its ARCHITECTURES list. The attnquant Gemma 3 registers under name="Gemma3ForCausalLM" and the attnquant Gemma 4 under name="Gemma4ForConditionalGeneration". Each registration name matches Hugging Face’s architecture class, which is also the name of MAX’s built-in architecture, so the custom architecture wins over the built-in of the same name. The package also clears MAX’s lazy registry entry for each overridden name so MAX doesn’t try to materialize the built-in afterward. For the override mechanism and the per-model file layout, see extending.

At load time the safetensors weight adapter auto-decomposes the Hugging Face checkpoint into theta bundles. If the checkpoint already carries theta keys, the adapter skips decomposition. The pipeline then builds a graph whose linear projections reconstruct their weights from theta, auto-selecting the faster of the two math-identical reconstruction paths for your accelerator. You don’t choose the path by hand for a normal run.

Set the fast path on NVIDIA

On NVIDIA GPUs, set two environment variables for the fastest decode with output that matches the original model:

ATTNQUANT_BLOCKS=<num_layers> ATTNQUANT_PRECOMPUTE_OUTSIDE=1

ATTNQUANT_BLOCKS set to the layer count (26 for Gemma 3 1B, 60 for Gemma 4 31B) selects the passthrough operating point, where each rebuilt weight equals the original. With ATTNQUANT_PRECOMPUTE_OUTSIDE=1, attnquant materializes each layer’s effective weight W_eff = s·W0 + U·Vᵀ once at load time, outside the graph. Every projection then runs a single x @ W_effᵀ matmul, the same launch count as the built-in dense path, with Q/K/V fused into one matmul. Leave ATTNQUANT_GEMM unset so device-aware selection picks the GEMM path on NVIDIA.

On one NVIDIA B200 serving Gemma 3 1B, with max serve + max benchmark at concurrency 1 pooled over 4 independent server sessions per cell, the fast path measured:

Device graph capture on versus off: 248.9 vs 191.3 tok/s (+30.1%). AttnQuant implements the buffer ABI that capture needs, so it can turn capture on; MAX’s built-in Gemma3ForCausalLM can’t capture at all.
Capture matched off, attnquant versus built-in: 191.3 vs 207.9 tok/s, no reproducible difference at this model’s ~20% minimum detectable effect.

The measured win is capture, not Q/K/V fusion. See throughput for the session-level data, the spread, and the memory footprint of this path.

Supported encodings and devices

Both attnquant architectures default to bfloat16 and also support float32; Gemma 4 additionally accepts float4_e2m1fnx2 checkpoints, and on the precompute-outside path ATTNQUANT_QUANT=nvfp4 adds 4-bit quantized weights for either model. Each model runs on a single device: tensor parallelism doesn’t yet compose with weight reconstruction, so multi_gpu_supported is false.

Next steps

Once the model is serving, tune it and read how it works:

Configuration: the canonical ATTNQUANT_* environment variables and config.json fields.
Reconstruction: the y = s·(W0 x) + U·(Vᵀ x) formula and how the path is auto-selected.
Extending: the six-file pattern to port a new decoder architecture.