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 function

enqueue_matmul2d_fp4

def enqueue_matmul2d_fp4[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)](c: TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[DType.bfloat16, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], packed: TileTensor[DType.uint8, Storage=packed.Storage, address_space=packed.address_space, linear_idx_type=packed.linear_idx_type], scales: TileTensor[DType.float8_e4m3fn, Storage=scales.Storage, address_space=scales.address_space, linear_idx_type=scales.linear_idx_type], ctx: DeviceContext)

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

a is the bf16 activation (M, K), packed the FP4 weight (N, K//2) (lo-nibble first), scales the FP8-E4M3 block scales (N, ceil(K/16)). C is (M, N). Interior-only fast path: caller must ensure N is a multiple of the threadgroup N-tile (32*num_sg_n*tn) and K % 16 == 0 (the W4A16 production shapes are aligned). M is rounded up and the store guards partial tiles.

The 16x32x16 MMA is the native matmul2d_mma_regc_bt_native tiling (two _mma_apple_transposable calls); pure Mojo over _mma_apple.

Parameters:

c_type (DType): Output element type; one of fp16, bf16, fp32 (defaults to float32). 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 applying a per-element transform to the output (defaults to None).
num_sg_m (Int): Simdgroup rows per threadgroup (defaults to 2).
num_sg_n (Int): Simdgroup cols per threadgroup (defaults to 2).
tm (Int): Output 16x32 tiles per simdgroup along M (defaults to 2).
tn (Int): Output 16x32 tiles per simdgroup along N (defaults to 2). tm*tn == 4 is the M5 register-cliff optimum.

Args:

c (TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output TileTensor view with shape (M, N).
a (TileTensor[DType.bfloat16, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Bf16 activation TileTensor view with shape (M, K).
packed (TileTensor[DType.uint8, Storage=packed.Storage, address_space=packed.address_space, linear_idx_type=packed.linear_idx_type]): Packed FP4 weight TileTensor view with shape (N, K//2) (uint8, lo-nibble first).
scales (TileTensor[DType.float8_e4m3fn, Storage=scales.Storage, address_space=scales.address_space, linear_idx_type=scales.linear_idx_type]): FP8-E4M3 block-scale TileTensor view with shape (N, ceil(K/16)).
ctx (DeviceContext): Device context used to enqueue the kernel.

Raises:

If the attached GPU is not Apple M5 (compute_capability == 5).