Add a new model architecture

Make any decoder attnquant by mirroring the six per-model files under models/gemma3 (uniform attention) or models/gemma4 (mixed attention), reusing the model-agnostic core and kernels unchanged.

To make a new decoder architecture attnquant, mirror the six files under models/gemma3/ (or models/gemma4/ for models with mixed attention) and reuse the model-agnostic core/ and kernels/ unchanged.

This guide comes from diffing models/gemma3/ (the simple, uniform-attention template) against models/gemma4/ (the mixed-attention case: dual head dims, attention_k_eq_v tying, and per-layer_type W0 grouping). If the theta bundle is new to you, read the theta bundle first for the mental model.

For the pieces this guide references, see architecture for the field-level SupportedArchitecture description, decompose for the decomposition API, attnquant-linear for the linear layer, reconstruction for the math, and configuration for the env vars.

What you reuse versus what you write

The reconstruction math is model-agnostic and stays untouched. You never edit core/decompose.py, core/attnquant_linear.py, core/attnquant_mlp.py, core/gemm_formulation.py, or kernels/generated_linear.mojo. A new architecture is entirely a matter of wiring MAX’s native decoder to those primitives through six per-model files:

Copy the gemma3 files

Each file below names its real Gemma 3 path as the copy target and the Gemma 4 counterpart as the mixed-attention diff.

Export the architecture

Copy models/gemma3/__init__.py. It imports the arch object and exposes it as a module-level ARCHITECTURES list:

from attnquant.models.gemma3.arch import attnquant_gemma3_arch

ARCHITECTURES = [attnquant_gemma3_arch]

Gemma 4 also re-exports its config class (models/gemma4/__init__.py adds AttnQuantGemma4Config) because the multimodal batch processor imports it by name. The package-level attnquant/__init__.py concatenates every model’s ARCHITECTURES into the top-level list, then clears MAX’s lazy registration entry for each overridden name:

for _arch in ARCHITECTURES:
    PIPELINE_REGISTRY._lazy_architectures.pop(_arch.name, None)

Don’t skip this step. Without it, MAX registers your custom arch, then later tries to materialize its built-in of the same name and fails with Refusing to override existing architecture. Popping the lazy entry is what lets the same-name override win. For the serving view of this override, see serving.

Construct the SupportedArchitecture

Copy models/gemma3/arch.py. Change only these fields per model:

• name= must equal the HF architecture class name so the custom arch overrides MAX’s built-in of the same name: "Gemma3ForCausalLM" for Gemma 3, "Gemma4ForConditionalGeneration" for Gemma 4.
• example_repo_ids: the Hugging Face repos this arch serves.
• pipeline_model: your AttnQuant<Model>PipelineModel.
• config and weight_adapters={WeightsFormat.safetensors: convert_safetensor_state_dict}: your subclasses.

Everything else is copy-paste identical between Gemma 3 and Gemma 4: default_encoding="bfloat16", PipelineTask.TEXT_GENERATION, TextTokenizer, TextContext, multi_gpu_supported=False, PagedMemoryPlanner.with_activation_reservation(0, always_signal_buffers=True), and rope_type="normal".

Add the attnquant fields

Copy models/gemma3/model_config.py. Subclass the model’s MAX config (Gemma3Config, or Gemma4TextConfig for Gemma 4) and add four attnquant fields (lora_rank, tie_families, shared_blocks, tied_projections) plus band_index() and an initialize_from_config() override that reads ATTNQUANT_RANK and ATTNQUANT_BLOCKS. Here shared_blocks (the config field), ATTNQUANT_BLOCKS (the env override), and G (the symbol) are the same value. The canonical description of these four fields lives in configuration.

FAMILIES, ATTN_FAMILIES, and tied_families_for() are identical between the two models. A projection family is one of the seven projections that repeats in every decoder layer (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj). The Gemma 4 diffs to call out:

• Gemma4TextConfig overrides rope_theta and rope_scaling as properties that raise, so initialize_from_config() skips them in the getattr loop and passes sentinels (rope_theta=-1, rope_scaling=None).
• Gemma 4 adds vision_config=None, unquantized_dtype=DType.bfloat16, and a text_config property that returns self: a shim so the inherited multimodal batch processor can read config.text_config.hidden_size.
• Gemma 4’s band_index() is combined with layer_types downstream so that W0 is grouped per (band, layer_type, family). Layers of different types never share a W0 because their projection dimensions differ.

Decompose HF weights: the uniform case (gemma3)

models/gemma3/weight_adapters.py is the template for any model whose attention projections have a single, uniform head dimension. convert_safetensor_state_dict runs seven steps:

1. Remap HF prefixes (GEMMA3_SAFETENSOR_MAP, model. to language_model.) via str.removeprefix().
2. Strip the language_model. prefix so the model-agnostic decompose can match layers.{i}.self_attn.{fam}.weight.
3. Run the is_attnquant_theta_dict() passthrough check. Pre-converted theta bundles skip decomposition.
4. Run decompose_standard_state_dict_to_attnquant(rank, shared_blocks).
5. Remap flat keys to nested FQNs via _LAYER_FAM_RE (layers.{i}.q_proj.s to layers.{i}.self_attn.q_proj.s; MLP families go under mlp).
6. Cast decomposed numpy arrays to the model dtype so load_state_dict() validation passes.
7. Re-add the language_model. prefix to match the wrapper module’s FQNs.

Decompose HF weights: the mixed-attention case (gemma4)

When a model has mixed attention (dual head dims), you can’t call the model-agnostic decompose directly, because q/k/v have different output dimensions on sliding versus full-attention layers. models/gemma4/weight_adapters.py is the file a new mixed-attention model copies. It adds _decompose_with_layer_type_w0(), which:

