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

QuadrantMmaOp

struct QuadrantMmaOp[out_type: DType, in_type: DType, shape: IndexList[Int(3)], k_group_size: Int, num_k_groups: Int, num_m_mmas: Int, num_n_mmas: Int, swizzle: Optional[Swizzle] = None]

MMA operator for AMD matmul ping-pong schedule.

Owns A/B/C register tiles in LOCAL address space. Provides quadrant load/compute methods for the ping-pong double-buffering schedule: load_a_quadrant/load_b_quadrant fill half the register tile from SMEM via load_lds_fragment, then mma_quadrant computes on it.

Parameters

• ​out_type (DType): Accumulator data type (typically float32).
• ​in_type (DType): Input element data type (bfloat16 or float8).
• ​shape (IndexList[Int(3)]): MMA instruction shape [M, N, K].
• ​k_group_size (Int): Number of MMA K-groups per fragment load.
• ​num_k_groups (Int): Number of k-groups across the full BK dimension.
• ​num_m_mmas (Int): MMA tiles along M within the warp tile.
• ​num_n_mmas (Int): MMA tiles along N within the warp tile.
• ​swizzle (Optional[Swizzle]): Optional SMEM swizzle for load helpers.

Implemented traits

AnyType, Deinitable, Movable

comptime members

accum_width

comptime accum_width = QuadrantMmaOp[out_type, in_type, shape, k_group_size, num_k_groups, num_m_mmas, num_n_mmas, swizzle].c_frag_size

BK

comptime BK = (QuadrantMmaOp[out_type, in_type, shape, k_group_size, num_k_groups, num_m_mmas, num_n_mmas, swizzle].num_k_mmas * shape[Int(2)])

c_frag_size

comptime c_frag_size = num_matrix_reg[shape[Int(0)], shape[Int(1)]]()

MMA_K

comptime MMA_K = shape[Int(2)]

MMA_M

comptime MMA_M = shape[Int(0)]

MMA_N

comptime MMA_N = shape[Int(1)]

num_k_mmas

comptime num_k_mmas = (num_k_groups * k_group_size)

quad_m

comptime quad_m = (num_m_mmas // Int(2))

quad_n

comptime quad_n = (num_n_mmas // Int(2))

quad_WM

comptime quad_WM = (QuadrantMmaOp[out_type, in_type, shape, k_group_size, num_k_groups, num_m_mmas, num_n_mmas, swizzle].quad_m * shape[Int(0)])

quad_WN

comptime quad_WN = (QuadrantMmaOp[out_type, in_type, shape, k_group_size, num_k_groups, num_m_mmas, num_n_mmas, swizzle].quad_n * shape[Int(1)])

WM

comptime WM = (num_m_mmas * shape[Int(0)])

WN

comptime WN = (num_n_mmas * shape[Int(1)])

Methods

__init__

def __init__(out self)

load_a_quadrant

def load_a_quadrant[which: Int](self, smem_tile: TileTensor[in_type, address_space=AddressSpace.SHARED])

Load A quadrant which from SMEM sub-tile to registers.

Tiles a_reg as [quad_m, reg_cols](which, 0) to get the register sub-tile for this quadrant, then loads via load_lds_fragment.

Parameters:

• ​which (Int): Quadrant index selecting which half of the A register tile to fill.

Args:

• ​smem_tile (TileTensor[in_type, address_space=AddressSpace.SHARED]): SMEM sub-tile for this A quadrant of shape quad_WM x BK.

load_b_quadrant

def load_b_quadrant[which: Int](self, smem_tile: TileTensor[in_type, address_space=AddressSpace.SHARED])

Load B quadrant which from SMEM sub-tile to registers.

Parameters:

• ​which (Int): Quadrant index selecting which half of the B register tile to fill.

Args:

• ​smem_tile (TileTensor[in_type, address_space=AddressSpace.SHARED]): SMEM sub-tile for this B quadrant of shape quad_WN x BK.

mma_quadrant

def mma_quadrant[which_a: Int, which_b: Int](self)

Execute MMA for quadrant (which_a, which_b) via TiledMma.

Slices A/B/C register tiles to the quadrant and delegates to TiledMma for stateless computation.

Parameters:

• ​which_a (Int): Quadrant index along M selecting the A and C row half.
• ​which_b (Int): Quadrant index along N selecting the B and C column half.

accum_tile

def accum_tile(self) -> TileTensor[out_type, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

Return the accumulator register tile.

Returns:

TileTensor[out_type, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]