Serve and generate with weightless

Run a weightless Gemma 3 or Gemma 4 model on MAX through run_max.py, register the custom architecture with --custom-architectures, and set WEIGHTLESS_PRECOMPUTE_OUTSIDE=1 for the fast path on NVIDIA.

A weightless 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. Weightless rebuilds each linear projection’s weight from a compact theta bundle (a shared block W0 per band, a per-row scale s, and a low-rank delta U, V of rank r) instead of loading a dense matrix. You point MAX at a Hugging Face checkpoint, add --custom-architectures weightless_fw to register the architecture, and run everything through the run_max.py wrapper.

This page covers the run commands, how the custom architecture loads over MAX’s built-in, and the flag to set on NVIDIA. For the full flag reference, see configuration.

Run through run_max.py

Invoke every weightless command through run_max.py, not the bare max CLI. On some machines the max._core C extension resolves the Mojo import path to a stale location from a previous install and writes it into the environment, overriding the pixi SDK. run_max.py imports max._core to trigger that behavior, strips the stale MODULAR_* variables so the pixi defaults take effect, adds the repo root to sys.path so import weightless_fw resolves, then calls the MAX entrypoint. The shape of every command is:

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

Generate a completion

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

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

Serve an endpoint

To bring up an OpenAI-compatible endpoint, pass a port. On NVIDIA, set WEIGHTLESS_PRECOMPUTE_OUTSIDE=1 (see Set the fast path on NVIDIA):

WEIGHTLESS_PRECOMPUTE_OUTSIDE=1 \
  pixi run python run_max.py serve \
    --model google/gemma-3-1b-it \
    --custom-architectures weightless_fw \
    --port 3100 \
    --max-length 8192 \
    --device-memory-utilization 0.9

The Gemma 4 31B checkpoint (Gemma4ForConditionalGeneration, 60 layers: 10 full-attention plus 50 sliding-window) serves through the same command with a different --model:

WEIGHTLESS_PRECOMPUTE_OUTSIDE=1 \
  pixi run python run_max.py serve \
    --model google/gemma-4-31B-it \
    --custom-architectures weightless_fw \
    --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 weightless_fw, MAX imports the package and reads its ARCHITECTURES list. The weightless Gemma 3 registers under name="Gemma3ForCausalLM" and the weightless 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 (convert_safetensor_state_dict) auto-decomposes the Hugging Face checkpoint into the theta bundle. 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. Two math-identical reconstruction paths exist: the generated_linear Mojo custom op, and a GEMM path built from stock max.graph matmuls. Weightless auto-selects between them by accelerator and rank (GEMM on NVIDIA and other non-Metal devices, the custom op on Metal at low rank). You don’t choose the path by hand for a normal run.

Set the fast path on NVIDIA

On NVIDIA GPUs, set one environment variable for the fastest decode:

WEIGHTLESS_PRECOMPUTE_OUTSIDE=1

With WEIGHTLESS_PRECOMPUTE_OUTSIDE=1, weightless 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 MAX built-in (the same model served on MAX’s built-in dense path, with normal dense weights and no weightless decomposition), with Q/K/V fused into one matmul. Leave WEIGHTLESS_GEMM unset so device-aware selection picks the GEMM path on NVIDIA.

These settings give native-equivalent output at the passthrough operating point (G = num_layers: one W0 per layer, so each reconstructed weight equals the original). Parity is verified by logit-parity: the reconstructed model’s logits match the original within a max logprob delta below 1.0, with greedy/argmax token match above 80-90%. This is not bit-exact.

On one NVIDIA B200 serving Gemma 3 1B, weightless with the fast path measured:

• Single-stream decode (batch 1, one request at a time): 439 vs 392 tok/s (+12%).
• Concurrent throughput at c=32 (32 requests in flight): 68.0 vs 63.8 req/s (+6.5%, median of 3, steady-state, non-overlapping).

The win comes from Q/K/V fusion: one matmul instead of three. The cost is memory: precompute-outside keeps both the full W_eff per layer and the theta bundle resident, so peak GPU memory is about 1.9x the MAX built-in baseline (32.7 vs 17.6 GB). See throughput for the full scorecard.

Supported encodings and devices

Both weightless architectures default to bfloat16 and also support float32. Both run on a single device: multi_gpu_supported is false, because tensor parallelism doesn’t yet compose with reconstructing each weight from its theta bundle, and WeightlessLinear.shard() raises NotImplementedError.

Current limits

Two limits shape what you can deploy today:

• There’s no compressed operating point yet. Shared-W0 configurations (G < num_layers) produce incoherent output and need a training or distillation step, so passthrough is the only coherent setting today, and it stores a full weight per layer.
• Tensor parallelism is unsupported (multi_gpu_supported=False), and the Gemma 4 31B throughput head-to-head is still pending.

Next steps

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

• Configuration: the canonical WEIGHTLESS_* environment variables and config.json fields, including WEIGHTLESS_RANK, WEIGHTLESS_BLOCKS, and the precompute-outside details.
• 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.