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_tma

def q_tma[dtype: DType, //, swizzle_mode: TensorMapSwizzle, *, BM: Int, depth: Int, q_num_heads: Int, group: Int, decoding: Bool, fuse_gqa: Bool = False, num_qk_stages: Int = Int(1)](ctx: DeviceContext, ptr: Pointer[Scalar[dtype], _safe=False], rows: Int) -> TMATensorTile[dtype, Int(4) if decoding or fuse_gqa else Int(3), _padded_shape[Int(4) if decoding or fuse_gqa else Int(3), dtype, q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, fuse_gqa=fuse_gqa, num_qk_stages=num_qk_stages](), swizzle_mode](), _ragged_shape[Int(4) if decoding or fuse_gqa else Int(3), dtype, q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, fuse_gqa=fuse_gqa, num_qk_stages=num_qk_stages](), swizzle_mode]()]

Creates a split TMA descriptor for the Q tensor, pairing the shared-memory tile shape with the global-memory layout.

Parameters:

dtype (DType): Element type of the Q tensor (inferred).
swizzle_mode (TensorMapSwizzle): TMA swizzle mode for the Q tensor tile.
BM (Int): Tile block size in the query (row) dimension, in elements.
depth (Int): Head dimension of the attention layer, in elements.
q_num_heads (Int): Number of query attention heads.
group (Int): Grouped-query attention group size, in query heads per KV head.
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).
num_qk_stages (Int): Number of pipeline stages used to split the Q shared-memory tile along the depth dimension (defaults to 1).

Args:

ctx (DeviceContext): Device context used to create the TMA descriptor.
ptr (Pointer[Scalar[dtype], _safe=False]): Base pointer to the Q tensor in global memory.
rows (Int): Number of rows in the Q tensor exposed via the TMA descriptor.

Returns:

TMATensorTile[dtype, Int(4) if decoding or fuse_gqa else Int(3), _padded_shape[Int(4) if decoding or fuse_gqa else Int(3), dtype, q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, fuse_gqa=fuse_gqa, num_qk_stages=num_qk_stages](), swizzle_mode](), _ragged_shape[Int(4) if decoding or fuse_gqa else Int(3), dtype, q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, fuse_gqa=fuse_gqa, num_qk_stages=num_qk_stages](), swizzle_mode]()]