Add a new model architecture

Make any decoder weightless 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 weightless, mirror the six files under models/gemma3/ (or models/gemma4/ for models with mixed attention) and reuse the model-agnostic core/ and kernels/ unchanged.

Weightless replaces each dense projection weight with a theta bundle: a shared block W0, a per-row scale s, and a low-rank delta (U, V) that together rebuild one projection’s weight as y = s·(W0 x) + U·(Vᵀ x), described in the theta bundle. This guide is derived by 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).

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

What you reuse versus what you write

The reconstruction math is model-agnostic and stays untouched. You never edit core/decompose.py, core/weightless_linear.py, core/weightless_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:

The six files, copied from gemma3

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 weightless_fw.models.gemma3.arch import weightless_gemma3_arch

ARCHITECTURES = [weightless_gemma3_arch]

Gemma 4 also re-exports its config class (models/gemma4/__init__.py adds WeightlessGemma4Config) because the multimodal batch processor imports it by name. The package-level weightless_fw/__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)

This is the load-bearing gotcha. 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 Weightless<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 weightless fields

Copy models/gemma3/model_config.py. Subclass the model’s MAX config (Gemma3Config, or Gemma4TextConfig for Gemma 4) and add four weightless fields (lora_rank, tie_families, shared_blocks, tied_projections) plus band_index() and an initialize_from_config() override that reads WEIGHTLESS_RANK and WEIGHTLESS_BLOCKS. Here shared_blocks (the config field), WEIGHTLESS_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_weightless_theta_dict() passthrough check. Pre-converted theta bundles skip decomposition.
4. Run decompose_standard_state_dict_to_weightless(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 section 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. to bare keys and adds MoE router/expert remaps (router.proj.weight to moe_block.gate.gate_score.weight, experts. to moe_block.experts.).

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 Weightless<Model> wrapper whose self.language_model is the weightless 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, where GEMM is general matrix-multiply), 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 path (WEIGHTLESS_PRECOMPUTE_OUTSIDE=1), which materializes the full effective weight W_eff = s·W0 + U·Vᵀ once at load time so each projection is a single matmul, the same launch count as the MAX built-in dense 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 weightless layers

Copy models/gemma3/gemma3_weightless.py. Subclass the MAX text model and attention layer, replace the seven Linear projections with WeightlessLinear (via a linear_cls partial) and the MLP with WeightlessMLP (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

WeightlessMLP must apply config.hidden_activation, never a hardcoded activation. This was a real bug: WeightlessMLP.__call__ previously hardcoded ops.silu, which produced incoherent output even at the passthrough operating point (G = num_layers: one W0 per layer, so each reconstructed weight equals the original and there’s no compression), because Gemma 3 and Gemma 4 both use gelu_tanh. The fix lives in core/weightless_mlp.py, where self.activation_function is resolved from the config:

# core/weightless_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 WeightlessMLP. For the parity impact of this fix, see quality.

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). This is the activation fix above.

Validate

Validate in three tiers, cheapest first.

1. 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: one W0 per layer, so each reconstructed weight equals the original and there’s no compression):

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; below passthrough it rises steeply (0.66 at G=13, rank 4). No coherent-and-compressed operating point exists today. See quality and reconstruction.

2. Parity and quality: serve tests

Run the parity and quality suites, which check logit-parity (output logits match the MAX built-in Gemma within a max logprob delta below 1.0, native-equivalent rather than bit-exact) and greedy-decode match against the MAX built-in model on the B200:

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

The MAX built-in baseline is the same model served on MAX’s built-in dense path: normal dense weights, no weightless decomposition. Logit-parity means the reconstructed model’s output logits match that baseline within a max logprob delta below 1.0, so the output is 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). The stated parity thresholds are max logprob delta < 1.0 and 80-90% greedy/argmax match; concrete measured deltas are pending re-measure with the activation fix.

3. Smoke: generate at passthrough

Finally, generate at the passthrough operating point and confirm coherent output before trusting any compressed config:

pixi run python run_max.py generate \
  --model <repo> \
  --custom-architectures weightless_fw \
  --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 after the activation fix. If the smoke test produces coherent text, the wiring is correct.