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_smem

def enqueue_matmul2d_fp4_smem[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(16), num_sg_n: Int = Int(1), tm: Int = Int(4), tn: Int = Int(1), smem_bk: Int = Int(256)](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 cooperative-decode matmul2d W4A16 GEMM (run_smem_decode).

Coalesced FP4 decode -> SMEM -> register B fragment (the fast W4A16 path; the per-lane run is decode-load-bound). Interior-only: caller must ensure N is a multiple of the threadgroup N-tile (32*num_sg_n*tn) and K % smem_bk == 0 (smem_bk is the cooperative-decode strip depth).

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 16).
num_sg_n (Int): Simdgroup cols per threadgroup (defaults to 1).
tm (Int): Output 16x32 tiles per simdgroup along M (defaults to 4).
tn (Int): Output 16x32 tiles per simdgroup along N (defaults to 1).
smem_bk (Int): Cooperative-decode strip depth in K; one decode feeds smem_bk//16 MMA K-steps (defaults to 256).

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).