Version: Nightly

For the complete Mojo documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /docs/manual/basics.md).

create_split_tma

def create_split_tma[rank: Int, dtype: DType, //, smem_shape: IndexList[rank], gmem_shape: IndexList[rank], swizzle_mode: TensorMapSwizzle, fold_chunks: Int = Int(1), row_major: Bool = False](ctx: DeviceContext, ptr: UnsafePointer[Scalar[dtype]], runtime_dim0: Int, out res: TMATensorTile[dtype, rank, _padded_shape[rank, dtype, smem_shape, swizzle_mode](), _ragged_shape[rank, dtype, smem_shape, swizzle_mode]()])

Creates a TMA tensor tile assuming that the first dimension in global memory has UNKNOWN_VALUE.

This function creates a TMATensorTile that optionally splits the last dimension of the tensor into multiples of swizzle granularity. This functionality is currently disabled because it was not found to improve performance.

When fold_chunks >= 2, the contiguous depth chunks are folded into a single rank-4/rank-5 TMA (see the 2-runtime-dim overload's docstring and _create_split_tma_folded). This overload is used by the cache-backed builders where num_heads is the static gmem_shape[1].

Parameters:

rank (Int): The number of dimensions of the tensor.
dtype (DType): The data type of the tensor elements.
smem_shape (IndexList[rank]): The shape of the tile in shared memory.
gmem_shape (IndexList[rank]): The shape of the global memory tensor.
swizzle_mode (TensorMapSwizzle): The swizzling mode for memory access optimization.
fold_chunks (Int): Number of depth chunks to fold into one rank-4 TMA (1 = original 3D behavior).
row_major (Bool): When True (and fold_chunks >= 2), build the rank-5 chunk-inner (row-major-atoms) box so one TMA writes a whole multi-atom-row page; False (default) keeps the rank-4 chunk-outer box.

Args:

ctx (DeviceContext): The CUDA device context used to create the TMA descriptor.
ptr (UnsafePointer[Scalar[dtype]]): Pointer to the global memory tensor data.
runtime_dim0 (Int): The runtime size of the first dimension of the global tensor.

Returns:

TMATensorTile[dtype, rank, _padded_shape[rank, dtype, smem_shape, swizzle_mode](), _ragged_shape[rank, dtype, smem_shape, swizzle_mode]()]: The resulting TMA tensor tile with split layout.

Raises:

If TMA descriptor creation fails.

def create_split_tma[rank: Int, dtype: DType, //, smem_shape: IndexList[rank], gmem_shape: IndexList[rank], swizzle_mode: TensorMapSwizzle, fold_chunks: Int = Int(1), row_major: Bool = False](ctx: DeviceContext, ptr: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], runtime_dim0: Int, runtime_dim1: Int, out res: TMATensorTile[dtype, rank, _padded_shape[rank, dtype, smem_shape, swizzle_mode](), _ragged_shape[rank, dtype, smem_shape, swizzle_mode]()])

Creates a TMA tensor tile assuming that the first two dimensions in global memory has UNKNOWN_VALUE.

When fold_chunks >= 2, the contiguous innermost (depth) dimension — which the swizzle hardware forces to be split into swizzle_granularity-sized chunks — is folded into an extra, non-innermost box dimension so that a single rank-4 cp.async.bulk.tensor copies all fold_chunks depth chunks at once instead of one TMA per chunk. The PUBLIC return type stays rank-3 (SplitLastDimTMATensorTile); the rank-4 CUDA descriptor is built internally and its opaque 128 B TMADescriptor blob (which is rank-agnostic — see TMADescriptor) is wrapped into the rank-3 tile via its @implicit constructor. The issue site (PagedRowIndices._tma_copy_kv_impl) must agree by issuing rank-4 coords; the shared kv_tma_fold_chunks predicate is the single source of truth that keeps the baked rank and the issue rank from drifting. fold_chunks == 1 reproduces exactly the original 3D behavior.

Folding is byte-equivalent to the per-chunk loop ONLY when the box's per-chunk SMEM stride (box_rows * swizzle_granularity) equals the consumer/producer chunk stride (smem_j_stride_rows * swizzle_granularity) — i.e. box_rows == smem_j_stride_rows — and the tile occupies a single page (pages_per_iter == 1). The caller is responsible for only passing fold_chunks >= 2 when those hold; here box_rows equals smem_shape[0].