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

BlockScaledMmaOp

struct BlockScaledMmaOp[mma_shape: IndexList[Int(3)], num_m_mmas: Int, num_n_mmas: Int, num_k_tiles: Int, num_b_slots: Int = Int(1), matrix_format: CDNA4F8F6F4MatrixFormat = CDNA4F8F6F4MatrixFormat.FLOAT4_E2M1]

Register ownership + block-scaled MFMA execution.

Loads packed uint8 A/B fragments from SMEM or GMEM and executes cdna4_block_scaled_mfma with per-lane E8M0 scale values.

Scale operand model: Each lane holds 32 FP4 elements and one E8M0 scale byte, matching the MX format's per-32-element granularity exactly. For 16x16x128: 64 lanes cover 16 rows x 4 K-groups. lane_row = lane_id % 16 (matrix row) lane_k_group = lane_id / 16 (K-group 0..3)

Scale packing: 4 spatial MMA tiles' scale bytes are packed into one Int32 VGPR: byte i holds the scale for m_mma=i (A) or n_mma=i (B). The MFMA byte-index selector (OP_SEL) picks the correct byte for each MMA tile, so one scale load covers all 4 m_mma or n_mma positions with zero overhead.

Parameters​

  • ​mma_shape (IndexList[Int(3)]): MFMA tile shape as (M, N, K) in logical FP4 elements, (16, 16, 128).
  • ​num_m_mmas (Int): Number of spatial M MMA tiles per warp tile (WM // MMA_M). Must be <= 4.
  • ​num_n_mmas (Int): Number of spatial N MMA tiles per warp tile (WN // MMA_N). Must be <= 4.
  • ​num_k_tiles (Int): Number of K sub-tiles within one BK iteration (BK_BYTES // packed_k_per_mma).
  • ​num_b_slots (Int): Number of B register slots for depth-2 prefetch (defaults to 1).
  • ​matrix_format (CDNA4F8F6F4MatrixFormat): f8f6f4 operand encoding for A and B. A lane covers 32 K-elements in every format; what changes is the bytes that occupies -- 16 (FP4), 24 (FP6), 32 (FP8) -- which lane_bytes below derives.

Implemented traits​

AnyType, Deinitable, Movable

comptime members​

a_bits​

comptime a_bits = matrix_format.bits_per_element()

A_BK_BYTES​

comptime A_BK_BYTES = (num_k_tiles * (Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)))

a_frag_width_bytes​

comptime a_frag_width_bytes = ((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)]))

A_K_HALF_STRIDE​

comptime A_K_HALF_STRIDE = ((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) // Int(16)))

a_num_frag_halves​

comptime a_num_frag_halves = (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) // Int(16))

a_packed_k_per_mma​

comptime a_packed_k_per_mma = (Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8))

a_reg_frag_bytes​

comptime a_reg_frag_bytes = SIMD(matrix_format.simd_width())

A_SMEM_ROW_BYTES​

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

b_bits​

comptime b_bits = matrix_format.bits_per_element()

B_BK_BYTES​

comptime B_BK_BYTES = (num_k_tiles * (Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)))

b_frag_width_bytes​

comptime b_frag_width_bytes = ((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)]))

B_K_HALF_STRIDE​

comptime B_K_HALF_STRIDE = ((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) // Int(16)))

b_num_frag_halves​

comptime b_num_frag_halves = (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) // Int(16))

b_packed_k_per_mma​

comptime b_packed_k_per_mma = (Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8))

b_reg_frag_bytes​

comptime b_reg_frag_bytes = SIMD(matrix_format.simd_width())

B_SMEM_ROW_BYTES​

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

bits_per_element​

comptime bits_per_element = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].a_bits

BK_BYTES​

comptime BK_BYTES = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].A_BK_BYTES

c_frag_size​

comptime c_frag_size = (Int((mul mma_shape[Int(0)], mma_shape[Int(1)])) // _resolve_warp_size())

FRAG_HALF_BYTES​

comptime FRAG_HALF_BYTES = 16

K_HALF_STRIDE​

comptime K_HALF_STRIDE = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].A_K_HALF_STRIDE

