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

cuda_mha_decoding_num_partitions

def cuda_mha_decoding_num_partitions(batch_size: Int, num_keys: Int, heads_per_group: Int, sm_count: Int) -> Int

Returns the CUDA decode split-K partition count for the given shape.

Takes the num_keys-independent upper bound from cuda_mha_decoding_max_num_partitions and further limits it so each partition spans at least 512 keys. The max(1, ...) floor preserves the

= 1 guard when num_keys < 512.

Args:

batch_size (Int): Number of decode requests in the batch.
num_keys (Int): Number of key cache entries to scan.
heads_per_group (Int): Query heads sharing each kv-head group (kv_num_heads for MHA).
sm_count (Int): Device multiprocessor count used to size the SM-fill target.

Returns:

Int: Partition count in [1, 32] further limited by num_keys // 512.