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_scales_type: DType, a_tile_rank: Int, a_tile_shape: IndexList[a_tile_rank], a_desc_shape: IndexList[a_tile_rank], b_tile_rank: Int, b_tile_shape: IndexList[b_tile_rank], b_desc_shape: IndexList[b_tile_rank], a_scales_tile_rank: Int, a_scales_tile_shape: IndexList[a_scales_tile_rank], a_scales_desc_shape: IndexList[a_scales_tile_rank], num_pipeline_stages: Int, expert_ids_layout: Layout, /, *, a_smem_layout: Layout, b_smem_layout: Layout, a_scales_smem_layout: Layout, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], cta_group: Int = Int(1)](a_tma_op: TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape], b_tma_op: TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape], a_scales_tma_op: TMATensorTile[a_scales_type, a_scales_tile_rank, a_scales_tile_shape, a_scales_desc_shape], a_smem_base: Pointer[Scalar[a_type], address_space=AddressSpace.SHARED, _safe=False], b_smem_base: Pointer[Scalar[b_type], address_space=AddressSpace.SHARED, _safe=False], a_scales_smem_base: Pointer[Scalar[a_scales_type], address_space=AddressSpace.SHARED, _safe=False], load_mma_pipeline: ProducerConsumerPipeline[num_pipeline_stages], peer_cta_coord: Tuple[Int, Int, Int], work_tile_coord: Tuple[Int, Int], a_multicast_mask: UInt16, b_multicast_mask: UInt16, iter_idx: Int, elect_one_cta: Bool, scheduler: TileScheduler[static_MN=scheduler.static_MN, tile_shape=scheduler.tile_shape, cluster=scheduler.cluster, cta_group=scheduler.cta_group, swizzle=scheduler.swizzle, swapAB=scheduler.swapAB], expert_ids: LayoutTensor[DType.int32, expert_ids_layout, ImmutAnyOrigin])

Issues multicast TMA loads for A, B, and A scales into the producer stage of the load-MMA pipeline.

Waits for the consumer (MMA) to release the buffer, then launches async multicast TMA copies of A, B, and the matching A scales strip into the current pipeline stage's SMEM and signals completion through the stage's mbarrier.

Parameters:

a_type (DType): Element dtype of the A activation operand (inferred).
b_type (DType): Element dtype of the B weight operand (inferred).
a_scales_type (DType): Element dtype of the A per-block scales (inferred).
a_tile_rank (Int): Rank of the A TMA tile descriptor (inferred).
a_tile_shape (IndexList[a_tile_rank]): Shape of each A TMA tile copy (inferred).
a_desc_shape (IndexList[a_tile_rank]): Descriptor shape of the A TMA tensor (inferred).
b_tile_rank (Int): Rank of the B TMA tile descriptor (inferred).
b_tile_shape (IndexList[b_tile_rank]): Shape of each B TMA tile copy (inferred).
b_desc_shape (IndexList[b_tile_rank]): Descriptor shape of the B TMA tensor (inferred).
a_scales_tile_rank (Int): Rank of the A scales TMA tile descriptor (inferred).
a_scales_tile_shape (IndexList[a_scales_tile_rank]): Shape of each A scales TMA tile copy (inferred).
a_scales_desc_shape (IndexList[a_scales_tile_rank]): Descriptor shape of the A scales TMA tensor (inferred).
num_pipeline_stages (Int): Number of load-MMA pipeline buffer stages (inferred).
expert_ids_layout (Layout): Layout of the expert-IDs mapping tensor (inferred).
a_smem_layout (Layout): SMEM layout for the A tile buffer.
b_smem_layout (Layout): SMEM layout for the B tile buffer.
a_scales_smem_layout (Layout): SMEM layout for the A scales strip.
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): CTA group size, 1 or 2 (defaults to 1).

Args:

a_tma_op (TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape]): TMA descriptor for loading A activation tiles.
b_tma_op (TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape]): TMA descriptor for loading B weight tiles.
a_scales_tma_op (TMATensorTile[a_scales_type, a_scales_tile_rank, a_scales_tile_shape, a_scales_desc_shape]): TMA descriptor for loading A scales tiles.
a_smem_base (Pointer[Scalar[a_type], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the A SMEM buffer.
b_smem_base (Pointer[Scalar[b_type], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the B SMEM buffer.
a_scales_smem_base (Pointer[Scalar[a_scales_type], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the A scales SMEM buffer.
load_mma_pipeline (ProducerConsumerPipeline[num_pipeline_stages]): Producer-consumer pipeline between load and MMA warps.
peer_cta_coord (Tuple[Int, Int, Int]): Peer CTA coordinates as (peer_id, mma_coord_m, mma_coord_n).
work_tile_coord (Tuple[Int, Int]): Work tile coordinates as (m, n).
a_multicast_mask (UInt16): Multicast mask selecting CTAs for A TMA loads.
b_multicast_mask (UInt16): Multicast mask selecting CTAs for B TMA loads.
iter_idx (Int): K-iteration index of the current tile.
elect_one_cta (Bool): Whether this CTA is the elected leader for the cluster.
scheduler (TileScheduler[static_MN=scheduler.static_MN, tile_shape=scheduler.tile_shape, cluster=scheduler.cluster, cta_group=scheduler.cta_group, swizzle=scheduler.swizzle, swapAB=scheduler.swapAB]): Tile scheduler tracking the current group and work tile.
expert_ids (LayoutTensor[DType.int32, expert_ids_layout, ImmutAnyOrigin]): Mapping from group index to expert row offset in B.