lane_bytes​

comptime lane_bytes = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].mma_frag_width_bytes

lanes_per_row​

comptime lanes_per_row = (_resolve_warp_size() // mma_shape[Int(0)])

mma_frag_width_bytes​

comptime mma_frag_width_bytes = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].a_frag_width_bytes

MMA_K​

comptime MMA_K = mma_shape[Int(2)]

MMA_M​

comptime MMA_M = mma_shape[Int(0)]

MMA_N​

comptime MMA_N = mma_shape[Int(1)]

num_frag_halves​

comptime num_frag_halves = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].a_num_frag_halves

packed_k_per_mma​

comptime packed_k_per_mma = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].a_packed_k_per_mma

reg_frag_bytes​

comptime reg_frag_bytes = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].a_reg_frag_bytes

scales_per_mma​

comptime scales_per_mma = (mma_shape[Int(2)] // Int(32))

SMEM_ROW_BYTES​

comptime SMEM_ROW_BYTES = BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].A_SMEM_ROW_BYTES

smem_swizzle​

comptime smem_swizzle = Optional(_full_row_lds_swizzle[_smem_row_bytes[Int((mul (Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)), num_k_tiles))]()]()) if BlockScaledMmaOp[mma_shape, num_m_mmas, num_n_mmas, num_k_tiles, num_b_slots, matrix_format].use_smem_swizzle else Optional()

use_smem_swizzle​

comptime use_smem_swizzle = True if (num_k_tiles == Int(1)) else (num_k_tiles == Int(1)) and (mma_shape[Int(0)] == Int(16)) and True and (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(16)) if (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(16)) else (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(32)) if (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) // Int(16))) == Int(64)) else (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) // Int(16))) == Int(64)) and (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(16)) if (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(16)) else (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(32)) or (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(24)) if (matrix_format.bits_per_element() == Int(6)) else (matrix_format.bits_per_element() == Int(6)) and (((Int((mul mma_shape[Int(2)], matrix_format.bits_per_element())) // Int(8)) // (_resolve_warp_size() // mma_shape[Int(0)])) == Int(24))

Methods​

__init__​

def __init__(out self)

accum_tile​

