IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /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. /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), 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 with schedule-driven prologue/kernel/epilogue. 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.
  • ​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))

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

bits_per_element​

comptime bits_per_element = BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, 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, 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, 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, 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, matrix_format, elementwise_lambda_fn].num_warps_m * BlockScaledMatmulAMD[BM, BN, BK_ELEMS, WM, WN, MMA_M, MMA_N, MMA_K, 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, 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, num_splits: Int = Int(1)](c: TileTensor[out_dtype, c_layout, MutAnyOrigin], a: TileTensor[DType.uint8, a_layout, ImmutAnyOrigin], b: TileTensor[DType.uint8, b_layout, ImmutAnyOrigin], sfa: TileTensor[DType.float8_e8m0fnu, sfa_layout, ImmutAnyOrigin], sfb: TileTensor[DType.float8_e8m0fnu, sfb_layout, ImmutAnyOrigin])

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.
  • ​num_splits (Int): Number of disjoint K-bands the K dimension is partitioned into (defaults to 1, no split).

Args:

Models

Cookbook

Puzzles

Blog

Community

Contact