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

create_tma_descriptor

def create_tma_descriptor[dtype: DType, rank: Int, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE, l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE, unpack_fp4: Bool = False](global_buf: DeviceBuffer[dtype], global_shape: IndexList[rank], global_strides: IndexList[rank], shared_mem_shape: IndexList[rank]) -> TMADescriptor

Creates a TMA descriptor for tiled memory operations.

Encodes tensor layout information into a 128-byte TMA descriptor that can be used with TMA hardware instructions to efficiently copy data between global and shared memory on NVIDIA GPUs.

The descriptor specifies a mapping from a tile in shared memory to a region in global memory, including dimensions, strides, data type, and optional swizzling for bank conflict avoidance.

Parameters:

• ​dtype (DType): The element data type of the tensor.
• ​rank (Int): The number of dimensions (1-5).
• ​swizzle_mode (TensorMapSwizzle): The swizzle pattern to apply in shared memory.
• ​l2_promotion (TensorMapL2Promotion): L2 cache promotion hint for TMA loads. Defaults to NONE.
• ​unpack_fp4 (Bool): When True, global_buf holds nibble-packed E2M1 (two values per uint8) and the copy pads it on the way into shared memory: each 16-value group lands in 16 bytes as 8 packed bytes followed by an 8-byte gap (see PACKED_FP4_ALIGN16B). The values stay nibble-packed; it is the padding that makes a K extent span one byte per element. global_shape and shared_mem_shape are then counted in FP4 ELEMENTS on the innermost dimension, while global_strides stays in uint8 units.

Args:

• ​global_buf (DeviceBuffer[dtype]): Device buffer containing the global memory tensor.
• ​global_shape (IndexList[rank]): Dimensions of the tensor in global memory.
• ​global_strides (IndexList[rank]): Strides (in elements) for each dimension in global memory. The tensor must be row-major (stride at innermost dimension equals 1).
• ​shared_mem_shape (IndexList[rank]): Dimensions of the tile to be copied to shared memory.

Returns:

TMADescriptor: A TMA descriptor configured for the specified tensor layout.

Raises:

An error if the descriptor creation fails.