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
matmul2d_fp4
Apple M5 weight-only NVFP4 (W4A16) matmul on the 16x32x16 matmul2d tile.
Target: Apple M5 (compute_capability == 5) only.
out = activation @ weight^T where the FP4 weight stays PACKED in DRAM
(uint8 [N, K//2] lo-nibble first + float8_e4m3fn [N, K//16] block scales).
This is the 16x32x16 matmul2d-tile sibling of the committed AppleM5Fp4MatMul
(which runs on the narrower simdgroup_matrix 16x16x16 MMA). Both paths coexist;
this one uses the wider 16x32x16 tile with the coalesced transpose_right=1 NT
feed. "matmul2d" here names the 16x32x16 tile SHAPE (Apple's term), implemented
natively as two air.simdgroup_matrix_16x16x16 MMAs -- one per 16-wide N half --
via the _mma_apple_transposable stdlib intrinsic (matmul2d_mma_regc_bt_native).
Pure Mojo over _mma_apple: no external library or codegen dependency, so it
builds on a stock toolchain.
Why this is fast: the coalesced NT decode
The FP4 weight is the RIGHT operand with transpose_b=True (W is (N, K)).
On the base matmul2d swizzle an NT B fragment is a stride-K scattered gather
(31 TF/s). The transpose_right=1 op reinterprets the right operand as (N, K)
so the B-fragment's fast axis is K -- a coalesced, K-contiguous load
(bt_frag_coord, below). For each B fragment this kernel decodes the packed FP4
weight DIRECTLY into the register fragment, K-contiguous:
- Per lane, the
transpose_rightB fragment reads 4 N-rows (rb, rb+8, rb+16, rb+24), each contributing 4 CONTIGUOUS K values (cb, cb+1, cb+2, cb+3,cba multiple of 4). Those 4 K nibbles are exactly ONE 2-byte packed load (packed[n, (k0+cb)//2]asuint16), and (becausecb+3 < 16) share ONE FP8 block scale (scales[n, (k0+cb)//16]). - Decode is the FTZ-safe branch-free
decode_e2m1_to_f32(bit-identical to theE2M1_TO_FLOAT32LUT the materialize oracle uses; the exponent-injection trick is WRONG on M5's flush-to-zero -- seefp4_utils.mojo), scaled by|block_scale|in f32, cast to bf16 -> the B fragment. So this path is BIT-EXACT vs the materialize -> dense oracle (same dequant arithmetic, same bf16matmul2dMMA).
No SMEM, no barrier (Apple's idiom): DRAM packed weight decoded straight to the
register B fragment, MLX NAXTile register-C accumulators (tm=tn=2, 4
accumulators -- the M5 register-cliff optimum, do NOT exceed). Interior-only fast
path; the launcher rounds the grid up and the store guards < M/< N, so the
last partial tile writes only in-bounds. The packed-weight and scale loads on the
last K strip read up to the tile boundary; callers pass tile-aligned N and
K % 16 == 0 (the W4A16 production shapes are aligned; a bounded edge/tail variant
mirrors AppleM5Fp4MatMul if needed).
The NVFP4 per-tensor weight_scale_2 scalar folds in OUTSIDE the kernel (a
post-matmul multiply by the graph lowering), identically to the committed path.
comptime values
M2D_K
comptime M2D_K = 16
M2D_M
comptime M2D_M = 16
M2D_N
comptime M2D_N = 32
Structs
-
Fp4WeightLoader: Owner of the packed-FP4 weight -> dequant -> {register B | SMEM} transition. -
Matmul2dFp4: W4A16 GEMM on the native 16x32x16matmul2dMMA: coalesced NT FP4 decode.
Functions
-
a_frag_coord: (row, col) in the 16x16 A tile for this lane's A elementi(0..7). -
bc_frag_coord: (row, col) in the 16x32 B/C tile for this lane's elementi(0..15). -
bt_frag_coord: (n, k) in the 16x32 right operand for elementiundertranspose_right=1. -
enqueue_matmul2d_fp4: Enqueue thematmul2dW4A16 GEMM:out = a @ dequant(packed, scales)^T. -
enqueue_matmul2d_fp4_smem: Enqueue the cooperative-decodematmul2dW4A16 GEMM (run_smem_decode). -
frag_col_base: Base MMA-tile column forlane(before the per-element col offset). -
frag_row_base: Base MMA-tile row forlane(before the per-element row jump). -
matmul2d_mma_regc_bt_native: Pure-Mojotranspose_right=116x32x16 MMA (nativesimdgroup_matrix).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!