IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/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. /max/get-started.md).

Mojo struct

AppleM5Fp4MatMul

struct AppleM5Fp4MatMul[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, *, BM: Int = Int(128), BK: Int = Int(64), coalesce_scales: Bool = True]

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

The whole threadgroup cooperatively decodes each (BN, BK) packed-FP4 weight sub-tile (+ its per-16-col FP8 block scales) into a row-major (BN, BK) bf16 SMEM tile, then runs the simdgroup MMA (BK//16 K-steps per call) reading B from SMEM; A loads from DRAM. Double-buffered. Deeper BK feeds more MMA K-steps per cooperative decode, halving the barrier + decode-phase count per unit of MMA work; BK=64 (at BM=128) measured +34-45% over BK=32 at large M, bit-exact, and is the production default for the BM=128 path. The hard ceiling is the M5 constraint that each thread's decode run stay <= 16 bf16 cols (16-bit SIMD arithmetic wider than 16 lanes crashes the Metal backend) -- BK=64/BM=128 gives exactly 16, and BK=128 regresses on SMEM pressure -- so BK=64 is the sweet spot. BN=64 is the Morton grid_n divisor (shared with dense).

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; receives SIMD[c_type, width] at absolute (row, col) (AMD's contract).
  • ​BM (Int): Threadgroup M-tile height (also the Morton grid_m divisor). The simdgroup tile is a square 32x32 (2x2 MmaOpApple), so BM=128 stacks 4 M-simdgroups (8 simdgroups/threadgroup) and BM=64 stacks 2 (4 simdgroups). BM=128 amortizes the cooperative decode over more MMA work and wins at large M; BM=64 keeps occupancy up at small M. The enqueue_* path selects between them by M.
  • ​BK (Int): K-strip depth (cols of the weight decoded per cooperative phase). One cooperative decode feeds BK//16 16x16 MMA K-steps, so a deeper BK halves the barrier + decode-phase count per unit of MMA work. BK=64 measured +34-45% over BK=32 at large M on M5 Max (deeper amortization of the un-overlappable decode), bit-exact. The decode keeps each thread's run to COLS_PER_THREAD <= 16 bf16 cols (one FP8 scale block) -- a hard M5 constraint: the Metal backend crashes on SIMD arithmetic over 16-bit-element vectors wider than 16 lanes (>= 48 bytes), so BK/BM must satisfy COLS_PER_THREAD <= 16. BK=64 at BM=128 gives COLS_PER_THREAD=16 (the sweet spot); BK=128 regresses (SMEM pressure) and BK=64 at BM=64 would force COLS_PER_THREAD=32 (the crash), so the enqueue_* path pairs BK=64 with BM=128 and BK=32 with BM=64.
  • ​coalesce_scales (Bool): On the interior fast path, stage the strip's FP8 block scales into SMEM with one coalesced cooperative load (reused by all threads), instead of each thread issuing a scattered 1-byte scale load. Measured +7-12% at large M (the scale load is uncoalesced- throughput-bound, the dominant reducible decode cost); but it adds a per-strip barrier that regresses small M, so the enqueue_* path enables it only for BM=128. No effect on the bounded edge/tail path (which keeps the per-thread scale load). Bit-exact either way.

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

BN​

comptime BN = 64

BYTES_PER_THREAD​

comptime BYTES_PER_THREAD = (AppleM5Fp4MatMul[c_type, elementwise_lambda_fn, BM=BM, BK=BK, coalesce_scales=coalesce_scales].PACKED_COLS // (Int((mul (BM // Int(32)), _resolve_warp_size(), 2)) // Int(64)))

COLS_PER_THREAD​

comptime COLS_PER_THREAD = (((BK // Int(2)) // (Int((mul (BM // Int(32)), _resolve_warp_size(), 2)) // Int(64))) * Int(2))

in_type​

comptime in_type = DType.bfloat16

Mma​

comptime Mma = MmaOpApple[DType.float32, DType.bfloat16, Int(2), Int(2), transpose_b=True]

NBLK_PER_STRIP​

comptime NBLK_PER_STRIP = (BK // Int(16))

NUM_MMA_M​

comptime NUM_MMA_M = 2

NUM_MMA_N​

comptime NUM_MMA_N = 2

NUM_SG​

comptime NUM_SG = (AppleM5Fp4MatMul[c_type, elementwise_lambda_fn, BM=BM, BK=BK, coalesce_scales=coalesce_scales].NUM_SG_M * Int(2))

NUM_SG_M​

comptime NUM_SG_M = (BM // Int(32))

NUM_SG_N​

comptime NUM_SG_N = 2

PACKED_COLS​

comptime PACKED_COLS = (BK // Int(2))

SF​

comptime SF = NVFP4_SF_VECTOR_SIZE

SG_M​

comptime SG_M = 32

SG_N​

comptime SG_N = 32

THREADS_PER_BLOCK​

comptime THREADS_PER_BLOCK = (AppleM5Fp4MatMul[c_type, elementwise_lambda_fn, BM=BM, BK=BK, coalesce_scales=coalesce_scales].NUM_SG * _resolve_warp_size())

THREADS_PER_ROW​

comptime THREADS_PER_ROW = (Int((mul (BM // Int(32)), _resolve_warp_size(), 2)) // Int(64))

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], log2_grid_m: UInt32, log2_grid_n: UInt32)

W4A16 GEMM kernel entry; M/N/K derive from C, A, and packed.

C is (M, N) row-major, A is (M, K) row-major (bf16 activation), packed is the FP4 weight (N, K//2), scales (N, ceil(K/16)). Grid is (1<<log2_grid_m) * (1<<log2_grid_n) threadgroups; OOB threadgroups early-return after Morton decode.

Edge/tail bounds (bit-exact vs materialize -> dense bf16 matmul): the cooperative stage zero-fills SMEM entries whose n_abs >= N (OOB N-row) or whose K column runs past k_valid (the partial last strip), so the MMA reads exact dequant for valid (n, k) and zero elsewhere -- a zero B contributes nothing, matching the dense reference's sum over k < K. Ragged M is handled on the A side via mma[bounded=True] (a_valid_rows zero-fills A rows past M); B from SMEM is treated fully valid because the stage already zeroed the OOB/tail entries.

Parameters:

  • ​c_layout (TensorLayout): Compile-time TensorLayout of the output tile c.
  • ​a_layout (TensorLayout): Compile-time TensorLayout of the activation tile a.
  • ​packed_layout (TensorLayout): Compile-time TensorLayout of the packed FP4 weight packed.
  • ​scale_layout (TensorLayout): Compile-time TensorLayout of the FP8 block scales scales.

Args:

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact