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 struct

TMAReduceExecutor

struct TMAReduceExecutor[c_type: DType, c_smem_dim0: Int, c_smem_dim1: Int, epc: EpilogueConfig, stage_contiguous_size: Int, c_swizzle: TensorMapSwizzle, batched: Bool = False]

Execute TMA reduce-add from SMEM to GMEM.

Mirrors TMAStoreExecutor but uses cp_async_bulk_tensor_reduce_global_shared_cta (add) instead of cp_async_bulk_tensor_global_shared_cta (store). Takes a typed TMATensorTile value (not a raw pointer) so the descriptor keeps its grid_constant provenance end-to-end -- otherwise the compiler drops the constant-memory optimization and each TMA issue refetches the descriptor. Only supports non-transpose path.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

`comptime` members

`c_smem_shape0`

comptime c_smem_shape0 = c_smem_dim0

`CG1_TMA_BM`

comptime CG1_TMA_BM = TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].c_smem_shape0

`CG2_TMA_BM`

comptime CG2_TMA_BM = TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].c_smem_shape0 if (epc == Int(256)) else epc.BM

`cta_group`

comptime cta_group = epc.cta_group

`stageN`

comptime stageN = epc.stageN

`TMA_BM`

comptime TMA_BM = TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].c_smem_shape0 if (eq epc.MMA_M, 256) else epc.BM if (epc == Int(2)) else TMAReduceExecutor[c_type, c_smem_dim0, c_smem_dim1, epc, stage_contiguous_size, c_swizzle, batched].CG1_TMA_BM

Methods

`execute`

static def execute[tma_rank: Int, tile_shape: IndexList[tma_rank], desc_shape: IndexList[tma_rank]](c_smem_tile: TileTensor[c_type, address_space=AddressSpace.SHARED, linear_idx_type=c_smem_tile.linear_idx_type], store_coords: TMAStoreCoords[epc, Self.c_smem_shape0, batched], c_tma_op: TMATensorTile[c_type, tma_rank, tile_shape, desc_shape], warp_id: UInt32, lane: UInt32)

Execute TMA reduce-add from SMEM to GMEM via typed descriptor.

Implemented traits​

comptime members​

c_smem_shape0​

CG1_TMA_BM​

CG2_TMA_BM​

cta_group​

stageN​

TMA_BM​

Methods​

execute​