RaggedTMA3DTile

struct RaggedTMA3DTile[dtype: DType, swizzle_mode: TensorMapSwizzle, *, BM: Int, BN: Int, middle_dim: Int, group: Int = Int(1), tma_blocks_per_op: Int = Int(0)]

Creates a TMA descriptor for loading/storing from ragged 3D arrays with a ragged leading dimension. This loads 2D tiles, indexing into the middle dim. When using this loads, it is essential that at least BM_seq * stride space has been allocated in front of the gmem pointer, otherwise CUDA_ERROR_ILLEGAL_ADDRESS may result.

The (middle_dim, rows) selector dims are always folded into one outermost descriptor dim (both are box == 1 and GMEM-contiguous), dropping one rank: each copy issues coordinate (ragged_idx + dynamic_dim) * middle_dim + middle_idx on that dim. Fewer descriptor dims means fewer per-issue UMOVs into uniform registers (which ptxas allocates poorly), at no offsetting cost.

When group > 1, the gmem is treated as 4D (rows, middle_dim, group, depth). The smem tile has BM_seq * group = BM rows, where BM_seq = BM // group is the number of distinct sequence positions. The dynamic_dim parameter in copy methods represents valid sequence positions. The descriptor is rank-4 (merged, BM_seq, group, depth), or rank-5 when tma_blocks_per_op > 0.

When tma_blocks_per_op > 0 (only valid for swizzle_mode == SWIZZLE_NONE), the contiguous depth dimension is split into (depth // swizzle_granularity, swizzle_granularity) and a blocks dimension is added to the descriptor box, so a single async_copy_batched copies tma_blocks_per_op swizzle-granularity blocks at once rather than one block per async_copy_from_col. With the selector merge this is rank-4 for group == 1 and rank-5 for group > 1. The blocks dimension's global extent is the true block count, so a box that overhangs the end is masked off by the TMA.

Parameters

• ​dtype (DType): The data type of the tensor.
• ​swizzle_mode (TensorMapSwizzle): The swizzling mode to use for memory access.
• ​BM (Int): The number of rows of the corresponding 2D shared memory tile.
• ​BN (Int): The number of columns of the corresponding 2D shared memory tile.
• ​middle_dim (Int): The middle (head) extent, folded into the ragged-selector coordinate. Each copy issues coordinate (ragged_idx + dynamic_dim) * middle_dim + middle_idx on the merged outermost descriptor dim.
• ​group (Int): The number of heads fused into each sequence position (default 1).
• ​tma_blocks_per_op (Int): Swizzle-granularity blocks copied per async_copy_batched (0 = disabled, use the per-block async_copy_from_col path).

Fields

• ​descriptor (TMADescriptor): The TMA descriptor that will be used to store the ragged tensor.

Implemented traits

AnyType, Copyable, DevicePassable, ImplicitlyCopyable, ImplicitlyDeletable, Movable

comptime members

BM_seq

comptime BM_seq = (BM // group)

Number of distinct sequence positions per tile.

device_type

comptime device_type = RaggedTMA3DTile[dtype, swizzle_mode, BM=BM, BN=BN, middle_dim=middle_dim, group=group, tma_blocks_per_op=tma_blocks_per_op]

The device-side type representation.

layout

comptime layout = tile_layout_k_major[dtype, BM, BN, swizzle_mode]()

The unswizzled-smem layout copied to/from by this tma op.

swizzle_granularity

comptime swizzle_granularity = (swizzle_mode.bytes() // size_of[dtype]())

The number of columns that must be copied at a time due to the swizzle size.

Methods

__init__

def __init__(out self, descriptor: TMADescriptor)

Initializes a new RaggedTMA3DTile with the provided TMA descriptor.

Args:

• ​descriptor (TMADescriptor): The TMA descriptor that defines the memory access pattern.

def __init__(out self, *, copy: Self)

Copy initializes this RaggedTMA3DTile from another instance.

Args:

• ​copy (Self): The other RaggedTMA3DTile instance to copy from.

get_type_name

static def get_type_name() -> String

Returns a string representation of the RaggedTMA3DTile type.

Returns:

String: A string containing the type name with all template parameters.

create

static def create[*, depth: Int = BN](ctx: DeviceContext, ptr: UnsafePointer[Scalar[dtype]], *, rows: Int) -> Self

Create a RaggedTMA3DTile.

Parameters:

• ​depth (Int): The size of the inner-most, contiguous, dimension.

Args:

• ​ctx (DeviceContext): The device context used to create the TMA descriptors.
• ​ptr (UnsafePointer[Scalar[dtype]]): The global memory pointer.
• ​rows (Int): The size of the ragged dimension.

Returns:

Self: A RaggedTMA3DTile corresponding to the gmem.

Raises:

If TMA descriptor creation fails.

async_copy_from_col

def async_copy_from_col[col: Int, eviction_policy: CacheEviction = CacheEviction.EVICT_FIRST](self, src: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED], *, ragged_idx: UInt32, dynamic_dim: UInt32, middle_idx: UInt32, elect: Int32)

Copy a single swizzle_granularity-wide column chunk from smem to gmem.

The TMA store is PTX-predicated on elect, so call this unconditionally from every lane (no warp-divergent if elect != 0: wrapper); only the elected lane issues the copy.

Parameters:

• ​col (Int): Which column chunk (0-indexed, each chunk is swizzle_granularity columns).
• ​eviction_policy (CacheEviction): Optional cache eviction policy that controls how the data is handled in the cache hierarchy. Defaults to EVICT_FIRST.

Args:

• ​src (UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]): Source shared memory pointer (base of the full tile).
• ​ragged_idx (UInt32): Index into the ragged dimension.
• ​dynamic_dim (UInt32): Number of rows (or seq positions when group > 1) to copy.
• ​middle_idx (UInt32): Index into the middle (generally head) dimension.
• ​elect (Int32): 0 on non-elected lanes (skip the TMA), non-zero on the single elected lane (issue the TMA).

async_copy_batched

def async_copy_batched[col_start: Int, eviction_policy: CacheEviction = CacheEviction.EVICT_FIRST](self, src: UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED], *, ragged_idx: UInt32, dynamic_dim: UInt32, middle_idx: UInt32, elect: Int32)

