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

tma_tile_qo

def tma_tile_qo[dtype: DType, //, swizzle_mode: TensorMapSwizzle, *, BM: Int, BK: Int, depth: Int](ctx: DeviceContext, ptr: Pointer[Scalar[dtype], _safe=False], rows: Int, out res: TMATensorTile[dtype, Int(2), IndexList(BM, BK, __list_literal__=NoneType(None)), _default_desc_shape[Int(2), dtype, IndexList(BM, BK, __list_literal__=NoneType(None)), swizzle_mode]()])

Creates a TMA descriptor for the Q or output tensor used in MLA decode.

Parameters:

dtype (DType): Element type of the Q or output tensor (inferred).
swizzle_mode (TensorMapSwizzle): TMA swizzle mode applied to the descriptor.
BM (Int): Tile height in rows for each TMA copy.
BK (Int): Tile width in columns for each TMA copy.
depth (Int): Column count of the full Q or output tensor.

Args:

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

Returns:

TMATensorTile[dtype, Int(2), IndexList(BM, BK, __list_literal__=NoneType(None)), _default_desc_shape[Int(2), dtype, IndexList(BM, BK, __list_literal__=NoneType(None)), swizzle_mode]()]