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_num_partitions

def mha_decoding_num_partitions(batch_size: Int, num_keys: Int, heads_per_group: Int, ctx: DeviceContext, is_mla: Bool = False) -> Int

Dispatches to the backend-specific decode split-K partition heuristic.

Queries the device multiprocessor count lazily and routes CUDA and HIP to their tuned heuristics; every other backend (Metal, accelerators with no split-K decode path) returns 1 so the decode kernel runs unsplit.

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, num_heads for MLA).
ctx (DeviceContext): Device context used to query the backend API and SM count.
is_mla (Bool): Whether the caller is the MLA decode path (HIP only).

Returns:

Int: Number of split-K partitions for the decode kernel's split-K loop.