Understanding TMA loads in CuTe DSL

11 Oct, 2026

Moving a matrix tile into shared memory sounds simple. Reading a TMA kernel introduces a few less obvious questions. Why does the copy use a coordinate tensor? Why group the tensor modes before partitioning? Why does the barrier need both an arrival count and a byte count?

In this post, I work through those questions using Yifan Yang's CuTe DSL example. His original article explains TMA loads, prefetch, and multicast in CuTe C++. Here, I focus on the Python load example and its tma_partition API.

This is a code walkthrough. I have not benchmarked the kernel. The snippets below use shorter variable names and omit surrounding setup; the linked file contains the complete example.

What TMA does

The Tensor Memory Accelerator, introduced with NVIDIA Hopper, lets a thread request a multidimensional transfer between global and shared memory. The hardware handles the tensor addressing and moves the data without using registers as an intermediate buffer. NVIDIA describes support for tensors with up to five dimensions. Hopper tuning guide

For a load, the useful sequence is:

describe the tensor and tile
          |
select a tile coordinate
          |
issue the asynchronous transfer
          |
wait for completion
          |
read the tile from shared memory

The request and its completion are separate events. Issuing a load does not make its destination immediately safe to read.

CuTe DSL expresses these operations in Python syntax, then compiles the kernel into GPU code. A CuTe tensor combines an iterator with layout information, and a copy atom describes a hardware copy operation. CUTLASS Python overview

The work assigned to each block

The example uses a contiguous FP16 matrix and a fixed tile size. Its launch has one warp per block. Each block owns a strip of rows and loops across the columns.

These are calculations from the example's constants and launch, rather than measured results.

For block b and loop iteration j, the selected region is:

rows:    [64*b, 64*(b+1))
columns: [256*j, 256*(j+1))

For example, block 1 at iteration 2 selects rows 64 through 127 and columns 512 through 767.

                  seven tiles across K
             0    1    2    3    4    5    6
block 0    [   ][   ][   ][   ][   ][   ][   ]
block 1    [   ][   ][ X ][   ][   ][   ][   ]
block 2    [   ][   ][   ][   ][   ][   ][   ]
   ...
block 131  [   ][   ][   ][   ][   ][   ][   ]

Each box represents 64 rows by 256 columns.

There are 924 tile loads in this decomposition. The blocks can run concurrently; within each block, the loop handles the seven tiles sequentially.

Describe the memory layout first

A layout maps a logical coordinate to an element offset. For a row-major matrix with cols columns:

offset(row, col) = row * cols + col

The corresponding CuTe construction is:

matrix_layout = cute.make_layout((rows, cols), stride=(cols, 1))
matrix_view = cute.make_tensor(input_tensor.iterator, matrix_layout)

The strides are in elements. Moving one column advances by one element; moving one row advances by cols elements. Creating this view describes existing storage. It does not allocate another matrix. CuTe layouts

The shared-memory tile has its own layout:

tile_shape = (64, 256)
tile_layout = cute.make_layout(tile_shape, stride=(256, 1))
shared_layout = cute.make_composed_layout(
    cute.make_swizzle(0, 4, 3),
    0,
    tile_layout,
)

A swizzle permutes layout offsets. Here, its first argument specifies zero swizzle bits, so this composition leaves the offsets unchanged. The tile remains row-major. This gives the example a composed layout without introducing a nontrivial shared-memory permutation. CuTe swizzle API

Describing storage explicitly also creates an assumption: the input must agree with that description. Rebuilding a row-major view around the pointer of a transposed tensor would not make its storage row-major.

Construct the copy atom and coordinate tensor

The host-side setup passes the global tensor, shared layout, and tile shape to make_tiled_tma_atom:

load_atom, coordinate_view = cpasync.make_tiled_tma_atom(
    cpasync.CopyBulkTensorTileG2SOp(),
    matrix_view,
    shared_layout,
    tile_shape,
)

CopyBulkTensorTileG2SOp selects a bulk tensor transfer from global to shared memory. The helper returns a copy atom and a TMA coordinate tensor in the API used by this example. CUTLASS 4.2.1 TMA API

The coordinate tensor deserves a closer look. An ordinary tensor view uses its layout to locate a value in memory. A TMA tensor uses its layout to calculate the coordinates consumed by a TMA instruction. It lets CuTe apply tiling and partitioning algebra to those coordinates. It is not an extra global-memory array of coordinate values. CuTe TMA tensors

This explains why the device code selects a tile from coordinate_view:

tile_coordinates = cute.local_tile(
    coordinate_view,
    tile_shape,
    (row_block, column_tile),
)

local_tile selects a tile using coordinates in the tiled problem. The pair (1, 2) means tile 1 along rows and tile 2 along columns. The corresponding element origin is (64, 512). CuTe tensor tiling

Keeping tile coordinates separate from element coordinates makes the block assignment easier to follow. The loop chooses a tile; the tensor layout supplies the coordinates needed for the transfer.

Allocate the destination and barrier

