GPU and cache acceleration

How weightless achieves fast decode on GPU — CUDA graph capture, QKV fusion, and NVFP4 quantization — with measured throughput on NVIDIA B200.

Weightless delivers fast GPU decode through three mechanisms: CUDA graph capture (eliminates launch overhead), QKV fusion (one matmul where vanilla MAX issues three), and NVFP4 quantization (4x weight compression). This document explains each mechanism and traces the throughput numbers to measured data on an NVIDIA B200.

Related: Banded sharing | Reconstruction | Throughput

CUDA graph capture

At batch-1 decode, a transformer model issues 168-504 kernel dispatches per token. On B200, each launch costs ~5-10 us of host overhead, so 300+ launches = ~1.5-3 ms of pure dispatch overhead per token before any compute. CUDA graph capture eliminates this entirely by replaying the entire decode step as a single graph.

Capture is enabled on both supported models:

• Gemma 3 1B: 2.45x speedup (282 to 691 tok/s batch-1 decode), 4.6x lower TPOT. Shipped in commit 0394499.
• Gemma 4 31B: enabled with supports_device_graph_capture=True in models/gemma4/arch.py. Generates at 66.9 tok/s with 98.3% GPU utilization.

QKV fusion

The weightless precompute-outside path fuses the query, key, and value projections into a single matmul. Vanilla MAX issues three separate matmuls for attention projections; weightless issues one. This reduces both dispatch count and HBM bandwidth (one weight read instead of three).

The throughput impact, measured on B200 with capture enabled:

On the 1B model, QKV fusion gives 12% over vanilla MAX. On the 31B model in bf16, the matmuls are large enough that the extra launches are negligible, so weightless is at parity. NVFP4 quantization changes the calculus: with 4x smaller weights, the HBM bandwidth saved by QKV fusion is a larger fraction of total, giving 10%.

NVFP4 quantization

NVFP4 quantization (4-bit weights with per-block float8 scales) reduces HBM bytes per projection 4x. The weightless precompute-outside path quantizes W_eff to NVFP4 at load time (WEIGHTLESS_QUANT=nvfp4), storing packed uint8 weights + float8 block scales instead of bf16. At passthrough (G=60), W_eff = W0 = the original weight, so the quantized path is mathematically identical to vanilla NVFP4 — but the weightless graph adds QKV fusion and capture on top.

Measured on Gemma 4 31B (passthrough, capture enabled):

Weightless NVFP4 is 10% faster than vanilla NVFP4 and 1.44x over vanilla bf16 (120.0 vs 83.4 tok/s). The roofline predicts 2.50x (matmul is 83.5% of GPU time, NVFP4 is 4x fewer HBM bytes); the gap from 1.44x to 2.50x is dequantization overhead and suboptimal batch-1 GEMV kernel efficiency. A specialized NVFP4 GEMV kernel that dequantizes in registers could close this gap.

The precompute takes _{22 min for 410 projections — a one-time load cost. A cache mechanism saves the quantized state dict to disk for fast subsequent loads (}3 min from cache). Quantized weights are 4x smaller (RSS drops from 246GB to 76GB after precompute).

Profiling: where GPU time goes

nsys profiling on the B200 shows the kernel breakdown:

On the 31B model, matmul kernels dominate at 83.5% of GPU time. These are HBM-bandwidth-bound GEMVs reading weights from HBM. The GPU is 98.3% utilized — this is a bandwidth problem, not a compute problem. NVFP4 quantization directly attacks this bottleneck by reducing HBM bytes 4x.

Architecture decisions

Precompute-outside (the fast path)

WEIGHTLESS_PRECOMPUTE_OUTSIDE=1 precomputes W_eff = s*W0 + U@V.T at load time, reducing each projection to a single vanilla matmul. The graph is identical to the base model, so MAX’s own optimized B200 GEMM kernels apply. This is the recommended decode path.

GEMM formulation (active, committed)

The GEMM path (WEIGHTLESS_GEMM=1 or auto-selected on non-Metal GPU) uses stock ops.matmul which MAX lowers to its own optimized B200 GEMM kernels with tensor cores. Recent improvements include 128-byte alignment padding on fused W0.T+V matrix, bf16 elementwise combine, and device-resident theta weights.

Custom-op kernel (for low-rank Metal path)

The generated_linear.mojo custom op fuses the W0 dot product and LoRA delta in a single kernel with no intermediate materialization. On Metal at low rank (r=4), this wins because the per-element reconstruction is cheaper than 2-3 stock matmuls. On B200, the GEMM path wins at every rank because MAX’s GEMM kernels use tensor cores efficiently.

Future work

• Gate/up fusion: Add gate/up fusion to gemm_formulation.py following the same pattern as QKV fusion. Fewer, larger matmuls improve compute efficiency even with capture on.
• Better NVFP4 GEMV kernel: A specialized kernel that dequantizes 4-bit weights in registers (not SMEM) and streams at full HBM bandwidth could close the gap from 1.44x to 2.50x roofline. Target: 2.0x+ over bf16.
• Trained weight sharing: With a training step to make the residual low-rank, shared NVFP4 W0 + bf16 LoRA could give 6.6x sharing x 1.29x NVFP4 = 8.5x over bf16.

References

• GEMM formulation: weightless_fw/core/gemm_formulation.py
• Weightless linear: weightless_fw/core/weightless_linear.py
• Quantization: weightless_fw/core/quantize.py
• B200 benchmark: docs/benchmark-results.md