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

q_gmem_shape

def q_gmem_shape[dtype: DType, swizzle_mode: TensorMapSwizzle, *, group: Int, q_num_heads: Int, depth: Int, decoding: Bool, fuse_gqa: Bool = False]() -> IndexList[Int(4) if decoding or fuse_gqa else Int(3)]

Computes the global-memory shape for a Q tensor TMA tile based on the tile configuration.

Parameters:

dtype (DType): Element type of the Q tensor.
swizzle_mode (TensorMapSwizzle): TMA swizzle mode for the Q tensor tile.
group (Int): Grouped-query attention group size, in query heads per KV head.
q_num_heads (Int): Number of query attention heads.
depth (Int): Head dimension of the attention layer, in elements.
decoding (Bool): Whether the kernel runs in single-token decoding mode.
fuse_gqa (Bool): Whether to fuse grouped-query attention into the tile shape (defaults to False).

Returns:

IndexList[Int(4) if decoding or fuse_gqa else Int(3)]