def accum_tile(self) -> ref[self._c_reg] TileTensor[.float32, Layout[TypeList[ComptimeInt[num_m_mmas], ComptimeInt[(num_n_mmas * (Int((mul mma_shape[Int(0)], mma_shape[Int(1)])) // _resolve_warp_size()))]](), TypeList[ComptimeInt[Int((mul (Int((mul mma_shape[Int(0)], mma_shape[Int(1)])) // _resolve_warp_size()), num_n_mmas))], ComptimeInt[Int(1)]]()], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

Returns:

ref[self._c_reg] TileTensor[.float32, Layout[TypeList[ComptimeInt[num_m_mmas], ComptimeInt[(num_n_mmas * (Int((mul mma_shape[Int(0)], mma_shape[Int(1)])) // _resolve_warp_size()))]](), TypeList[ComptimeInt[Int((mul (Int((mul mma_shape[Int(0)], mma_shape[Int(1)])) // _resolve_warp_size()), num_n_mmas))], ComptimeInt[Int(1)]]()], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

load_frag_from_smem​

def load_frag_from_smem[k_tile_idx: Int, slot: Int = Int(0)](mut self, a_smem_warp: TileTensor[.uint8, Engine=a_smem_warp.Engine, address_space=AddressSpace.SHARED, linear_idx_type=a_smem_warp.linear_idx_type], b_smem_warp: TileTensor[.uint8, Engine=b_smem_warp.Engine, address_space=AddressSpace.SHARED, linear_idx_type=b_smem_warp.linear_idx_type])

Loads A/B fragments from row-major SMEM for k-tile k_tile_idx.

Dispatches on the operand format at compile time. The dispatch lives here rather than at the call sites so a new call site cannot forget it and silently instantiate the power-of-two loader for FP6, whose 24-byte fragment vectorize/distribute cannot express.

Parameters:

  • ​k_tile_idx (Int): K-tile index within the current BK iteration.
  • ​slot (Int): B-fragment ring slot (default 0). Depth-2 writes the next tile's B here so it does not WAR against the current MFMAs.

Args:

load_a_frag_from_smem​

def load_a_frag_from_smem[k_tile_idx: Int](self, a_smem_warp: TileTensor[.uint8, Engine=a_smem_warp.Engine, address_space=AddressSpace.SHARED, linear_idx_type=a_smem_warp.linear_idx_type])

A-only variant of load_frag_from_smem for callers that source B elsewhere (e.g. preshuffled DRAM via PreshuffledBLoader).

Parameters:

  • ​k_tile_idx (Int): K-tile index within the current BK iteration.

Args:

load_b_frag_preshuffled​

def load_b_frag_preshuffled[k_tile_idx: Int, N: Int, K_BYTES: Int, slot: Int = Int(0)](self, b_loader: PreshuffledBLoader[N, K_BYTES], warp_n_off: Int, k_byte_base: Int)

Load B fragments directly from preshuffled DRAM into b_reg slot slot.

Each lane issues one buffer_load_dwordx4 per (k_tile, n_mma) at the per-lane MFMA mapping (lane%16 → n-row, lane//16 → k-group). The slot parameter selects which b_reg half to write into when num_b_slots > 1 (depth-2 prefetch).

Parameters:

  • ​k_tile_idx (Int): K-tile index within the current BK iteration.
  • ​N (Int): Total N dimension of the B matrix in output columns.
  • ​K_BYTES (Int): Total K dimension in packed bytes (K // 2).
  • ​slot (Int): B register slot to write into (defaults to 0).

Args:

  • ​b_loader (PreshuffledBLoader[N, K_BYTES]): Preshuffled B DRAM loader issuing per-lane buffer_load_dwordx4 reads.
  • ​warp_n_off (Int): Starting N-row offset of this warp's B tile within the block.
  • ​k_byte_base (Int): Base byte offset in K for the current BK iteration.

load_scales_from_smem​

def load_scales_from_smem[k_tile_idx: Int](mut self, a_scale_smem_warp: TileTensor[.uint8, Engine=a_scale_smem_warp.Engine, address_space=AddressSpace.SHARED, linear_idx_type=a_scale_smem_warp.linear_idx_type], b_scale_smem_warp: TileTensor[.uint8, Engine=b_scale_smem_warp.Engine, address_space=AddressSpace.SHARED, linear_idx_type=b_scale_smem_warp.linear_idx_type])

Load packed scale VGPRs for k-tile k_tile_idx from SMEM.

Packs num_m_mmas (A) or num_n_mmas (B) scale bytes into one Int32 each using the same col_major[MMA_M, WARP_SIZE/MMA_M] distribute pattern as load_frag_from_smem. Each lane picks one scale byte via (lane_row, lane_k_group). TileTensor's stride handling means this works for any parent SMEM layout.

The MFMA byte-index selector (a_scale_byte_index=m_mma, b_scale_byte_index=n_mma) picks the correct byte: no shifts or masks at consumption time.

Parameters:

  • ​k_tile_idx (Int): K-tile index within the current BK iteration.

Args:

mma​

def mma[k_tile_idx: Int, slot: Int = Int(0)](self)

Execute block-scaled MFMA for k-tile k_tile_idx using B from slot.

B→src_a, A→src_b (AMD MFMA convention). The packed scale VGPRs hold one byte per spatial MMA tile. a_scale_byte_index=m selects byte m from _a_scale_packed, b_scale_byte_index=n selects byte n from _b_scale_packed.

slot selects which b_reg half to read when num_b_slots > 1.

Parameters:

  • ​k_tile_idx (Int): K-tile index within the current BK iteration.
  • ​slot (Int): B register slot to read from when num_b_slots > 1 (defaults to 0).

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community