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 struct

Matmul2dFp4

struct Matmul2dFp4[c_type: DType = DType.float32, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, num_sg_m: Int = Int(2), num_sg_n: Int = Int(2), tm: Int = Int(2), tn: Int = Int(2), smem_bk: Int = Int(16)]

W4A16 GEMM on the native 16x32x16 matmul2d MMA: coalesced NT FP4 decode.

Each 16x32x16 MMA is two native _mma_apple_transposable (air.simdgroup_matrix_16x16x16) calls (matmul2d_mma_regc_bt_native) -- pure Mojo over _mma_apple, no external library or codegen dependency.

Parameters

c_type (DType): Output element type (fp16, bf16, fp32). Accumulation is fp32.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional fused epilogue (AMD's (row, col) contract). Currently unused on the interior store; wired for parity.
num_sg_m (Int): Simdgroup rows per threadgroup.
num_sg_n (Int): Simdgroup cols per threadgroup.
tm (Int): Output 16x32 tiles per simdgroup along M (register-C accumulators).
tn (Int): Output 16x32 tiles per simdgroup along N. tm*tn == 4 is the M5 register-cliff optimum; do NOT exceed 4 (measured collapse past it).
smem_bk (Int): Cooperative-decode strip depth for run_smem_decode (the bf16 weight sub-tile columns decoded per phase; one decode feeds smem_bk//16 MMA K-steps). Unused by the per-lane run kernel. The production launcher pins 256, the M5 optimum: a deeper strip amortizes the decode over more MMA work; 512 spills.

Implemented traits

AnyType, ImplicitlyDeletable

`comptime` members

`BK`

comptime BK = smem_bk

`BN`

comptime BN = Matmul2dFp4[c_type, elementwise_lambda_fn, num_sg_m, num_sg_n, tm, tn, smem_bk].TG_N

`BYTES_PER_THREAD`

comptime BYTES_PER_THREAD = (Matmul2dFp4[c_type, elementwise_lambda_fn, num_sg_m, num_sg_n, tm, tn, smem_bk].PACKED_COLS // (Int((mul _resolve_warp_size(), num_sg_m, num_sg_n)) // Int((mul num_sg_n, tn, 32))))

`COLS_PER_THREAD`

comptime COLS_PER_THREAD = (((smem_bk // Int(2)) // (Int((mul _resolve_warp_size(), num_sg_m, num_sg_n)) // Int((mul num_sg_n, tn, 32)))) * Int(2))

`in_type`

comptime in_type = DType.bfloat16

`MMA_K`

comptime MMA_K = M2D_K

`MMA_M`

comptime MMA_M = M2D_M

`MMA_N`

comptime MMA_N = M2D_N

`NUM_SG`

comptime NUM_SG = (num_sg_m * num_sg_n)

`PACKED_COLS`

comptime PACKED_COLS = (Matmul2dFp4[c_type, elementwise_lambda_fn, num_sg_m, num_sg_n, tm, tn, smem_bk].BK // Int(2))

`SF`

comptime SF = NVFP4_SF_VECTOR_SIZE

`TG_M`

comptime TG_M = ((Int(16) * num_sg_m) * tm)

`TG_N`

comptime TG_N = ((Int(32) * num_sg_n) * tn)

`THREADS_PER_BLOCK`

comptime THREADS_PER_BLOCK = (Matmul2dFp4[c_type, elementwise_lambda_fn, num_sg_m, num_sg_n, tm, tn, smem_bk].NUM_SG * _resolve_warp_size())

`THREADS_PER_ROW`

comptime THREADS_PER_ROW = (Int((mul _resolve_warp_size(), num_sg_m, num_sg_n)) // Matmul2dFp4[c_type, elementwise_lambda_fn, num_sg_m, num_sg_n, tm, tn, smem_bk].BN)

Methods

`run`

static def run[c_layout: TensorLayout, a_layout: TensorLayout, packed_layout: TensorLayout, scale_layout: TensorLayout](c: TileTensor[c_type, c_layout, MutAnyOrigin], a: TileTensor[DType.bfloat16, a_layout, ImmutAnyOrigin], packed: TileTensor[DType.uint8, packed_layout, ImmutAnyOrigin], scales: TileTensor[DType.float8_e4m3fn, scale_layout, ImmutAnyOrigin], M: Int, N: Int, K: Int)

W4A16 kernel entry. C (M, N), A (M, K) bf16, packed (N, K//2), scales (N, ceil(K/16)). Interior fast path (tile-aligned N, K%16==0).

Parameters:

c_layout (TensorLayout): Layout of the output C (M, N) view.
a_layout (TensorLayout): Layout of the activation A (M, K) view.
packed_layout (TensorLayout): Layout of the packed FP4 weight (N, K//2) view.
scale_layout (TensorLayout): Layout of the block scales (N, ceil(K/16)) view.

Args:

c (TileTensor[c_type, c_layout, MutAnyOrigin]): Output TileTensor view with shape (M, N).
a (TileTensor[DType.bfloat16, a_layout, ImmutAnyOrigin]): Bf16 activation TileTensor view with shape (M, K).
packed (TileTensor[DType.uint8, packed_layout, ImmutAnyOrigin]): Packed FP4 weight TileTensor view with shape (N, K//2) (uint8, lo-nibble first).
scales (TileTensor[DType.float8_e4m3fn, scale_layout, ImmutAnyOrigin]): FP8-E4M3 block-scale TileTensor view with shape (N, ceil(K/16)).
M (Int): Number of rows in the activation and the output.
N (Int): Number of columns in the output and rows of the weight.
K (Int): Reduction dimension; inner size of a and the weight.

`run_smem_decode`

static def run_smem_decode[c_layout: TensorLayout, a_layout: TensorLayout, packed_layout: TensorLayout, scale_layout: TensorLayout](c: TileTensor[c_type, c_layout, MutAnyOrigin], a: TileTensor[DType.bfloat16, a_layout, ImmutAnyOrigin], packed: TileTensor[DType.uint8, packed_layout, ImmutAnyOrigin], scales: TileTensor[DType.float8_e4m3fn, scale_layout, ImmutAnyOrigin], M: Int, N: Int, K: Int)

Cooperative-decode W4A16: coalesced FP4 decode -> SMEM -> register B.

The run kernel decodes per-lane, but each lane's transpose_right B fragment reads 4 N-scattered rows (stride-K/2 in the packed (N, K//2) weight) -> the packed + scale loads are uncoalesced gathers and dominate (measured: they drop the ~60 TF/s pure-MMA feed to ~10). This kernel COOPERATIVELY decodes the (BN, BK) weight sub-tile into SMEM with ADJACENT threads reading ADJACENT packed bytes (coalesced K), then each simdgroup reads its tn register B fragments from SMEM via the transpose_right map (bt_frag_coord) and runs the native 16x32x16 MMA (matmul2d_mma_regc_bt_native, two _mma_apple_transposable calls). This is the committed AppleM5Fp4MatMul cooperative-decode idea on the matmul2d tile shape -- the MMA stays REGISTER-fed (a plain SMEM->register indexed load).

Interior fast path (tile-aligned N, K % BK == 0). A loads register-direct from DRAM (base map, K contiguous). Bit-exact vs run / the materialize oracle (same decode_e2m1_to_f32 * |scale|, same MMA), modulo the fp32 MMA reduction order (identical to run, so bit-exact vs run).

Parameters:

c_layout (TensorLayout): Layout of the output C (M, N) view.
a_layout (TensorLayout): Layout of the activation A (M, K) view.
packed_layout (TensorLayout): Layout of the packed FP4 weight (N, K//2) view.
scale_layout (TensorLayout): Layout of the block scales (N, ceil(K/16)) view.

Args:

c (TileTensor[c_type, c_layout, MutAnyOrigin]): Output TileTensor view with shape (M, N).
a (TileTensor[DType.bfloat16, a_layout, ImmutAnyOrigin]): Bf16 activation TileTensor view with shape (M, K).
packed (TileTensor[DType.uint8, packed_layout, ImmutAnyOrigin]): Packed FP4 weight TileTensor view with shape (N, K//2) (uint8, lo-nibble first).
scales (TileTensor[DType.float8_e4m3fn, scale_layout, ImmutAnyOrigin]): FP8-E4M3 block-scale TileTensor view with shape (N, ceil(K/16)).
M (Int): Number of rows in the activation and the output.
N (Int): Number of columns in the output and rows of the weight.
K (Int): Reduction dimension; inner size of a and the weight.

Parameters​

Implemented traits​

comptime members​

BK​

BN​

BYTES_PER_THREAD​

COLS_PER_THREAD​

in_type​

MMA_K​

MMA_M​

MMA_N​

NUM_SG​

PACKED_COLS​

SF​

TG_M​

TG_N​

THREADS_PER_BLOCK​

THREADS_PER_ROW​

Methods​

run​

run_smem_decode​