For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

copy_accum_to_gmem

def copy_accum_to_gmem[c_type: DType, c_tile_rank: Int, c_tile_shape: IndexList[c_tile_rank], c_desc_shape: IndexList[c_tile_rank], num_accum_pipeline_stages: Int, c_tensor_layout: Layout, /, *, c_smem_layout: Layout, repeat: Int, accum_type: DType, cta_group: Int, epilogue_dtype: DType, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], num_output_warps: Int, c_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, elementwise_compute_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]] = None, register_based_epilogue: Bool = True, transpose_c: Bool = False, scale_c_coord: Bool = True](c_smem_base: Pointer[Scalar[c_type], MutAnyOrigin, address_space=AddressSpace.SHARED], c_tma_op: TMATensorTile[c_type, c_tile_rank, c_tile_shape, c_desc_shape], c: LayoutTensor[c_type, c_tensor_layout, MutAnyOrigin], mma_output_pipeline: ProducerConsumerPipeline[num_accum_pipeline_stages], mma_output_stage: UInt32, tmem_offset: UInt32, c_coord: Tuple[UInt32, UInt32], c_shape: Tuple[UInt32, UInt32], expert_scale: Float32, group_end_idx: UInt32)

Copies one accumulator stage from TMEM to global memory via shared memory.

Loads the MMA result from tensor memory, applies the per-expert scale and an optional elementwise epilogue, packs the result through stmatrix into shared memory, and stores it to global memory with TMA (or a CUDA-core fallback for unaligned tails). Handles both the transposed (transpose_c) and non-transposed output layouts.

Parameters:

• ​c_type (DType): The data type of the output tensor C.
• ​c_tile_rank (Int): The rank of the C TMA tile.
• ​c_tile_shape (IndexList[c_tile_rank]): The shape of the C TMA tile.
• ​c_desc_shape (IndexList[c_tile_rank]): The descriptor shape of the C TMA tile.
• ​num_accum_pipeline_stages (Int): The number of accumulator pipeline stages.
• ​c_tensor_layout (Layout): The layout of the output LayoutTensor.
• ​c_smem_layout (Layout): The layout of the C shared memory tile.
• ​repeat (Int): The number of tcgen05_ld repetitions per stage.
• ​accum_type (DType): The accumulator data type held in TMEM.
• ​cta_group (Int): The number of CTAs cooperating per output tile (1 or 2).
• ​epilogue_dtype (DType): The dtype used for the cast output before storing.
• ​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.
• ​num_output_warps (Int): The number of epilogue warps.
• ​c_swizzle (TensorMapSwizzle): The TMA swizzle mode for C.
• ​elementwise_compute_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]]): Optional fused elementwise epilogue.
• ​register_based_epilogue (Bool): Whether the epilogue runs in registers.
• ​transpose_c (Bool): Whether the output is stored transposed.
• ​scale_c_coord (Bool): Whether to scale tile coordinates by block tile strides.

Args:

• ​c_smem_base (Pointer[Scalar[c_type], MutAnyOrigin, address_space=AddressSpace.SHARED]): Base pointer to the C shared memory buffer.
• ​c_tma_op (TMATensorTile[c_type, c_tile_rank, c_tile_shape, c_desc_shape]): The TMA tensor tile descriptor for C.
• ​c (LayoutTensor[c_type, c_tensor_layout, MutAnyOrigin]): The output LayoutTensor in global memory.
• ​mma_output_pipeline (ProducerConsumerPipeline[num_accum_pipeline_stages]): The producer/consumer pipeline for MMA output.
• ​mma_output_stage (UInt32): The accumulator pipeline stage to drain.
• ​tmem_offset (UInt32): The column offset into TMEM for this stage.
• ​c_coord (Tuple[UInt32, UInt32]): The (row, column) tile coordinate of the output tile.
• ​c_shape (Tuple[UInt32, UInt32]): The (M, N) extent of the output tensor.
• ​expert_scale (Float32): The per-expert scale factor applied to the accumulator.
• ​group_end_idx (UInt32): The exclusive end row index for the current expert group.