Gemma 3 and 4 architectures

API reference for attnquant_gemma3_arch and attnquant_gemma4_arch, the SupportedArchitecture objects that register Gemma 3 and Gemma 4 as attnquant custom architectures in MAX.

attnquant_gemma3_arch and attnquant_gemma4_arch are the two SupportedArchitecture objects that make attnquant Gemma 3 and Gemma 4 selectable from MAX. A SupportedArchitecture is a MAX registry entry that binds a model name to the pipeline model, config, tokenizer, and weight adapter MAX uses to serve it. These two objects wire the attnquant pipeline model and the decomposing weight adapter into MAX’s normal serving path.

This page documents both objects as an API surface: the fields each one wires, the same-name override that lets a custom architecture replace a MAX built-in, and the band_index() config helper. To register a new model of your own, see extending.

attnquant_gemma3_arch is defined in models/gemma3/arch.py; attnquant_gemma4_arch is defined in models/gemma4/arch.py. The Gemma 3 module also registers a third entry, attnquant_gemma3_multimodal_arch (name="Gemma3ForConditionalGeneration", for the multimodal Gemma 3 checkpoints such as google/gemma-3-4b-it and google/gemma-3-12b-it, with device graph capture disabled); this page focuses on the two text-generation architectures.

Fields

Both objects are constructed from the same SupportedArchitecture fields. This table lists the values that differ between the two:

Both objects share these values:

The three fields that carry the attnquant behavior are pipeline_model, config, and weight_adapters:

• pipeline_model: builds the graph with the attnquant text model and the reconstruction kernel wired in as a custom extension. AttnQuantGemma3PipelineModel overrides _build_graph; AttnQuantGemma4PipelineModel subclasses the MAX Gemma 4 multimodal pipeline model and overrides load_model and _build_language_graph.
• config: the standard MAX Gemma config plus the attnquant fields (lora_rank, tie_families, shared_blocks, tied_projections). See Config fields.
• weight_adapters: maps the safetensors format to convert_safetensor_state_dict, which decomposes each dense projection into its theta bundle as the checkpoint loads.

The transformer body (attention, RoPE, normalization, KV cache) is unchanged from MAX’s Gemma implementation. Only the projection layers are replaced with AttnQuantLinear. Because multi_gpu_supported=False, AttnQuantLinear.shard() raises NotImplementedError.

Same-name override

Each name is the Hugging Face architecture class string, and it deliberately matches a MAX built-in: "Gemma3ForCausalLM" and "Gemma4ForConditionalGeneration" are the same names MAX ships. When you serve with --custom-architectures attnquant, MAX imports the package, registers each entry from its ARCHITECTURES list with allow_override=True, and the attnquant architecture takes precedence over the built-in of the same name.

attnquant/__init__.py builds ARCHITECTURES from both model modules, then clears MAX’s lazy registration entry for each overridden name:

from max.pipelines.lib.registry import PIPELINE_REGISTRY

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

Without this cleanup, MAX’s later lazy materialization calls register(builtin_arch) without allow_override and fails with Refusing to override existing architecture. For the reasoning behind the override, see extending.

To serve Gemma 3 1B (Gemma3ForCausalLM, 26 layers) through the attnquant path:

ATTNQUANT_BLOCKS=26 ATTNQUANT_PRECOMPUTE_OUTSIDE=1 \
    python run_max.py serve --model google/gemma-3-1b-it \
    --custom-architectures attnquant --port 3100

To generate from Gemma 4 31B (Gemma4ForConditionalGeneration, 60 layers: 10 full-attention plus 50 sliding-window):

ATTNQUANT_BLOCKS=60 ATTNQUANT_PRECOMPUTE_OUTSIDE=1 \
    python run_max.py generate --model google/gemma-4-31B-it \
    --custom-architectures attnquant --prompt "..." --max-new-tokens 128

Config fields and band_index()

AttnQuantGemma3Config subclasses MAX’s Gemma3Config and AttnQuantGemma4Config subclasses Gemma4TextConfig, so each inherits every standard Gemma field and adds four attnquant fields:

• lora_rank (default 32): rank r of the per-(layer, family) low-rank delta U, V. Higher rank improves approximation quality at the cost of more shipped bits. Overridden by ATTNQUANT_RANK.
• tie_families (default "all"): which projection families are attnquant. "all" makes every q/k/v/o/gate/up/down projection attnquant; "attn" makes only q/k/v/o attnquant and builds the MLP projections as plain dense max.nn.Linear.
• shared_blocks (default 1): the number of contiguous layer bands G that each own a distinct shared W0 per family. G=1 is a single W0 per family.
• tied_projections (default all seven families): the families implied by tie_families, derived at initialization for inspection. The model recomputes this value from tie_families when building the graph, so setting it directly has no effect.

These fields are read from config.json with backward-compatible defaults, so an unmodified Gemma repo loads without error. For the full defaults, environment overrides (ATTNQUANT_RANK, ATTNQUANT_BLOCKS, and the rest), and precedence rules, see runtime configuration.

band_index() returns the layer-to-band map used to inject each band’s W0. It partitions [0, num_hidden_layers) into shared_blocks contiguous bands with sizes as equal as possible (earlier bands take the +1 remainder), and returns a list whose length is num_hidden_layers and whose value at index i is the band index of layer i. The config field shared_blocks, the environment override ATTNQUANT_BLOCKS, and the math symbol G are the same value.

For Gemma 3’s 26 layers:

config.band_index()   # shared_blocks=1 -> [0]*26
                       # shared_blocks=2 -> [0]*13 + [1]*13
                       # shared_blocks=3 -> [0]*9 + [1]*9 + [2]*8

For Gemma 4’s 60 layers, band_index() uses the same partition, but W0 routing combines it with layer_types so that full-attention and sliding-window layers never share a W0 group even within one band (their projection dimensions differ because the two layer types use different head dims):

config.band_index()   # shared_blocks=1 -> [0]*60
                       # shared_blocks=2 -> [0]*30 + [1]*30
                       # shared_blocks=3 -> [0]*20 + [1]*20 + [2]*20

At the passthrough operating point (G == num_hidden_layers: one W0 per layer, so each reconstructed weight equals the original), each layer maps to its own band and output is native-equivalent to the built-in path.

Add another model

Registering a new decoder reuses the core/ and kernels/ machinery unchanged and follows a six-file per-model pattern (arch.py, model.py, model_config.py, weight_adapters.py, {model}_attnquant.py, and __init__.py). For the full procedure, see extending.