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

def load_AB[a_type: DType, b_type: 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], 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)], 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], 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], 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)

TMA producer phase that loads A and B tiles from GMEM into SMEM via multicast.

Parameters:

a_type (DType): DType of the A operand elements.
b_type (DType): DType of the B operand elements.
a_rank (Int): Number of dimensions in the A TMA tensor map.
a_tile_shape (IndexList[a_rank]): Shape of each A tile loaded by TMA, in elements.
a_desc_shape (IndexList[a_rank]): Shape of the A TMA descriptor, in elements.
b_rank (Int): Number of dimensions in the B TMA tensor map.
b_tile_shape (IndexList[b_rank]): Shape of each B tile loaded by TMA, in elements.
b_desc_shape (IndexList[b_rank]): Shape of the B TMA descriptor, in elements.
a_dim0 (Int): Row extent of each A SMEM tile, in elements.
a_dim1 (Int): Column extent of each A SMEM tile, in elements.
a_num_tiles (Int): Number of A SMEM tiles in the pipeline buffer.
a_swizzle_bytes (Int): SMEM swizzle granularity for A tiles, in bytes.
b_dim0 (Int): Row extent of each B SMEM tile, in elements.
b_dim1 (Int): Column extent of each B SMEM tile, in elements.
b_num_tiles (Int): Number of B SMEM tiles in the pipeline buffer.
b_swizzle_bytes (Int): SMEM swizzle granularity for B tiles, in bytes.
num_pipeline_stages (Int): Number of stages in the load-to-MMA pipeline.
block_tile_shape (IndexList[Int(3)]): Block tile shape as (BM, BN, BK).
mma_shape (IndexList[Int(3)]): MMA instruction shape as (MMA_M, MMA_N, MMA_K).
cta_group (Int): Number of cooperating CTAs per multicast load (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 map descriptor for the A operand.
b_tma_op (TMATensorTile[b_type, b_rank, b_tile_shape, b_desc_shape]): TMA tensor map descriptor for the B operand.
a_smem_tiles (SMemTileArray2D[a_type, a_dim0, a_dim1, a_num_tiles, a_swizzle_bytes]): Pipeline buffer of A SMEM tiles.
b_smem_tiles (SMemTileArray2D[b_type, b_dim0, b_dim1, b_num_tiles, b_swizzle_bytes]): Pipeline buffer of B SMEM tiles.
load_mma_pipeline (ProducerConsumerPipeline[num_pipeline_stages]): Producer/consumer pipeline between TMA load and MMA stages.
peer_cta_coord (Tuple[Int, Int, Int]): Coordinate of the peer CTA within the cluster, used to offset GMEM read and SMEM write addresses.
work_tile_coord (Tuple[Int, Int, Int]): (M, N, batch) tile coordinate assigned to this CTA by the tile scheduler.
a_multicast_mask (UInt16): Multicast mask selecting CTAs that receive the A tile.
b_multicast_mask (UInt16): Multicast mask selecting CTAs that receive the B tile.
iter_idx (UInt32): Current K iteration index for this pipeline stage.
elect_one_cta (Bool): Whether this CTA sets up the TMA barrier byte expectation.