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, c_type: DType, sfa_dtype: DType, sfb_dtype: 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], sfa_tile_rank: Int, sfa_tile_shape: IndexList[sfa_tile_rank], sfa_desc_shape: IndexList[sfa_tile_rank], sfb_tile_rank: Int, sfb_tile_shape: IndexList[sfb_tile_rank], sfb_desc_shape: IndexList[sfb_tile_rank], num_pipeline_stages: Int, group_scale_offsets_layout: Layout, transpose_b: Bool, /, *, a_smem_layout: Layout, b_smem_layout: Layout, sfa_smem_layout: Layout, sfb_smem_layout: Layout, config: BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b], 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_tile_rank, a_tile_shape, a_desc_shape], b_tma_op: TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape], sfa_tma_op: TMATensorTile[sfa_dtype, sfa_tile_rank, sfa_tile_shape, sfa_desc_shape], sfb_tma_op: TMATensorTile[sfb_dtype, sfb_tile_rank, sfb_tile_shape, sfb_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], sfa_smem_base: Pointer[Scalar[sfa_dtype], address_space=AddressSpace.SHARED, _safe=False], sfb_smem_base: Pointer[Scalar[sfb_dtype], 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: UInt32, 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_id: Int32, group_scale_offsets: LayoutTensor[DType.uint32, group_scale_offsets_layout, MutAnyOrigin])

Issues multicast TMA loads for A, B, and their scale factors into shared memory.

Waits for the consumer (MMA) to release the current pipeline stage buffer, then launches asynchronous multicast TMA copies of the A and B tiles along with their corresponding block scale factors into shared memory for one pipeline stage, accounting for expert offsets and group scale offsets.

Parameters:

a_type (DType): The data type of input tensor A.
b_type (DType): The data type of input tensor B.
c_type (DType): The data type of the output tensor C.
sfa_dtype (DType): The data type of the A scale factors.
sfb_dtype (DType): The data type of the B scale factors.
a_tile_rank (Int): The rank of the A TMA tile.
a_tile_shape (IndexList[a_tile_rank]): The shape of the A TMA tile.
a_desc_shape (IndexList[a_tile_rank]): The descriptor shape of the A TMA tile.
b_tile_rank (Int): The rank of the B TMA tile.
b_tile_shape (IndexList[b_tile_rank]): The shape of the B TMA tile.
b_desc_shape (IndexList[b_tile_rank]): The descriptor shape of the B TMA tile.
sfa_tile_rank (Int): The rank of the A scale factor TMA tile.
sfa_tile_shape (IndexList[sfa_tile_rank]): The shape of the A scale factor TMA tile.
sfa_desc_shape (IndexList[sfa_tile_rank]): The descriptor shape of the A scale factor TMA tile.
sfb_tile_rank (Int): The rank of the B scale factor TMA tile.
sfb_tile_shape (IndexList[sfb_tile_rank]): The shape of the B scale factor TMA tile.
sfb_desc_shape (IndexList[sfb_tile_rank]): The descriptor shape of the B scale factor TMA tile.
num_pipeline_stages (Int): The number of load/MMA pipeline stages.
group_scale_offsets_layout (Layout): The layout of the group scale offsets tensor.
transpose_b (Bool): Whether B is stored transposed (K-major).
a_smem_layout (Layout): The shared memory layout for A.
b_smem_layout (Layout): The shared memory layout for B.
sfa_smem_layout (Layout): The shared memory layout for A scale factors.
sfb_smem_layout (Layout): The shared memory layout for B scale factors.
config (BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b]): The block-scaled matmul configuration.
block_tile_shape (IndexList[Int(3)]): The (BM, BN, BK) block tile shape.
mma_shape (IndexList[Int(3)]): The (MMA_M, MMA_N, MMA_K) MMA shape.
num_sf_k_tiles (Int): The number of scale factor tiles along K.
cta_group (Int): The number of CTAs cooperating per output tile (1 or 2).
k_group_size (Int): The number of K iterations loaded per pipeline stage.

Args:

a_tma_op (TMATensorTile[a_type, a_tile_rank, a_tile_shape, a_desc_shape]): The TMA tensor tile descriptor for A.
b_tma_op (TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape]): The TMA tensor tile descriptor for B.
sfa_tma_op (TMATensorTile[sfa_dtype, sfa_tile_rank, sfa_tile_shape, sfa_desc_shape]): The TMA tensor tile descriptor for A scale factors.
sfb_tma_op (TMATensorTile[sfb_dtype, sfb_tile_rank, sfb_tile_shape, sfb_desc_shape]): The TMA tensor tile descriptor for B scale factors.
a_smem_base (Pointer[Scalar[a_type], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the A shared memory buffer.
b_smem_base (Pointer[Scalar[b_type], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the B shared memory buffer.
sfa_smem_base (Pointer[Scalar[sfa_dtype], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the A scale factor shared memory buffer.
sfb_smem_base (Pointer[Scalar[sfb_dtype], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the B scale factor shared memory buffer.
load_mma_pipeline (ProducerConsumerPipeline[num_pipeline_stages]): The producer/consumer pipeline for TMA loads.
peer_cta_coord (Tuple[Int, Int, Int]): The (peer_id, m_coord, n_coord) of the peer CTA.
work_tile_coord (Tuple[Int, Int]): The (m, n) coordinate of the work tile.
a_multicast_mask (UInt16): The multicast mask for A TMA loads.
b_multicast_mask (UInt16): The multicast mask for B TMA loads.
iter_idx (UInt32): The K iteration index for this load.
elect_one_cta (Bool): Whether this CTA is the elected leader CTA.
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]): The tile scheduler tracking group and expert state.
expert_id (Int32): The expert ID for the current work tile.
group_scale_offsets (LayoutTensor[DType.uint32, group_scale_offsets_layout, MutAnyOrigin]): The per-group scale factor offsets tensor.