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 function

load_AB_SFA

def load_AB_SFA[a_type: DType, b_type: DType, sfa_dtype: DType, sfa_tma_dtype: DType, a_rank: Int, a_tile_shape: IndexList[a_rank], a_desc_shape: IndexList[a_rank], b_rank: Int, b_tile_shape: IndexList[b_rank], b_desc_shape: IndexList[b_rank], sfa_rank: Int, sfa_tile_shape: IndexList[sfa_rank], sfa_desc_shape: IndexList[sfa_rank], a_dim0: Int, a_dim1: Int, a_num_tiles: Int, a_swizzle_bytes: Int, b_dim0: Int, b_dim1: Int, b_num_tiles: Int, b_swizzle_bytes: Int, num_pipeline_stages: Int, /, *, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], num_sf_k_tiles: Int, cta_group: Int = Int(1), k_group_size: Int = Int(1)](a_tma_op: TMATensorTile[a_type, a_rank, a_tile_shape, a_desc_shape], b_tma_op: TMATensorTile[b_type, b_rank, b_tile_shape, b_desc_shape], sfa_tma_op: TMATensorTile[sfa_tma_dtype, sfa_rank, sfa_tile_shape, sfa_desc_shape], a_smem_tiles: SMemTileArray2D[a_type, a_dim0, a_dim1, a_num_tiles, a_swizzle_bytes], b_smem_tiles: SMemTileArray2D[b_type, b_dim0, b_dim1, b_num_tiles, b_swizzle_bytes], sfa_smem_tiles: SMemTileArrayWithLayout[sfa_dtype], load_mma_pipeline: ProducerConsumerPipeline[num_pipeline_stages], peer_cta_coord: Tuple[Int, Int, Int], work_tile_coord: Tuple[Int, Int, Int], a_multicast_mask: UInt16, b_multicast_mask: UInt16, iter_idx: UInt32, elect_one_cta: Bool)

Issues multicast TMA loads for one K-group of A, B, and SFA into SMEM.

Waits for the MMA consumer to release the pipeline slot, programs the producer mbarrier with the expected byte count, and fires 3D multicast loads for A and B plus a 4D uint16 load for SFA into the stage's SMEM tiles. SFB is not loaded here; it uses the cp.async path.

Parameters:

a_type (DType): Element dtype of the A operand matrix.
b_type (DType): Element dtype of the B operand matrix.
sfa_dtype (DType): Element dtype of the A scale factors.
sfa_tma_dtype (DType): Element dtype used for the SFA TMA descriptor; may differ from sfa_dtype (for example uint16 for 4D TMA).
a_rank (Int): Tensor rank of the A operand TMA descriptor.
a_tile_shape (IndexList[a_rank]): Per-tile shape of the A TMA load.
a_desc_shape (IndexList[a_rank]): Full descriptor shape of the A TMA load.
b_rank (Int): Tensor rank of the B operand TMA descriptor.
b_tile_shape (IndexList[b_rank]): Per-tile shape of the B TMA load.
b_desc_shape (IndexList[b_rank]): Full descriptor shape of the B TMA load.
sfa_rank (Int): Tensor rank of the SFA TMA descriptor.
sfa_tile_shape (IndexList[sfa_rank]): Per-tile shape of the SFA TMA load.
sfa_desc_shape (IndexList[sfa_rank]): Full descriptor shape of the SFA TMA load.
a_dim0 (Int): Row count of each A SMEM tile.
a_dim1 (Int): Column count of each A SMEM tile.
a_num_tiles (Int): Total number of A SMEM tiles across all pipeline stages.
a_swizzle_bytes (Int): Swizzle stride in bytes for the A SMEM tiles.
b_dim0 (Int): Row count of each B SMEM tile.
b_dim1 (Int): Column count of each B SMEM tile.
b_num_tiles (Int): Total number of B SMEM tiles across all pipeline stages.
b_swizzle_bytes (Int): Swizzle stride in bytes for the B SMEM tiles.
num_pipeline_stages (Int): Number of producer/consumer stages in the A/B/SFA load and MMA pipeline.
block_tile_shape (IndexList[Int(3)]): Block tile shape as (BM, BN, BK) in elements.
mma_shape (IndexList[Int(3)]): MMA atom shape as (MMA_M, MMA_N, MMA_K) in elements.
num_sf_k_tiles (Int): Number of scale-factor K-tiles loaded per K-group iteration.
cta_group (Int): Number of CTAs cooperating per MMA group (defaults to 1).
k_group_size (Int): Number of K-tiles loaded per pipeline stage (defaults to 1).

Args:

a_tma_op (TMATensorTile[a_type, a_rank, a_tile_shape, a_desc_shape]): TMA tensor tile descriptor for loading A from global memory.
b_tma_op (TMATensorTile[b_type, b_rank, b_tile_shape, b_desc_shape]): TMA tensor tile descriptor for loading B from global memory.
sfa_tma_op (TMATensorTile[sfa_tma_dtype, sfa_rank, sfa_tile_shape, sfa_desc_shape]): TMA tensor tile descriptor for loading SFA (A scale factors) from global memory.
a_smem_tiles (SMemTileArray2D[a_type, a_dim0, a_dim1, a_num_tiles, a_swizzle_bytes]): SMEM tile array holding the A operand tiles.
b_smem_tiles (SMemTileArray2D[b_type, b_dim0, b_dim1, b_num_tiles, b_swizzle_bytes]): SMEM tile array holding the B operand tiles.
sfa_smem_tiles (SMemTileArrayWithLayout[sfa_dtype]): SMEM tile array holding the A scale-factor tiles.
load_mma_pipeline (ProducerConsumerPipeline[num_pipeline_stages]): Producer/consumer pipeline synchronizing A/B/SFA loads with MMA consumption.
peer_cta_coord (Tuple[Int, Int, Int]): (v, m, n) coordinates of this CTA within the cluster, used to compute SMEM slice offsets for multicast distribution.
work_tile_coord (Tuple[Int, Int, Int]): (M, N, batch) coordinates of the output tile being computed.
a_multicast_mask (UInt16): Multicast bitmask selecting which CTAs receive the A TMA load.
b_multicast_mask (UInt16): Multicast bitmask selecting which CTAs receive the B TMA load.
iter_idx (UInt32): Current K-iteration index within the tile loop, in units of individual K-tiles.
elect_one_cta (Bool): Whether this CTA is elected as the leader for mbarrier byte-count programming.