Copy tma_blocks_per_op swizzle_granularity-wide column blocks from smem to gmem in a single TMA, starting at block col_start.

Only valid when tma_blocks_per_op > 0 (SWIZZLE_NONE). The descriptor box covers tma_blocks_per_op blocks; if col_start + tma_blocks_per_op overruns the true block count, the TMA masks the overhang off (no gmem write). With the (middle_dim, rows) selector merge the descriptor is rank-4 for group == 1 and rank-5 for group > 1.

The TMA store is PTX-predicated on elect, so call this unconditionally from every lane (no warp-divergent if elect != 0: wrapper); only the elected lane issues the copy.

Parameters:

• ​col_start (Int): First block (0-indexed, each block is swizzle_granularity columns) copied by this op.
• ​eviction_policy (CacheEviction): Optional cache eviction policy that controls how the data is handled in the cache hierarchy. Defaults to EVICT_FIRST.

Args:

• ​src (UnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]): Source shared memory pointer (base of the full blocked tile).
• ​ragged_idx (UInt32): Index into the ragged dimension.
• ​dynamic_dim (UInt32): Number of rows (or seq positions when group > 1) to copy.
• ​middle_idx (UInt32): Index into the middle (generally head) dimension.
• ​elect (Int32): 0 on non-elected lanes (skip the TMA), non-zero on the single elected lane (issue the TMA).

prefetch_descriptor

def prefetch_descriptor(self)

Prefetches the TMA descriptor into cache.