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

mha_decoding_max_num_partitions

def mha_decoding_max_num_partitions(batch_size: Int, num_keys: Int, heads_per_group: Int, ctx: DeviceContext) -> Int

Returns the num_keys-independent upper bound on mha_decoding_num_partitions.

Used to launch a graph-stable decode grid. Only the CUDA decode path over-launches and early-returns the extra partitions; every other backend keeps max == actual so the max >= actual invariant holds and no over-launch path is taken.

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).
ctx (DeviceContext): Device context used to query the backend API and SM count.

Returns:

Int: Upper bound on the partition count for any num_keys on this backend.