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_SFB

def load_AB_SFA_SFB[a_type: DType, b_type: DType, sfa_dtype: DType, sfb_dtype: DType, sfa_tma_dtype: DType, sfb_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], sfb_rank: Int, sfb_tile_shape: IndexList[sfb_rank], sfb_desc_shape: IndexList[sfb_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], sfb_tma_op: TMATensorTile[sfb_tma_dtype, sfb_rank, sfb_tile_shape, sfb_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], sfb_smem_tiles: SMemTileArrayWithLayout[sfb_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 A, B, and their scale factors (SFA, SFB) into a pipeline stage of shared memory.

Parameters:

a_type (DType): Element dtype of the A operand matrix (inferred).
b_type (DType): Element dtype of the B operand matrix (inferred).
sfa_dtype (DType): Element dtype of the A scale factors (inferred).
sfb_dtype (DType): Element dtype of the B scale factors (inferred).
sfa_tma_dtype (DType): Element dtype used for the SFA TMA descriptor; may differ from sfa_dtype (for example uint16 for 4D TMA) (inferred).
sfb_tma_dtype (DType): Element dtype used for the SFB TMA descriptor; may differ from sfb_dtype (for example uint16 for 4D TMA) (inferred).
a_rank (Int): Tensor rank of the A operand TMA descriptor (inferred).
a_tile_shape (IndexList[a_rank]): Per-tile shape of the A TMA load (inferred).
a_desc_shape (IndexList[a_rank]): Full descriptor shape of the A TMA load (inferred).
b_rank (Int): Tensor rank of the B operand TMA descriptor (inferred).
b_tile_shape (IndexList[b_rank]): Per-tile shape of the B TMA load (inferred).
b_desc_shape (IndexList[b_rank]): Full descriptor shape of the B TMA load (inferred).
sfa_rank (Int): Tensor rank of the SFA TMA descriptor (inferred).
sfa_tile_shape (IndexList[sfa_rank]): Per-tile shape of the SFA TMA load (inferred).
sfa_desc_shape (IndexList[sfa_rank]): Full descriptor shape of the SFA TMA load (inferred).
sfb_rank (Int): Tensor rank of the SFB TMA descriptor (inferred).
sfb_tile_shape (IndexList[sfb_rank]): Per-tile shape of the SFB TMA load (inferred).
sfb_desc_shape (IndexList[sfb_rank]): Full descriptor shape of the SFB TMA load (inferred).
a_dim0 (Int): Row count of each A SMEM tile (inferred).
a_dim1 (Int): Column count of each A SMEM tile (inferred).
a_num_tiles (Int): Total number of A SMEM tiles across all pipeline stages (inferred).
a_swizzle_bytes (Int): Swizzle stride in bytes for the A SMEM tiles (inferred).
b_dim0 (Int): Row count of each B SMEM tile (inferred).
b_dim1 (Int): Column count of each B SMEM tile (inferred).
b_num_tiles (Int): Total number of B SMEM tiles across all pipeline stages (inferred).
b_swizzle_bytes (Int): Swizzle stride in bytes for the B SMEM tiles (inferred).
num_pipeline_stages (Int): Number of producer/consumer stages in the A/B/SFA/SFB load and MMA pipeline (inferred).
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.
sfb_tma_op (TMATensorTile[sfb_tma_dtype, sfb_rank, sfb_tile_shape, sfb_desc_shape]): TMA tensor tile descriptor for loading SFB (B 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.
sfb_smem_tiles (SMemTileArrayWithLayout[sfb_dtype]): SMEM tile array holding the B scale-factor tiles.
load_mma_pipeline (ProducerConsumerPipeline[num_pipeline_stages]): Producer/consumer pipeline synchronizing A/B/SFA/SFB 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.