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

BlockScaledMatmulAMD

struct BlockScaledMatmulAMD[BM: Int = Int(128), BN: Int = Int(128), BK_ELEMS: Int = Int(128), WM: Int = Int(64), WN: Int = Int(64), MMA_M: Int = Int(16), MMA_N: Int = Int(16), MMA_K: Int = Int(128), num_stages: Int = Int(1), matrix_format: CDNA4F8F6F4MatrixFormat = CDNA4F8F6F4MatrixFormat.FLOAT4_E2M1, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None]

Native MXFP4 block-scaled matmul for AMD CDNA4.

Uses cdna4_block_scaled_mfma with FLOAT4_E2M1 format directly. Single-buffer pipeline, or at num_stages == 2 an LDS ping-pong with a depth-2 B-fragment ring. SMEM is plain row-major (no blocked-product), with a conditional XOR swizzle (BlockScaledMmaOp.use_smem_swizzle) removing LDS bank conflicts on the A/B fragment read/write when the tile config qualifies.

Parameters

• ​BM (Int): Block tile rows (output M per block). Default 128.
• ​BN (Int): Block tile cols (output N per block). Default 128.
• ​BK_ELEMS (Int): Block tile K in logical FP4 elements. Default 128.
• ​WM (Int): Warp tile rows. BM must be divisible by WM. Default 64.
• ​WN (Int): Warp tile cols. BN must be divisible by WN. Default 64.
• ​MMA_M (Int): MFMA tile rows. WM must be divisible by MMA_M. Default 16.
• ​MMA_N (Int): MFMA tile cols. WN must be divisible by MMA_N. Default 16.
• ​MMA_K (Int): MFMA K-depth in logical FP4 elements. Default 128.
• ​num_stages (Int): SMEM pipeline depth. 1 is the single-buffer schedule; 2 ping-pongs LDS (even tiles_per_split, VGPR-bound occupancy).
• ​matrix_format (CDNA4F8F6F4MatrixFormat): f8f6f4 operand encoding for A and B (FP4 E2M1 by default). BK_ELEMS counts ELEMENTS, so a given BK_ELEMS costs 1.5x the registers and LDS at MXFP6 and 2x at MXFP8.
• ​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. When set, each output fragment is handed to the lambda with its global (m, n) instead of being stored to c, which lets a caller route output bands elsewhere (e.g. scattering K/V into a paged KV cache). Requires MMA_M == 16 and num_splits == 1.

Implemented traits

AnyType, Deinitable, Movable

comptime members

a_bits

comptime a_bits = matrix_format.bits_per_element()

A_BK_BYTES

