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

mxfp4_block_scaled_matmul_amd

def mxfp4_block_scaled_matmul_amd[MMA_M: Int = Int(16), MMA_N: Int = Int(16), MMA_K: Int = Int(128)](c: TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[DType.uint8, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[DType.uint8, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type], a_scales: TileTensor[DType.float8_e8m0fnu, Storage=a_scales.Storage, address_space=a_scales.address_space, linear_idx_type=a_scales.linear_idx_type], b_scales: TileTensor[DType.float8_e8m0fnu, Storage=b_scales.Storage, address_space=b_scales.address_space, linear_idx_type=b_scales.linear_idx_type], ctx: DeviceContext)

Launch native MXFP4 block-scaled matmul on AMD CDNA4.

Uses cdna4_block_scaled_mfma with FLOAT4_E2M1 directly: no dequantization to FP8. Both A and B must be packed uint8 with E8M0 scaling factors. Accumulates in float32, casts to c.dtype during the store epilogue.

Tile shape is selected at runtime based on M to match the expected arithmetic intensity regime (decode vs. prefill). A proper (N, K)-keyed dispatch table is planned for a follow-up PR.

Parameters:

MMA_M (Int): MFMA tile rows. Default 16. Pass 32 to opt into the 32x32x64 MFMA shape (must be paired with MMA_N=32, MMA_K=64).
MMA_N (Int): MFMA tile cols. Default 16.
MMA_K (Int): MFMA K-depth in logical FP4 elements. Default 128.

Args:

c (TileTensor[Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output [M, N] (any float dtype, e.g. float32 or bfloat16).
a (TileTensor[DType.uint8, Storage=a.Storage, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Packed A [M, K//2] uint8 (two MXFP4 elements per byte).
b (TileTensor[DType.uint8, Storage=b.Storage, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Packed B [N, K//2] uint8 (transposed, two MXFP4 per byte).
a_scales (TileTensor[DType.float8_e8m0fnu, Storage=a_scales.Storage, address_space=a_scales.address_space, linear_idx_type=a_scales.linear_idx_type]): A scales [M, K//32] float8_e8m0fnu.
b_scales (TileTensor[DType.float8_e8m0fnu, Storage=b_scales.Storage, address_space=b_scales.address_space, linear_idx_type=b_scales.linear_idx_type]): B scales [N, K//32] float8_e8m0fnu.
ctx (DeviceContext): Device context for kernel launch.