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
decode_mma_shape
def decode_mma_shape[dtype: DType, depth: Int, num_heads: Int, mla_mode: Bool = False, fold_wide_mma: Bool = False]() -> IndexList[Int(3)]
Return the MFMA shape the gfx950 decode kernels use for this shape.
Split out of AMDStructuredConfig.get_mma_shape so host-side dispatch can
ask for the shape without building a full config, and so there is one
definition of the rule rather than two that can drift.
Parameters:
- dtype (
DType): Element type shared by Q, K, and V. - depth (
Int): Attention head depth. - num_heads (
Int): Number of query heads. - mla_mode (
Bool): Whether multi-latent attention tiling is active. - fold_wide_mma (
Bool): Whether the MHA token fold's wide arm applies, frommha_decode_fold_wide_mma. Defaults to False, which asks for the unfolded shape — what the host fold gate tests before it knows the width.
Returns:
IndexList[Int(3)]: The (M, N, K) MFMA shape.