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_partial

def load_AB_partial[a_type: DType, b_type: DType, a_scales_type: DType, b_tile_rank: Int, b_tile_shape: IndexList[b_tile_rank], b_desc_shape: IndexList[b_tile_rank], num_pipeline_stages: Int, expert_ids_layout: Layout, a_gmem_layout: Layout, a_scales_gmem_layout: Layout, /, *, a_smem_layout: Layout, b_smem_layout: Layout, a_scales_smem_layout: Layout, block_tile_shape: IndexList[Int(3)], cta_group: Int = Int(1), a_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE](a_gmem: LayoutTensor[a_type, a_gmem_layout, ImmutAnyOrigin], a_scales_gmem: LayoutTensor[a_scales_type, a_scales_gmem_layout, ImmutAnyOrigin], b_tma_op: TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_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], 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], expert_end_row: Int, m_tile_global_start: Int)

Sibling to load_AB for tiles the full-TMA path can't handle: fills A and a_scales SMEM via a cooperative warp copy from gmem and issues TMA only for B.

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).
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).
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_gmem_layout (Layout): Layout of the A activation tensor in global memory (inferred).
a_scales_gmem_layout (Layout): Layout of the A per-block scales tensor in global memory (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).
cta_group (Int): CTA group size, 1 or 2 (defaults to 1).
a_swizzle (TensorMapSwizzle): Swizzle mode for the A SMEM buffer, used to compute physical store offsets in the cooperative warp copy (defaults to SWIZZLE_NONE).

Args:

a_gmem (LayoutTensor[a_type, a_gmem_layout, ImmutAnyOrigin]): A activation tensor in global memory, read by the cooperative warp copy.
a_scales_gmem (LayoutTensor[a_scales_type, a_scales_gmem_layout, ImmutAnyOrigin]): A per-block scales tensor in global memory, read by the cooperative warp copy.
b_tma_op (TMATensorTile[b_type, b_tile_rank, b_tile_shape, b_desc_shape]): TMA descriptor for loading B weight 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).
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.
expert_end_row (Int): Global row index one past the last row of the current expert; rows at or past it are zeroed during the cooperative copy.
m_tile_global_start (Int): Global M-row start of the current tile within the A buffer.