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

consumer_main_loop

def consumer_main_loop[accum_type: DType, c_type: DType, a_type: DType, b_type: DType, sfa_dtype: DType, sfb_dtype: DType, a_swizzle: TensorMapSwizzle, b_swizzle: TensorMapSwizzle, transpose_b: Bool, pipeline_stages: Int, scaling_kind: UMMAKind, /, *, a_smem_layout: Layout, b_smem_layout: Layout, sfa_smem_layout: Layout, sfb_smem_layout: Layout, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], SFA_NUM_COLS: Int, SFB_NUM_COLS: Int, cta_group: Int = Int(1), cluster_shape: IndexList[Int(3)] = Index[Int, Int, Int](Int(1), Int(1), Int(1)), k_group_size: Int = Int(1)](tmem_addr: UInt32, sfa_tmem: UInt32, sfb_tmem: UInt32, 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[pipeline_stages], mma_op: MmaOpSM100_BlockScaled_SS[c_type, a_type, b_type, sfa_dtype, sfb_dtype, scaling_kind, block_tile_shape, mma_shape, accum_type=accum_type, cta_group=cta_group, cluster_shape=cluster_shape, a_swizzle=a_swizzle, b_swizzle=b_swizzle, transpose_b=transpose_b], elect_one_warp: Bool, iter_idx: UInt32, k_start: UInt32, work_tile_coord: Tuple[Int, Int])

Consumes one pipeline stage and issues the block-scaled MMA into TMEM.

Waits for the producer (TMA load) to fill the current stage, then reads A, B, and their scale factors from shared memory and launches the block-scaled MmaOpSM100_BlockScaled_SS into tensor memory, accumulating across K-group iterations and committing the result to the accumulator pipeline barrier.

Parameters:

accum_type (DType): The accumulator data type held in TMEM.
c_type (DType): The data type of the output tensor C.
a_type (DType): The data type of input tensor A.
b_type (DType): The data type of input tensor B.
sfa_dtype (DType): The data type of the A scale factors.
sfb_dtype (DType): The data type of the B scale factors.
a_swizzle (TensorMapSwizzle): The TMA swizzle mode for A.
b_swizzle (TensorMapSwizzle): The TMA swizzle mode for B.
transpose_b (Bool): Whether B is stored transposed (K-major).
pipeline_stages (Int): The number of load/MMA pipeline stages.
scaling_kind (UMMAKind): The UMMA scaling kind for block-scaled MMA.
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.
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.
SFA_NUM_COLS (Int): The number of TMEM columns for A scale factors.
SFB_NUM_COLS (Int): The number of TMEM columns for B scale factors.
cta_group (Int): The number of CTAs cooperating per output tile (1 or 2).
cluster_shape (IndexList[Int(3)]): The cluster shape for multicast MMA.
k_group_size (Int): The number of K iterations per pipeline stage.

Args:

tmem_addr (UInt32): The base TMEM address for the accumulator.
sfa_tmem (UInt32): The base TMEM address for A scale factors.
sfb_tmem (UInt32): The base TMEM address 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[pipeline_stages]): The producer/consumer pipeline for TMA loads.
mma_op (MmaOpSM100_BlockScaled_SS[c_type, a_type, b_type, sfa_dtype, sfb_dtype, scaling_kind, block_tile_shape, mma_shape, accum_type=accum_type, cta_group=cta_group, cluster_shape=cluster_shape, a_swizzle=a_swizzle, b_swizzle=b_swizzle, transpose_b=transpose_b]): The block-scaled SM100 MMA operation object.
elect_one_warp (Bool): Whether the calling warp is the elected leader warp.
iter_idx (UInt32): The K iteration index for this MMA.
k_start (UInt32): The K iteration at which to initialize the accumulator.
work_tile_coord (Tuple[Int, Int]): The (m, n) coordinate of the work tile.