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

hip_mha_decoding_num_partitions

def hip_mha_decoding_num_partitions(batch_size: Int, num_keys: Int, heads_per_group: Int, sm_count: Int, is_mla: Bool = False) -> Int

Wave-aligned split-K target for MI355X MHA + MLA decode.

Two regimes, distinguished by whether the kernel packs queries by BM (MLA) or spawns one CTA per kv-head (MHA):

MHA-style (heads_per_group < BM): the call comes from get_mha_decoding_num_partitions passing heads_per_group = num_heads // group = kv_num_heads, which is typically small (≤ 8). Each (kv_head, batch) is its own CTA in grid_y, so actual_ctas_per_partition = heads_per_group × batch_size. When work_items = heads_per_group × batch_size ≥ sm_count, one partition already fills the GPU; use few partitions, just enough to amortize key reads. This matches the original heuristic's HIGH_OCC branch (preserved verbatim to avoid regressing the MHA bench grid).
MLA-style (heads_per_group ≥ BM): the call comes from mla.mojo passing heads_per_group = num_heads (≥ BM=32 for h ∈ {32, 64, 128}). MLA packs BM queries into one CTA, so actual_ctas_per_partition = ceildiv(num_heads, BM) × batch_size. Even when work_items looks large (e.g. bs=8 h=64 → 512), actual CTAs are only 16; needs many partitions. Apply the 2-wave wave-aligned formula: one_wave = sm_count // ctas_per_partition two_wave = 2 × sm_count // ctas_per_partition work_floor = ceildiv(pages, MAX_PAGES_PER_SPLIT) np_target = clamp(work_floor, one_wave, two_wave)

EXCEPTION (num_heads <= 16, e.g. Kimi-K2.5 TP=4): the one_wave floor underfills. With num_blocks_y=1, ctas_per_partition = batch_size, so one wave (np = sm/bs) gives each CU exactly one CTA: no second block to overlap HBM-read latency. These shapes are latency-bound, so target two full waves instead: np_target = min(two_wave, pages) Measured on MI355: two-wave np is 5-10% faster than one-wave across bs=4 (32K-128K) and bs=8/16 short context; past two waves regresses on reduce cost. bs=1 (two_wave=512 -> clamps to 256 = one wave) unchanged.

Phase 0 sweep (PARTITIONING_PLAN.md) validated MLA-style at h=64: bs=1 ctx=131K → np=128 (capped, fills GPU at 1-wave + cap) bs=2 ctx=65K → np=64 (one_wave=64 dominates) bs=2 ctx=80K → np=64 (one_wave=64 dominates; work_floor 64 capped) bs=2 ctx=131K → np=128 (work_floor=103 → bucket to 128) bs=8 ctx=80K → np=32 (work_floor 64 capped by two_wave=32) bs=8 ctx=131K → np=32 (work_floor 103 capped by two_wave=32) bs=16 ctx=131K → np=16 (work_floor 103 capped by two_wave=16)

AMD reducer constraint: mla_splitk_reduce supports MAX_PARTITIONS up to parts_per_lane × WARP_SIZE; the 256-partition specialization (parts_per_lane=4) lifts the MLA-style cap to 256. Only nk >= 64K (pages >= 256) actually reaches np=256; smaller nk is page-limited.

Tunables (MLA-style): BM = 32 (MLA decode block-M on MI355) SPLIT_PAGE_SIZE = 256 (min keys per partition) MAX_PAGES_PER_SPLIT = 5 (= 1280 keys per partition cap) MAX_HIP_PARTITIONS = 256 (reducer's MAX_PARTITIONS limit; the MHA-style branch above stays pinned ≤64)

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).
sm_count (Int): Device multiprocessor count used to size the wave-fill target.
is_mla (Bool): Whether the caller is the MLA decode path (defaults to False).

Returns:

Int