In-kernel reconstruction

Compare the two math-identical paths that rebuild each weight at inference time, the generated_linear custom op and the stock-GEMM formulation.

A weightless layer never runs the standard x @ weight.T matmul, because the weight doesn’t exist in memory. Instead it computes the reconstruction and the projection together:

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

expanded per element as:

y[t, o] = s[o] · Σᵢ W0[o,i] · x[t,i]
        + Σₖ U[o,k] · Σᵢ V[i,k] · x[t,i]

The framework provides two implementations of this identical math. Which one runs is chosen automatically from the accelerator and the LoRA rank, and can be forced with an environment variable.

The custom-op path

The default at low rank reconstructs each weight inside the generated_linear Mojo custom op. It has CPU, Metal, and NVIDIA GPU (B200/SM100) implementations, all accumulating in float32 and casting back to the output dtype.

On CPU and Metal, the kernel uses one block per output element with a shared-memory tree reduction. On NVIDIA GPUs (B200), it uses an optimized kernel with vectorized 128-bit loads (bf16x8 / fp32x4), warp-level shuffle reduction, one warp per output row, cooperative vtx computation for the LoRA delta, and a single block-level barrier.

This path wins at low rank (for example rank = 4): the per-element reconstruction is cheap, and there’s no benefit from routing through large GEMMs.

The MLP can also route through a single fused generated_mlp custom op that combines the gate, up, activation, and down projections into one dispatch. This is on by default; set WEIGHTLESS_FUSED_MLP=0 to fall back to three separate generated_linear calls.

The GEMM path

The same output can be expressed entirely with stock max.graph operations, which MAX lowers to its optimized device GEMM kernels:

y = s[None, :] · (x @ W0.T) + (x @ V) @ U.T

The matmuls run in the input dtype (bf16 on GPU, which accumulates in float32 internally via Tensor Cores); only the final elementwise combine is done in float32 before the result is cast back to x.dtype. This path wins decisively at higher rank (rank = 16, rank = 32), where the launch-bound custom-op dispatch pattern costs more than a handful of large matmuls.

How the path is selected

Selection is controlled by WEIGHTLESS_GEMM, resolved in this order:

Explicit value wins

WEIGHTLESS_GEMM=1 (or true/yes/on) forces the GEMM path; 0 (or false/no/off) forces the custom op. An explicit value always overrides the auto-select.

Unset on a non-Metal accelerator

On NVIDIA GPUs (and any non-Metal accelerator), the GEMM path is the default at any rank. It lowers to stock device GEMM kernels that win at every rank measured on the B200, where the custom op’s per-element reconstruction is 10x slower at rank 32. See throughput.

Unset on Metal

The GEMM path is used when lora_rank >= 8 (the crossover measured on Apple Metal) and the custom op below that, where its per-element reconstruction is cheaper.

Unset and accelerator unknown

Falls back to the Metal rank crossover.

The accelerator is detected once via max.driver.accelerator_api, and the auto-selection decision is logged once at startup.

Q/K/V fusion

In the GEMM path only, the three attention projections that share the same input x can be batched into two dispatches: one fused GEMM (the shared W0 block and the LoRA down-projection concatenated along the input axis) and one block-diagonal LoRA up-projection GEMM — instead of separate matmuls per Q/K/V child. This fusion is on by default; set WEIGHTLESS_QKV_FUSE=0 to disable it. The custom-op path never fuses.

See runtime configuration for the full list of environment variables.