For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
mha_decode_fold_wide_mma
def mha_decode_fold_wide_mma[dtype: DType, num_heads: Int, group: Int, q_seq_len: Int]() -> Bool
Whether the MHA decode token fold runs the 32x32x64 MFMA.
MHA only: MLA decode folds too, but with its own block geometry.
At MMA_M == 32 two 16-f32 C-fragments join into the PV A-operand, so P
stays in registers where the 16-row shape must round-trip it through LDS.
_mha_decode_fold_warp_m reads this to pick WM, so the widths on the wide
MMA and the widths given a 32-row M-tile are one set by construction.
Rows must tile 32 — which also puts exactly one M-tile in a warp, all
PRegisterBuffer.mma_tile's gather handles — and leave at least THREE
warps. Two suffices for the kernel but not for the split-K partition count,
which callers derive from a 4-warp CTA (hip_mha_decoding_num_partitions):
a 2-warp CTA handed that count measures +4% against the -22% of its own.
group == num_heads selects the single-KV-head arm, the only one stacking
num_heads * q_seq_len rows and setting WN == BN.
Parameters:
- dtype (
DType): Element type shared by Q, K, and V. - num_heads (
Int): Number of query heads, all owned by the fold's one KV head. - group (
Int): Query heads per KV head. - q_seq_len (
Int): Token SLOTS the M dimension is built from. Callers pass the PADDED tile width, not the tokens a sequence carries — the two differ exactly whenmha_decode_fold_tile_q_seq_lenpads.
Returns:
Bool: True when this fold width takes the 32x32x64 MFMA.