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_mha_decoding_max_num_partitions

def get_mha_decoding_max_num_partitions[num_heads: Int, group: Int](batch_size: Int, num_keys: Int, ctx: DeviceContext) -> Int

Return the maximum number of split-K partitions for CUDA-graph-stable launches.

Returns an upper bound on the partition count that remains constant for a given batch size and key length, allowing the kernel grid to be captured in a CUDA graph. CTAs whose partition index exceeds the runtime num_partitions early-exit without doing work.

Parameters:

num_heads (Int): Total number of query heads.
group (Int): GQA group size (query heads per key/value head).

Args:

batch_size (Int): Number of sequences in the batch.
num_keys (Int): Maximum key-sequence length (cache length).
ctx (DeviceContext): GPU device context used to query SM count and other properties.

Returns:

Int: The stable upper bound on the number of split-K partitions.