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

splitk_partition_idx

def splitk_partition_idx(splitk_partitions: UInt32) -> UInt32

This CTA's split-K partition index [0, splitk_partitions).

Derived from the grid coordinate (block_idx.x % splitk_partitions), NOT block_rank_in_cluster(): the scheduler maps block_idx.x // cluster_size -> tile (cluster_size == splitk_partitions for split-K, since it forces pair_cta=False), so the low bits are the partition. This is correct and CTA-uniform whether or not the launch forms a hardware cluster -- M2 has no cross-CTA traffic, so it does not depend on cluster co-residency. (M4's DSMEM combine will additionally require a real cluster; that is where block_rank_in_cluster() / cluster co-residency re-enters.)

Args:

splitk_partitions (UInt32): Number of split-K partitions (P); the partition index is block_idx.x % splitk_partitions.

Returns:

UInt32