For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo module

fp4_matmul

Apple M5 weight-only NVFP4 (W4A16) matmul: out = activation @ weight^T.

Apple silicon GPU (Metal 4, compute_capability == 5). The FP4 weight stays packed in DRAM (uint8 [N, K//2] lo-nibble + float8_e4m3fn [N, K//16] block scales); each (BN, BK) weight sub-tile is COOPERATIVELY dequantized to bf16 once per K-strip into threadgroup memory (SMEM), then the dense bf16 simdgroup MMA reads B from SMEM. No FP4/FP8 MMA hardware is used -- the dequant target is bf16 and the MMA is the dense bf16 path.

A-vs-B decision: in out = x @ W^T the FP4 weight is the B operand (transpose_b=True, W is [N, K]). The SMEM B sub-tile is stored row-major (BN, BK) -- the same lane->element order the dense transpose_b=True path reads -- so feeding it to mma() reproduces the dense hw_transpose_b exactly. The activation A is the dense bf16 operand, loaded from DRAM each K-strip.

WHY cooperative-SMEM (not inline-per-fragment dequant): the cost of the W4A16 path over the dense bf16 path is entirely the FP4 dequant (the MMA + epilogue are identical to dense). The earlier inline kernel dequantized each B fragment in the MMA loop; this kernel decodes the whole (BN, BK) sub-tile ONCE per strip, cooperatively across the threadgroup, so one decode phase feeds BK//16 K-steps of MMA. At BK=32 (one cooperative decode feeds two 16x16 K-steps) this beats the inline kernel at every measured shape, bit-exact -- see benchmarks/gpu/linalg/bench_apple_fp4_smem_bk.mojo and the kernel docstring. The dequant phase is a barrier-fenced SCALAR-ALU compute phase that the M5 cannot overlap with the MMA via double-buffering (measured: double-buffer == single-buffer for FP4); the win is amortizing the un-overlapped decode over more MMA work, NOT hiding it. The double buffer is kept because it is free here (the next strip's decode is issued before the current strip's MMA) and matches the dense pipelining idiom.

This kernel mirrors AppleM5MatMul._run_gemm_body for everything but the B side (Morton scheduler, per-simdgroup 32x32 / 2x2 tiling, edge-tile bounds, K-strip + tail loop, cast/epilogue store). The dense fp16/bf16/fp32 path is untouched.

Structs

AppleM5Fp4MatMul: W4A16 GEMM: bf16 activation x packed-FP4 weight, cooperative-SMEM dequant.

Functions

enqueue_apple_fp4_matmul: Enqueue the W4A16 matmul: out = a @ dequant(packed, scales)^T.

Structs​

Functions​

Structs

Functions