Quality (Gemma 4 31B)

Decomposition quality findings for weightless Gemma 4 31B on B200 — why the analytic decomposition does not produce a coherent compressed operating point, and what would be needed to fix it.

The Gemma 4 31B model (google/gemma-4-31B-it) is a 60-layer transformer with dual attention types (10 full-attention + 50 sliding-window) and dual head dimensions (256 for sliding, 512 for full). Its weight shapes are significantly larger than Gemma 3 1B, making the decomposition quality problem more severe.

Activation function fix

A bug in WeightlessMLP.__call__ hardcoded ops.silu as the activation function, but Gemma 3 and 4 use gelu_tanh (configured via config.hidden_activation). The fix routes through self.activation_function, which reads the config. Before the fix, the weightless Gemma 4 model produced gibberish text even in passthrough mode (G=60, one W0 per layer). After the fix, the model produces coherent English: greedy decode, parity, and quality tests all pass on the B200. The decomposition quality problem (below) is separate from this bug and remains the gating issue for compressed operating points.

Architecture and weight shapes

The full and sliding attention layers have different q_proj and o_proj dimensions, so the adapter correctly separates them by layer_type for W0 grouping.

Decomposition quality: same structural failure as Gemma 3 1B

Measured on the real Gemma 4 31B safetensors (relative Frobenius error of diag(s) * W0 + U @ V^T against each true weight, averaged over all layers and families):

At G=1 (all layers share one W0 per family), the reconstruction error is 90.7% at rank 32 — the “reconstructed” weight keeps almost none of the original. Even at rank 256 with 5 bands (12 layers per band), the error is still 60.9%.

Residual energy spectrum

The residual weight - s * W0 is genuinely full-rank. Measured via SVD on the real weights:

The residual has 5376 singular values (the full input dimension). Rank 1024 captures only 62% of the energy, and the spectrum decays slowly — there is no “knee” where a practical rank captures most of the signal.

Per-layer SVD (no sharing, each layer gets its own W0) confirms the same: individual weight matrices also lack low-rank structure. Rank 32 captures 8-15%, rank 1024 captures 55-71%.

ALS optimization barely helps

Alternating Least Squares (ALS) was tested as a potential improvement over the mean W0:

ALS improves the error by only 3.7% (0.849 to 0.812), confirming that the problem is not the choice of W0 but the fundamental lack of low-rank structure in the residual.

Verdict: no coherent, compressed operating point

Same conclusion as Gemma 3 1B: the analytic decomposition has no operating point that is both coherent and compressed for Gemma 4 31B. The gap to a usable error (roughly under 10%) is 6x to 9x even at rank 256. A usable weightless Gemma 4 31B would need:

1. A training step (distillation) that fits the shared blocks and deltas to a distillation objective, rather than the one-shot analytic decomposition.
2. A fundamentally better shared representation — the mean W0 is a poor shared block because Gemma’s per-layer weights do not cluster around their family mean.
3. KV cache low-rank compression as a complementary axis — compressing the KV cache rather than the weights themselves.

Kernel optimizations

GEMM path

The weightless GEMM path has been optimized with two changes:

1. QKV fusion: WeightlessStackedLinear.__call__ now batches the shared-x matmuls across Q/K/V projections into 2 GEMM dispatches (one fused W0+V matmul, one block-diagonal LoRA U matmul) instead of separate matmuls per Q/K/V child. This reduces kernel launch overhead at batch size 1 decode.

2. bf16 matmuls: weightless_gemm and weightless_qkv_gemm now use the input dtype (bf16) for matmuls instead of casting to fp32. On GPU, bf16 matmuls use fp32 accumulation internally via Tensor Cores, so accuracy is preserved while halving memory bandwidth.

These optimizations reduce the 2x overhead vs vanilla dense, but the accuracy problem (decomposition quality) remains the gating issue.

Custom-op path (B200 kernel)

The generated_linear Mojo custom op has been optimized for NVIDIA B200 (SM100) with vectorized 128-bit loads (bf16x8 / fp32x4), warp-level shuffle reduction (warp.sum, lane_group_sum), one warp per output row (8 rows per block), cooperative vtx computation for the LoRA delta, and a single block-level barrier (versus _{25 in the baseline). The kernel passes all 9 correctness tests and 5 parity tests on the B200. On the Gemma 3 1B benchmark, custom-op throughput is unchanged (}146 tok/s) because the kernel is memory-bound at small dimensions; the optimization targets large-dimension workloads (Gemma 4 31B: hidden_size=5376, intermediate_size=21504) where vectorized loads and warp shuffle reduction have more impact.

Reproducing

# Run the weightless model with different G and rank
WEIGHTLESS_BLOCKS=<G> WEIGHTLESS_RANK=<r> pixi run python run_max.py generate \
  --model google/gemma-4-31B-it --custom-architectures weightless_fw \
  --prompt "The capital of France is" --max-new-tokens 48