comptime A_BK_BYTES = (Int((mul matrix_format.bits_per_element(), BK_ELEMS)) // Int(8))

a_packed_k_per_mma

comptime a_packed_k_per_mma = (Int((mul matrix_format.bits_per_element(), MMA_K)) // Int(8))

A_SMEM_ROW_BYTES

comptime A_SMEM_ROW_BYTES = _smem_row_bytes[(Int((mul matrix_format.bits_per_element(), BK_ELEMS)) // Int(8))]()

b_bits

comptime b_bits = matrix_format.bits_per_element()

B_BK_BYTES

comptime B_BK_BYTES = (Int((mul matrix_format.bits_per_element(), BK_ELEMS)) // Int(8))

b_packed_k_per_mma

comptime b_packed_k_per_mma = (Int((mul matrix_format.bits_per_element(), MMA_K)) // Int(8))

B_SMEM_ROW_BYTES

comptime B_SMEM_ROW_BYTES = _smem_row_bytes[(Int((mul matrix_format.bits_per_element(), BK_ELEMS)) // Int(8))]()

bits_per_element

comptime bits_per_element = BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, num_stages, matrix_format, elementwise_lambda_fn].a_bits

BK_BYTES

comptime BK_BYTES = BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, num_stages, matrix_format, elementwise_lambda_fn].A_BK_BYTES

c_frag_size

comptime c_frag_size = ((MMA_M * MMA_N) // _resolve_warp_size())

k_tile_size

comptime k_tile_size = BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, num_stages, matrix_format, elementwise_lambda_fn].BK_BYTES

lane_bytes

comptime lane_bytes = (Int((mul matrix_format.bits_per_element(), 32)) // Int(8))

num_k_tiles

comptime num_k_tiles = ((Int((mul matrix_format.bits_per_element(), BK_ELEMS)) // Int(8)) // (Int((mul matrix_format.bits_per_element(), MMA_K)) // Int(8)))

num_m_mmas

comptime num_m_mmas = (WM // MMA_M)

num_n_mmas

comptime num_n_mmas = (WN // MMA_N)

num_threads

comptime num_threads = (BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, num_stages, matrix_format, elementwise_lambda_fn].num_warps * _resolve_warp_size())

num_warps

comptime num_warps = (BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, num_stages, matrix_format, elementwise_lambda_fn].num_warps_m * BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, num_stages, matrix_format, elementwise_lambda_fn].num_warps_n)

num_warps_m

comptime num_warps_m = (BM // WM)

num_warps_n

comptime num_warps_n = (BN // WN)

packed_k_per_mma

comptime packed_k_per_mma = BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, num_stages, matrix_format, elementwise_lambda_fn].a_packed_k_per_mma

scales_per_mma

comptime scales_per_mma = (MMA_K // Int(32))

simd_width

comptime simd_width = simd_width_of[DType.uint8]()

Methods

run

static def run[out_dtype: DType, c_layout: TensorLayout, a_layout: TensorLayout, b_layout: TensorLayout, sfa_layout: TensorLayout, sfb_layout: TensorLayout, c_engine: TensorEngine, a_engine: TensorEngine, b_engine: TensorEngine, sfa_engine: TensorEngine, sfb_engine: TensorEngine, num_splits: Int = Int(1)](c: TileTensor[out_dtype, c_layout, MutAnyOrigin, Engine=c_engine], a: TileTensor[.uint8, a_layout, ImmutAnyOrigin, Engine=a_engine], b: TileTensor[.uint8, b_layout, ImmutAnyOrigin, Engine=b_engine], sfa: TileTensor[.float8_e8m0fnu, sfa_layout, ImmutAnyOrigin, Engine=sfa_engine], sfb: TileTensor[.float8_e8m0fnu, sfb_layout, ImmutAnyOrigin, Engine=sfb_engine])

MXFP4 block-scaled GEMM kernel with SMEM pipeline.

With num_splits > 1 this is the inter-block split-K body: each block_idx.z slice accumulates one disjoint K-band into its own [M, N] region of a stacked (num_splits * M, N) float32 workspace (out_dtype is float32 in that mode). A separate reduce kernel sums the num_splits partials and casts to the real output dtype. num_splits == 1 is byte-identical to the no-split path (split_id == 0, full K range, zero offset).

Parameters:

• ​out_dtype (DType): Element type of the output tensor c; must be float32 when num_splits > 1.
• ​c_layout (TensorLayout): Compile-time layout of the output tensor c.
• ​a_layout (TensorLayout): Compile-time layout of the A operand.
• ​b_layout (TensorLayout): Compile-time layout of the B operand.
• ​sfa_layout (TensorLayout): Compile-time layout of the A scales tensor sfa.
• ​sfb_layout (TensorLayout): Compile-time layout of the B scales tensor sfb.
• ​c_engine (TensorEngine): Engine of the output tensor c.
• ​a_engine (TensorEngine): Engine of the A operand.
• ​b_engine (TensorEngine): Engine of the B operand.
• ​sfa_engine (TensorEngine): Engine of the A scales tensor sfa.
• ​sfb_engine (TensorEngine): Engine of the B scales tensor sfb.
• ​num_splits (Int): Number of disjoint K-bands the K dimension is partitioned into (defaults to 1, no split).

Args:

• ​c (TileTensor[out_dtype, c_layout, MutAnyOrigin, Engine=c_engine]): Output matrix [M, N] of dtype out_dtype; in split-K mode a stacked (num_splits * M, N) float32 workspace.
• ​a (TileTensor[.uint8, a_layout, ImmutAnyOrigin, Engine=a_engine]): Packed A operand [M, K//2] uint8, two MXFP4 nibbles per byte.
• ​b (TileTensor[.uint8, b_layout, ImmutAnyOrigin, Engine=b_engine]): Packed B operand [N, K//2] uint8, transposed with two MXFP4 nibbles per byte.
• ​sfa (TileTensor[.float8_e8m0fnu, sfa_layout, ImmutAnyOrigin, Engine=sfa_engine]): A block scales [M, K//32] as float8_e8m0fnu, one scale per 32 MXFP4 elements.
• ​sfb (TileTensor[.float8_e8m0fnu, sfb_layout, ImmutAnyOrigin, Engine=sfb_engine]): B block scales [N, K//32] as float8_e8m0fnu, one scale per 32 MXFP4 elements.