API reference

The attnquant API surface: the model-agnostic core (decompose, AttnQuantLinear, AttnQuantMLP) and generated_linear kernel reused across models, plus the per-model architecture files.

Two parts make attnquant work: a model-agnostic core plus kernel that decompose weights into theta bundles and rebuild them at inference time, and a thin per-model adapter that registers a specific architecture as a MAX custom architecture. These reference pages document that API surface: the classes, functions, and config fields you call. For the walkthrough of how the pieces fit together, follow the extending guide: the reference is the spec, the guide is the tutorial.

Reference pages

Start with the page that matches the layer you’re calling:

core.decompose

Turn a Hugging Face state dict into a flat theta bundle at load: decompose_standard_state_dict_to_attnquant, the single-weight helper, and the reconstruction-error checks. Analytic, no training step.

AttnQuantLinear

The drop-in linear_cls layer and its stacked variant, the owned (s, U, V) versus injected (W0) weights, and the reconstruction path each one drives.

Gemma 3 / Gemma 4 architecture

The SupportedArchitecture objects that register each model, the package hook that overrides MAX’s built-in of the same name, and the config fields.

Extending to a new model

The canonical how-to: the six-file per-model pattern that ports a new decoder while reusing core/ and kernels/ unchanged.

Package layout

The core/ and kernels/ code is model-agnostic: it works for any decoder with the standard projection layout. Each supported model adds a thin adapter under models/:

attnquant/
├─ core/          # model-agnostic decomposition, layers, GEMM formulation
│  ├─ decompose.py           # HF weights -> theta bundle (and back)
│  ├─ attnquant_linear.py   # AttnQuantLinear, AttnQuantStackedLinear
│  ├─ attnquant_mlp.py      # AttnQuantMLP (gate/up/down projections)
│  ├─ quantize.py            # NVFP4 quantization of the effective weight
│  ├─ compile_cache.py       # compiled-graph (.mef) disk cache
│  └─ gemm_formulation.py    # GEMM-path reconstruction and QKV fusion
├─ kernels/       # generated_linear Mojo custom op (CPU + Metal + NVIDIA B200)
│  └─ generated_linear.mojo
├─ models/
│  ├─ gemma3/     # Gemma 3 1B (Gemma3ForCausalLM, 26 layers)
│  └─ gemma4/     # Gemma 4 31B (Gemma4ForConditionalGeneration, 60 layers)
└─ run_max.py     # MAX CLI wrapper: run_max.py serve|generate ...

The model-agnostic core exports its public surface from core/__init__.py: AttnQuantLinear and AttnQuantStackedLinear, AttnQuantMLP, the decompose helpers, and the GEMM-path functions (use_gemm_formulation, attnquant_gemm, attnquant_qkv_gemm). AttnQuantMLP subclasses max.nn.MLP and dispatches the gate, up, and down projections as AttnQuantLinear children so each one runs its own reconstruction. The generated_linear Mojo custom op in kernels/ is one of the two math-identical reconstruction paths (auto-selected by accelerator and rank); the other is the stock max.graph GEMM path in gemm_formulation.py.

Per-model adapters

Each model under models/<name>/ follows the same six-file pattern and reuses core/ and kernels/ without changes. The reference for the two shipped models is on the architecture page; the extending guide is canonical for the how. One detail is worth stating here: each arch.py registers a SupportedArchitecture whose name is the Hugging Face architecture class (for example Gemma3ForCausalLM), and the top-level ARCHITECTURES list is what MAX reads when you pass --custom-architectures attnquant. Because the name matches a built-in, the attnquant architecture overrides MAX’s built-in of the same name. See the extending guide for why the override is needed and how to register your own.