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

get_start_and_end_for_partitions

def get_start_and_end_for_partitions[tile_size: Int](num_keys: Int, num_partitions: Int, partition_idx: Int) -> Tuple[Int, Int]

Calculate start and end indices for a partition.

Non-empty partitions are packed at low indices 0..N-1 with partition_size = max(tile_size, align_up(ceildiv(num_keys, num_partitions), tile_size)); partitions >= N are empty (start == end == num_keys).

Args:

num_keys (Int): Total number of keys (sequence length).
num_partitions (Int): Number of partitions to split keys into.
partition_idx (Int): Index of current partition (0 to num_partitions-1).

Returns:

Tuple[Int, Int]: Tuple of (start_idx, end_idx) for the partition, aligned to tile_size.