IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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, Deinitable, Movable

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_arg: Int32, N_arg: Int32, K_arg: Int32)

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:

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_arg: Int32, N_arg: Int32, K_arg: Int32)

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:

Models

Cookbook

Puzzles

Blog

Community

Contact