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 struct

BlockScaledMmaOp_PreB

struct BlockScaledMmaOp_PreB[mma_shape: IndexList[Int(3)], warp_tile: IndexList[Int(3)], num_b_slots: Int = Int(1), num_scale_slots: Int = Int(1)]

Per-warp register state + MFMA dispatch for the preb (preshuffled-B, preshuffled-scale) kernel.

warp_tile is the (M, N, K) region this warp computes per outer-K iteration, in the same element units as mma_shape. Per-warp MFMA counts are derived as warp_tile[i] // mma_shape[i].

Asserted in __init__: warp_tile[i] % mma_shape[i] == 0 per axis, and num_k_mmas % 2 == 0 (k_pack=2 cell halves). num_m_mmas / num_n_mmas may be odd; the constructor rotates the scale i32 per CTA so OPSEL keeps the same comptime formula. See module-level comment for the scale-cell byte ordering.

Parameters

• ​mma_shape (IndexList[Int(3)]): The MFMA instruction shape as (M, N, K) in MFMA-native element units, for example (16, 16, 128) for MXFP4.
• ​warp_tile (IndexList[Int(3)]): The (M, N, K) region this warp computes per outer-K iteration, in the same element units as mma_shape. Per-warp MFMA counts are derived as warp_tile[i] // mma_shape[i].
• ​num_b_slots (Int): Number of _b_reg slots for software pipelining (defaults to 1). Set to 2 to double-buffer B fragments across outer-K iterations.
• ​num_scale_slots (Int): Number of A/B scale-register ring slots for software pipelining (defaults to 1). At the depth-2 default this is 1 (1-deep, unchanged behavior); the depth-3 co-deepened prefetch sets it to 2 so the scale ring double-buffers alongside the B fragments.

Implemented traits

AnyType, ImplicitlyDeletable

comptime members

BK_BYTES

comptime BK_BYTES = (warp_tile[Int(2)] // Int(2))

c_frag_size

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

mma_frag_width_bytes

comptime mma_frag_width_bytes = Int(16)

MMA_K

comptime MMA_K = mma_shape[Int(2)]

MMA_K_BYTES

comptime MMA_K_BYTES = (mma_shape[Int(2)] // Int(2))

MMA_M

comptime MMA_M = mma_shape[Int(0)]

MMA_N

comptime MMA_N = mma_shape[Int(1)]

num_k_mmas

comptime num_k_mmas = (warp_tile[Int(2)] // mma_shape[Int(2)])

num_m_mmas

comptime num_m_mmas = (warp_tile[Int(0)] // mma_shape[Int(0)])

num_n_mmas

comptime num_n_mmas = (warp_tile[Int(1)] // mma_shape[Int(1)])

Methods

__init__

def __init__(out self, warp_m_off: Int, warp_n_off: Int)

accum_tile

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

Returns:

ref[self._c_reg] TileTensor[DType.float32, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

load_a_frag_from_smem

def load_a_frag_from_smem[mma_k_idx: Int](self, a_smem_warp: TileTensor[DType.uint8, Storage=a_smem_warp.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a_smem_warp.linear_idx_type])

Load A fragment for MFMA-K position mma_k_idx from row-major SMEM.

XOR-16 swizzled read (matches the write in copy_a_tile_to_smem): each lane reads the 16B vec at slot-tile (row, col_byte), then swizzles the flat in-tile byte offset before the raw_load. WM==BM so a_smem_warp IS the contiguous [BM, BK_BYTES] slot tile and raw_load indexes it directly.

Parameters:

• ​mma_k_idx (Int): Index of the MFMA step along K within the warp tile, in [0, num_k_mmas).

Args:

• ​a_smem_warp (TileTensor[DType.uint8, Storage=a_smem_warp.Storage, address_space=AddressSpace.SHARED, linear_idx_type=a_smem_warp.linear_idx_type]): The shared-memory A tile for this warp, a contiguous [BM, BK_BYTES] slot tile indexed directly by raw_load.

load_b_frag_preshuffled

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

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

Parameters:

• ​mma_k_idx (Int): Index of the MFMA step along K within the warp tile, in [0, num_k_mmas).
• ​slot (Int): The _b_reg slot to load into (defaults to 0).

Args:

• ​b_loader (PreshuffledBLoader): The PreshuffledBLoader for the preshuffled B tensor.
• ​warp_n_off (Int): Global N offset of this warp's tile.
• ​k_byte_base (Int): Base byte offset along K for the current outer-K tile.

load_a_scales_preshuffled

def load_a_scales_preshuffled[k_pair: Int, slot: Int = Int(0)](mut self, a_scale_loader: PreshuffledScaleLoader, warp_m_off: Int, k_pair_idx: Int)

Issue per-lane i32 scale loads for A at one k_pair slot.

Caller provides the absolute k_pair_idx (= k_iter * (num_k_mmas / 2) + k_pair); each step advances by 8 K-scales (= 2 MFMAs along K). One i32 per (mi_pair, k_pair) per lane.

Parameters:

• ​k_pair (Int): Index of the k_pair slot within the current outer-K tile, in [0, num_k_mmas // 2).
• ​slot (Int): Scale-ring slot (comptime; 0 at the 1-deep default).

Args:

• ​a_scale_loader (PreshuffledScaleLoader): The PreshuffledScaleLoader for the A scale tensor.
• ​warp_m_off (Int): Global M offset of this warp's tile.
• ​k_pair_idx (Int): Absolute k_pair index across all outer-K iterations, equal to k_iter * (num_k_mmas / 2) + k_pair; each step advances by 8 K-scales.

load_b_scales_preshuffled

def load_b_scales_preshuffled[k_pair: Int, slot: Int = Int(0)](mut self, b_scale_loader: PreshuffledScaleLoader, warp_n_off: Int, k_pair_idx: Int)

Mirror of load_a_scales_preshuffled along N.

Parameters:

• ​k_pair (Int): Index of the k_pair slot within the current outer-K tile, in [0, num_k_mmas // 2).
• ​slot (Int): Scale-ring slot (comptime; 0 at the 1-deep default).

Args:

• ​b_scale_loader (PreshuffledScaleLoader): The PreshuffledScaleLoader for the B scale tensor.
• ​warp_n_off (Int): Global N offset of this warp's tile.
• ​k_pair_idx (Int): Absolute k_pair index across all outer-K iterations, equal to k_iter * (num_k_mmas / 2) + k_pair; each step advances by 8 K-scales.

mma

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

Block-scaled MFMA at MFMA-K position mma_k_idx using B from slot.

B-major / n-outer / m-inner: hoist the B fragment + b_byte + b_scale (VMEM-loaded _b_reg) once per n, then cycle A (m-inner, LDS-loaded _a_reg). Keeping B resident across the m-loop improves MFMA ILP.

OPSEL byte selection from the 2x2 cell: a_byte = (mma_k_idx % 2) * 2 + (m % 2) b_byte = (mma_k_idx % 2) * 2 + (n % 2) Scale dword lives at _*_scale_packed[mn // 2, mma_k_idx // 2]. WM/WN=16 CTAs see only m=0 / n=0, so the constructor shrui (_a_scale_shift / _b_scale_shift) rotates the i32 to the right OPSEL byte.

Parameters:

• ​mma_k_idx (Int): Index of the MFMA step along K within the warp tile, in [0, num_k_mmas).
• ​slot (Int): The _b_reg slot to read B fragments from (defaults to 0).
• ​scale_slot (Int): The scale-ring slot to read A/B scales from (defaults to 0; non-zero only under the depth-3 co-deepened prefetch).

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit