Weightless inference for MAX

Decompose pretrained weights into a compact theta bundle and reconstruct them inside GPU kernels at inference time, so no full weight matrix lives in DRAM.

Weightless is an inference framework for MAX that never stores a model’s full weight matrices. It decomposes each linear projection into a small theta bundle (a shared block W0, a per-layer scale s, and a low-rank delta U, V), then reconstructs the original weight on the fly inside a Mojo custom op:

y = s * (W0 x) + U (Vᵀ x)

The [out_dim, in_dim] weight is never materialized. The framework registers as a MAX custom architecture, so you serve a weightless model with the same max generate and max serve commands you already use. Only the weights are handled differently.

How it fits together

The decomposition is analytic, so there’s no training step. The framework transforms a standard HuggingFace checkpoint into theta at load time (per-row least-squares scale plus a truncated SVD of the residual), and each WeightlessLinear layer reconstructs its weight in-kernel during generation.

What’s supported today

• Gemma 3 1B (Gemma3ForCausalLM) as a weightless custom architecture.
• Gemma 4 31B (Gemma4ForConditionalGeneration) as a weightless custom architecture — research/parity only; see Gemma 4 quality (no compressed operating point yet).
• CPU, Metal, and NVIDIA GPU (B200 / SM100) execution via the generated_linear Mojo kernel and a stock-GEMM path.
• Two math-identical reconstruction paths (an in-kernel custom op and a stock-GEMM formulation), auto-selected by accelerator and LoRA rank.
• max generate and max serve on osx-arm64 and Linux NVIDIA GPUs.