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: TensorLayout, a_scales_type: DType, b_scales_type: DType, b_scales_layout: TensorLayout, /, *, 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: TileTensor[b_scales_type, b_scales_layout, Storage=b_scales.Storage, address_space=b_scales.address_space, linear_idx_type=b_scales.linear_idx_type], a_scales_smem_tiles: SMemTileArray2DRowMajor[a_scales_type], c_upper_main_tile: TileTensor[accum_type, accum_layout, Storage=c_upper_main_tile.Storage, address_space=AddressSpace.LOCAL, linear_idx_type=c_upper_main_tile.linear_idx_type], c_lower_main_tile: TileTensor[accum_type, accum_layout, Storage=c_lower_main_tile.Storage, address_space=AddressSpace.LOCAL, linear_idx_type=c_lower_main_tile.linear_idx_type], 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)])

Promotes FP8 MMA partial products by applying blockwise A and B scales to the tensor-memory accumulators.

Parameters:

pipeline_stages (Int): Number of double-buffered stages in the load-MMA pipeline.
num_accum_pipeline_stages (Int): Number of double-buffered stages in the MMA-output pipeline.
accum_type (DType): Element type of the accumulators (float32).
accum_layout (TensorLayout): Memory layout of the accumulator TileTensor.
a_scales_type (DType): Element type of the A blockwise scales (float32).
b_scales_type (DType): Element type of the B blockwise scales (float32).
b_scales_layout (TensorLayout): Memory layout of the B scales TileTensor.
block_tile_shape (IndexList[Int(3)]): GEMM block tile shape (BM, BN, BK).
mma_shape (IndexList[Int(3)]): Tensor-core MMA shape (MMA_M, MMA_N, MMA_K).
cta_group (Int): Number of CTAs cooperating per MMA group.
CLUSTER_SIZE (Int32): Total number of CTAs in the cluster.
is_lower_frag_required (Bool): Whether the lower accumulator fragment must be processed.
num_output_warps (Int): Number of warps participating in the epilogue.

Args:

b_scales (TileTensor[b_scales_type, b_scales_layout, Storage=b_scales.Storage, address_space=b_scales.address_space, linear_idx_type=b_scales.linear_idx_type]): B blockwise scales used to rescale FP8 partial products.
a_scales_smem_tiles (SMemTileArray2DRowMajor[a_scales_type]): Shared-memory tile array holding A blockwise scales loaded for the current pipeline stage.
c_upper_main_tile (TileTensor[accum_type, accum_layout, Storage=c_upper_main_tile.Storage, address_space=AddressSpace.LOCAL, linear_idx_type=c_upper_main_tile.linear_idx_type]): Upper accumulator fragment to accumulate scaled results into.
c_lower_main_tile (TileTensor[accum_type, accum_layout, Storage=c_lower_main_tile.Storage, address_space=AddressSpace.LOCAL, linear_idx_type=c_lower_main_tile.linear_idx_type]): Lower accumulator fragment to accumulate scaled results into.
mma_output_pipeline (ProducerConsumerPipeline[num_accum_pipeline_stages]): Producer/consumer pipeline synchronizing MMA production with accumulator consumption.
tmem_addr (UInt32): Base tensor-memory address for loading accumulator fragments.
load_mma_pipeline (ProducerConsumerPipeline[pipeline_stages]): Producer/consumer pipeline synchronizing TMA loads with MMA consumption.
work_tile_coord (Tuple[Int, Int]): Work tile coordinate (m, n) of the current output tile.
elect_one_warp (Bool): Whether this warp is the elected leader for barrier signaling.
stage_stride_cols (Int): Column stride between accumulator pipeline stages in tensor memory.
k_iter (Int): K-dimension iteration index for the current scale application.
problem_shape (StaticTuple[Int32, Int(3)]): Full GEMM problem dimensions (M, N, K).