Inside the kernel, SmemAllocator supplies storage for the tile and a 64-bit memory barrier. Allocation creates the destination view; the later TMA copy fills it. The allocator supports an element type, layout, alignment, and swizzle. Shared-memory allocator API

For this compact FP16 tile, the payload is 32 KiB. A 64-bit barrier adds eight bytes, giving the example's launch request of 32776 bytes. That arithmetic applies to this layout and allocation order. A different layout or alignment can require additional space.

Two byte counts appear in the program:

shared-memory allocation: tile storage + barrier storage
expected transaction:    tile payload only

The barrier tracks transferred data. Its own storage is not part of the transfer.

Initialize the barrier before using it

The initialization pattern is:

with cute.arch.elect_one():
    cute.arch.mbarrier_init(barrier_ptr, 1)

cute.arch.mbarrier_init_fence()
cute.arch.barrier()

elect_one selects one thread per warp. Here, one thread initializes the barrier with one expected arrival. Other threads can wait on it. The fence and block barrier establish initialization before use. CuTe architecture API

Increasing the block size would require revisiting this code. elect_one elects one thread per warp, so several warps could execute an initialization intended to happen once per block.

Before each copy, the elected thread arrives and declares the expected payload:

with cute.arch.elect_one():
    cute.arch.mbarrier_arrive_and_expect_tx(barrier_ptr, payload_bytes)

This combines a software arrival with transaction accounting. The byte count must match the asynchronous work tracked by this phase. Arrival and transaction API

Why group modes before tma_partition?

The copy atom describes the operation. The selected coordinate tile and shared-memory destination still need views arranged for that operation.

The example groups both tile dimensions into the first mode:

shared_group = cute.group_modes(shared_tile, 0, 2)
global_group = cute.group_modes(tile_coordinates, 0, 2)

shared_partition, global_partition = cpasync.tma_partition(
    load_atom,
    0,
    cute.make_layout(1),
    shared_group,
    global_group,
)

group_modes groups modes starting at 0 and stopping before 2. Schematically, (64, 256) becomes ((64, 256),). It preserves the underlying layout mapping; it does not copy data or flatten the tile into a newly allocated vector. Grouping API

The grouping presents the whole two-dimensional tile as the mode to partition for TMA. tma_partition then produces compatible shared and global views for the atom. Its return order is shared first, global second. TMA partition API

There are two coordinate systems here. local_tile chooses a tile in the matrix. The coordinate 0 and size-one CTA layout passed to tma_partition describe participation in this single-CTA arrangement. They do not select matrix tile 0.

Issue the load and wait for its phase

The global partition is the source, so it comes first in cute.copy:

cute.copy(
    load_atom,
    global_partition,
    shared_partition,
    tma_bar_ptr=barrier_ptr,
    mcast_mask=None,
)

cute.arch.mbarrier_wait(barrier_ptr, phase_bit)
phase_bit ^= 1

The barrier pointer connects transfer completion to the wait. The DSL TMA copy path handles single-thread issuance internally, so this call stays outside the explicit election used for barrier setup. CuTe single-thread execution guidance

An mbarrier phase completes when its pending arrival count and transaction count reach zero. For this transfer, the software arrival satisfies the arrival requirement, while TMA completion satisfies the transaction requirement. A successful wait makes the transferred data available to the participating threads. PTX barrier semantics

The phase parity starts at 0 and alternates as the barrier is reused:

The program tracks which phase it needs to finish. Waiting with the same parity on every iteration would lose that distinction. PTX parity waits

If computation consumes the tile, it belongs after the successful wait. Buffer reuse also needs a guarantee that all consumers have finished before a later transfer overwrites their data.

What this example establishes

The source contains loads and waits, with diagnostic tensor printing commented out. It has no output tensor, result comparison, or timing measurement. Its floor-divided launch and loop cover complete tiles only. The supplied dimensions divide evenly. These are boundaries visible in the source.

A correctness experiment could use position-dependent input values, copy each loaded tile to an output tensor after the wait, and compare every element. That would check tile origins and coverage. Timing would be a separate experiment with compilation excluded from the measured region.

The loop immediately waits after issuing each transfer. To overlap loading with computation, an extended kernel would need multiple buffer stages and synchronization governing when each stage is ready and when it is free to overwrite. Asynchronous hardware supplies the mechanism; the kernel schedule determines how much useful overlap occurs.

There is also a version boundary. The source commit mentions the 4.1 API, and the 4.2.1 documentation describes its two-result atom construction. Current documentation describes a TmaInfo return object that still supports two-value unpacking. It also marks cutlass.utils.SmemAllocator as a deprecated alias. Check the installed CUTLASS version when running the example. Current TMA API, current allocator API

The part I find most useful is the relationship between the objects. The layout describes storage. The coordinate tensor lets tiling select the transfer origin. Partitioning prepares the views for the copy atom. The barrier tells the threads when the destination is ready to read. Following those responsibilities makes the kernel much easier to reason about.