• Reads layer_types and attention_k_eq_v from the config.
• Splits attention weights by (band, layer_type, family) so sliding (head_dim 256) and full (512) layers get separate shared blocks, with keys W0.sliding_attention.q_proj and W0.full_attention.q_proj.
• Keeps MLP families grouped per (band, family) since gate/up/down dimensions match across layer types (single W0.gate_proj, and so on).
• Skips v_proj on full-attention layers when attention_k_eq_v is set. Those layers reuse K, so no V weight exists in the checkpoint. Sliding layers always keep v_proj.

The prefix map also differs: GEMMA4_LANGUAGE_SAFETENSOR_MAP strips model.language_model. (and the bare-model. wrapper variants) to bare keys. It deliberately omits the MoE router/expert remaps MAX’s native adapter carries (router.proj.weight to moe_block.gate.gate_score.weight, experts. to moe_block.experts.): the attnquant decoder hardcodes enable_moe_block=False, so there’s no moe_block.* submodule to load into, and a key rename alone wouldn’t split the stacked expert tensors.

Override the graph build

Copy models/gemma3/model.py. Subclass the MAX pipeline model (Gemma3Model, or Gemma3_MultiModalModel for Gemma 4) and override _build_graph (Gemma 4 splits this into _build_language_graph and load_model) to:

1. Run the adapter to produce the theta state dict.
2. Build the AttnQuant<Model> wrapper whose self.language_model is the attnquant text model.
3. Pass custom_extensions=[KERNELS_DIR] only on the custom-op path. Omit it when use_gemm_formulation() selects the GEMM path (stock max.graph matmuls), which has no kernel to compile.
4. Call load_state_dict(..., strict=False) because the theta keys replace .weight and the placeholder .weight keys are absent (_strict_state_dict_loading = False).
5. Call precompute_weff_outside_graph(state_dict, self.dtype) when use_precomputed_weff_outside() is true. This is the precompute-outside fast path.
6. Compile-cache the graph as .mef via load_or_compile / cached_mef_path keyed on config_dims, so repeat runs skip the multi-minute graph build.

The per-model diff: Gemma 4 handles the multimodal wrapper and builds MultiKVCacheParams (separate sliding and global KV caches); Gemma 3 uses a single KV cache.

Swap in the attnquant layers

Copy models/gemma3/gemma3_attnquant.py. Subclass the MAX text model and attention layer, replace the seven Linear projections with AttnQuantLinear (via a linear_cls partial) and the MLP with AttnQuantMLP (attn-only mode leaves the MLP dense), declare the W0 Weights at model level (self.W0[band][fam]), and inject them with set_w0() before every forward. The layer reads band_index(), and for Gemma 4 also layer_type, to pick the right W0.

Everything non-linear stays MAX-native: norms, RoPE, embeddings, attention scoring, and the sliding-window mask. The per-model diffs, taken from the file docstrings:

• Gemma 3: Gemma3RMSNorm with weight_offset=1, dual RoPE (global/local), sliding_window_pattern=6, four norms per block, single KV cache.
• Gemma 4: dual head dims (256 sliding / 512 full), attention_k_eq_v (fused qk_proj with 2 children versus qkv_proj with 3), ProportionalRotaryEmbedding for global layers, Gemma4RMSNorm with weight_offset=0, MultiKVCacheParams, and a per-layer layer_scalar.

Route the MLP through config.hidden_activation

AttnQuantMLP must apply config.hidden_activation, never a hardcoded activation. Gemma 3 and Gemma 4 both use gelu_tanh, not silu. The activation function is resolved from the config in core/attnquant_mlp.py:

# core/attnquant_mlp.py resolves the model's hidden_activation
gate_out = self.activation_function(gate_out)

When you bring up a new arch, confirm its hidden_activation flows through the config to AttnQuantMLP. A wrong activation breaks output quality even when every weight is correct.

Per-model checklist

Before serving, enumerate these for the new model:

• Projection families present: the usual seven (q/k/v/o/gate/up/down), or a subset.
• Head dims: uniform (Gemma 3) or dual (Gemma 4). Dual head dims require the layer_type W0 grouping in weight_adapters.py.
• k == v tying: if the model reuses K for V on some layers (attention_k_eq_v), drop v_proj on those layers in the adapter.
• MLP intermediate_size: must match the checkpoint.
• hidden_activation: set correctly (gelu_tanh for Gemma 3/4), per the activation rule above.

Validate

Validate in three tiers, cheapest first.

Analytic: Frobenius error at passthrough

Run the decomposition tests and confirm the reconstruction error is approximately zero at the passthrough operating point (G = num_layers):

pixi run pytest tests/test_decompose.py

The relative_frobenius_error() should be ~0 at G = num_layers for every rank. At Gemma 3’s G=26 the measured error is 0 across ranks 4-256.

Parity and quality: serve tests

Run the parity and quality suites, which check logit-parity and greedy-decode match against the MAX built-in model:

pixi run pytest tests/test_quality.py -v -m serve -s

Logit-parity means the attnquant model’s output logits match the built-in baseline within a max logprob delta below 1.0 (native-equivalent, not bit-exact). The custom-op path also carries kernel correctness and parity tests (tests/test_kernel.py, tests/test_parity.py, plus tests/test_gemma4_cross.py for Gemma 4). See quality for what the thresholds mean.

Smoke: generate at passthrough

Finally, generate at the passthrough operating point and confirm sensible output:

pixi run python run_max.py generate \
  --model <repo> \
  --custom-architectures attnquant \
  --prompt 'The capital of France is' \
  --max-new-tokens 16

At G = num_layers Gemma 3 generates correct English (The capital of France is Paris.), and Gemma 4 passes greedy decode, parity, and quality. If the smoke test produces coherent text, the wiring is correct.