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

promote_accumulators

def promote_accumulators[pipeline_stages: Int, num_accum_pipeline_stages: Int, accum_type: DType, accum_layout: Layout, a_scales_type: DType, b_scales_type: DType, b_scales_layout: Layout, expert_ids_layout: Layout, /, *, a_scales_smem_layout: Layout, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], cta_group: Int, CLUSTER_SIZE: Int32, is_lower_frag_required: Bool, num_output_warps: Int](b_scales: LayoutTensor[b_scales_type, b_scales_layout, ImmutAnyOrigin], b_scales_n: Int, a_scales_smem_base: Pointer[Scalar[a_scales_type], address_space=AddressSpace.SHARED, _safe=False], c_upper_main_tile: LayoutTensor[accum_type, accum_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=c_upper_main_tile.element_layout, layout_int_type=c_upper_main_tile.layout_int_type, linear_idx_type=c_upper_main_tile.linear_idx_type, masked=c_upper_main_tile.masked, alignment=c_upper_main_tile.alignment], c_lower_main_tile: LayoutTensor[accum_type, accum_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=c_lower_main_tile.element_layout, layout_int_type=c_lower_main_tile.layout_int_type, linear_idx_type=c_lower_main_tile.linear_idx_type, masked=c_lower_main_tile.masked, alignment=c_lower_main_tile.alignment], mma_output_pipeline: ProducerConsumerPipeline[num_accum_pipeline_stages], tmem_addr: UInt32, load_mma_pipeline: ProducerConsumerPipeline[pipeline_stages], work_tile_coord: Tuple[Int, Int], elect_one_warp: Bool, stage_stride_cols: Int, k_iter: Int, problem_shape: StaticTuple[Int32, Int(3)], expert_ids: LayoutTensor[DType.int32, expert_ids_layout, ImmutAnyOrigin], 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])

Loads MMA outputs from TMEM and accumulates them into register C fragments with blockwise FP8 scaling.

Reads the per-stage TMEM accumulator fragments, fetches the matching A scales from SMEM and B scales from global memory, multiplies them into a per-fragment scale, and accumulates the scaled products into the upper and lower register tiles for the epilogue.

Parameters:

pipeline_stages (Int): Number of load-MMA pipeline buffer stages (inferred).
num_accum_pipeline_stages (Int): Number of MMA-output pipeline buffer stages (inferred).
accum_type (DType): Accumulator dtype for the MMA result (inferred); must be float32.
accum_layout (Layout): Layout of the upper and lower accumulator register tiles (inferred); shape is (num_stages, num_elements).
a_scales_type (DType): Element dtype of the A per-block scales (inferred).
b_scales_type (DType): Element dtype of the B per-block scales (inferred).
b_scales_layout (Layout): Layout of the B per-block scales tensor (inferred).
expert_ids_layout (Layout): Layout of the expert-IDs mapping tensor (inferred).
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.
CLUSTER_SIZE (Int32): Number of CTAs in the cluster, used to gate the load-mbar arrive.
is_lower_frag_required (Bool): Whether the lower accumulator fragment must be loaded and accumulated.
num_output_warps (Int): Number of epilogue output warps participating in the promotion.

Args:

b_scales (LayoutTensor[b_scales_type, b_scales_layout, ImmutAnyOrigin]): Per-block scales for B of shape [num_experts, N // BN, K // BK], read to fetch the B scale for the current K block.
b_scales_n (Int): N dimension of the B scales tensor (columns per expert), used to compute the expert row offset into b_scales.
a_scales_smem_base (Pointer[Scalar[a_scales_type], address_space=AddressSpace.SHARED, _safe=False]): Base pointer to the A scales SMEM buffer.
c_upper_main_tile (LayoutTensor[accum_type, accum_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=c_upper_main_tile.element_layout, layout_int_type=c_upper_main_tile.layout_int_type, linear_idx_type=c_upper_main_tile.linear_idx_type, masked=c_upper_main_tile.masked, alignment=c_upper_main_tile.alignment]): Upper accumulator register tile of shape (num_stages, num_elements) that the scaled MMA fragments are accumulated into.
c_lower_main_tile (LayoutTensor[accum_type, accum_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=c_lower_main_tile.element_layout, layout_int_type=c_lower_main_tile.layout_int_type, linear_idx_type=c_lower_main_tile.linear_idx_type, masked=c_lower_main_tile.masked, alignment=c_lower_main_tile.alignment]): Lower accumulator register tile of shape (num_stages, num_elements) that the scaled MMA fragments are accumulated into when is_lower_frag_required.
mma_output_pipeline (ProducerConsumerPipeline[num_accum_pipeline_stages]): Producer-consumer pipeline carrying TMEM accumulator fragments from the MMA warp.
tmem_addr (UInt32): Base TMEM address of the accumulator allocation.
load_mma_pipeline (ProducerConsumerPipeline[pipeline_stages]): Producer-consumer pipeline between the load and MMA warps, used to track the A scales SMEM stage.
work_tile_coord (Tuple[Int, Int]): Work tile coordinates as (m, n); n selects the B scale index.
elect_one_warp (Bool): Whether this warp is the elected leader warp.
stage_stride_cols (Int): TMEM column stride between consecutive MMA-output pipeline stages.
k_iter (Int): K-iteration index of the current tile, used to index the A and B scales.
problem_shape (StaticTuple[Int32, Int(3)]): Problem shape as (M, N, K) of the grouped GEMM.
expert_ids (LayoutTensor[DType.int32, expert_ids_layout, ImmutAnyOrigin]): Mapping from group index to expert row offset in B.
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.