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

mla_decode_dispatch_scalars

def mla_decode_dispatch_scalars(batch_size: Int, max_cache_valid_length: Int, q_max_seq_len: Int, num_heads: Int, is_fp8_kv: Bool, ctx: DeviceContext) -> Tuple[Int, Int, Int]

Compute the packed MLA decode dispatch metadata for the active device.

Returns (batch_size, q_max_seq_len, num_partitions). Dispatches on ctx.api() internally, mirroring mha_decoding_num_partitions:

HIP (AMD): the AMD MLA decode kernel bakes its split-K grid from mha_decoding_num_partitions(..., is_mla=True) at dispatch time and never reads num_partitions from the scalar buffer; the metadata num_partitions is used only as the HIP device-graph capture/replay selection key. To keep the capture key byte-for-byte equal to the baked grid, compute the SAME value the kernel bakes: pass RAW max_cache_valid_length and num_heads (kv_heads == 1 for MLA so heads_per_group == num_heads), with NO q_max_seq_len adjustment.
CUDA (NVIDIA/SM100): delegate to the SM100 runtime heuristic compute_mla_dispatch_scalars_runtime, which reads the device SM count.
Any other api (e.g. Metal): raise: there is no MLA decode dispatch path for non-HIP/non-CUDA devices, so route them to a clear error rather than silently into the SM100 heuristic.

Args:

batch_size (Int): Current decode batch size.
max_cache_valid_length (Int): Max valid KV-cache length across the batch.
q_max_seq_len (Int): Max query sequence length (1 for decode).
num_heads (Int): Number of Q attention heads.
is_fp8_kv (Bool): Whether the KV cache is FP8 (selects the SM100 sub-heuristic).
ctx (DeviceContext): Device context used to read api() and the SM count.

Returns:

Tuple[Int, Int, Int]: The packed (batch_size, q_max_seq_len, num_partitions) metadata.