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

compute_tma_tile_dims

def compute_tma_tile_dims[BM: Int, BN: Int, MMA_M: Int, OutputM: Int, CLUSTER_M: Int, CLUSTER_N: Int, cta_group: Int, AB_swapped: Bool = False]() -> StaticTuple[Int, Int(3)]

Compute TMA tile dimensions (a_tile_dim0, b_tile_dim0, c_tile_dim0).

Parameters:

BM (Int): Block size along the M dimension of the matmul tile.
BN (Int): Block size along the N dimension of the matmul tile.
MMA_M (Int): M dimension of the MMA instruction shape, used to select the output tile strategy.
OutputM (Int): M dimension of the output tensor, used as the C tile size when the MMA shape or CTA group allows it.
CLUSTER_M (Int): Number of CTAs in the cluster along the M dimension.
CLUSTER_N (Int): Number of CTAs in the cluster along the N dimension.
cta_group (Int): Number of cooperating CTAs in a group, either 1 or 2 for dual-CTA mode.
AB_swapped (Bool): Whether the A and B operands are swapped, selecting the full-output tile path (defaults to False).

Returns:

StaticTuple[Int, Int(3)]: StaticTuple of (a_tile_dim0, b_tile_dim0, c_tile_dim0).