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_window

def splitk_window(T: UInt32, num_partitions: UInt32, partition_idx: UInt32) -> Tuple[UInt32, UInt32]

Front-loaded balanced split of the combined K-tile range [0, T).

Partition p owns the BN-tile window [cb, ce) where the first r = T % num_partitions partitions get q+1 = ceil(T/P) tiles and the rest get q = floor(T/P):

cb = p*q     + min(p,   r)
ce = (p+1)*q + min(p+1, r)

The chunks differ by at most one tile (balanced load), but tiles are filled leading partition first: for T >= 1 partition 0 is always non-empty (it owns tile 0), and only trailing partitions go empty (cb == ce) once T < num_partitions. Keeping the writer (rank 0) non-empty lets the cross-CTA combine (which hardcodes own = rank 0) stay valid; empty trailing partitions stage a neutral identity and the writer weights them to zero. M6 routes idle CTAs (partition_idx >= num_partitions) through that same neutral path.

T is a tile count (small) and num_partitions <= 16, so the products cannot overflow UInt32. num_partitions is comptime at every call site, so the ///% lower to multiply-shift, not real divides.

Args:

T (UInt32): Total number of K-tiles to divide across partitions.
num_partitions (UInt32): Number of split-K partitions (P).
partition_idx (UInt32): This CTA's partition index [0, num_partitions).

Returns:

Tuple[UInt32, UInt32]