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 module

tmem_engine

Defines the TensorEngine that views Blackwell Tensor Memory (TMEM).

TMEM is the SM100 accumulator memory: a grid of TMEM_NUM_LANES lanes by TMEM_NUM_COLS columns per CTA, each cell 32 bits. A tile over TMemEngine places its elements on that grid, so its layout has a lane stride of 1 and a column stride of TMEM_NUM_LANES: the whole accumulator is (128, 512):(1, 128), an N-column allocation is (128, N):(1, 128), and nested shapes express the lane placement of other MMA configurations. Lanes are the stride-1 dimension because every allocation has 128 of them, so the conversion between a flat index and a cell is a shift and a mask that no column count enters. The hardware names a cell by a UInt32 whose upper half is the lane and whose lower half is the column; that encoding stays inside the engine, whose offset and distance translate between it and the layout's flat indices.

TMEM has no pointer representation, so a tile over this engine has no element loads or stores; unsafe_ptr rejects them at compile time. Data moves through copy_from. Copying a register or shared-memory tile into a TMEM tile issues tcgen05.st, copying a TMEM tile into one issues tcgen05.ld. Each instruction names at most 64 registers, and a copy moves a row in slices of 64 columns with one wait per slice, so a wide row never holds more than 64 staging registers live. The _async variants issue every instruction of a copy without waiting, so several copies can share one wait issued through wait_store or wait_load. An async copy stages its whole row, so its row is capped at 64 columns; a wider row is tiled into slices with one copy each.

Access is warp-collective: tcgen05.ld and tcgen05.st take a collective base in the address and place each thread's registers relative to it, so a TMEM operand is always a base, never base + lane, and its layout says which cells this thread's registers hold. Two views exist. The row view of a warp's (32, N) tile is its row 0 (tile.tile[1, N](coord[0, 0])): the 32x32b shape hands thread t lane base + t, so the copy means "these columns of the lane this thread owns". The 16x256b view of a (16, N) half of that tile, tmem_16x256b(block), has layout TMem16x256bLayout[N], a column pair by two lanes 8 apart, repeated every 8 columns: the 16x256b shape hands thread t the cells of that layout from lane base + t // 4, column 2 * (t % 4), in the layout's flat order, which is the MMA accumulator fragment order. tmem_cell turns a view coordinate back into the lane and column this thread's element occupies.

The design is written up in Mojo/docs/stdlib/internal/gpua/tmem_engine.md.

comptime values

TMem16x256bLayout

comptime TMem16x256bLayout[num_cols: Int] = Layout(Coord(Coord(), ComptimeInt()), Coord(Coord(), ComptimeInt()))

This thread's 16x256b fragment of a (16, num_cols) lane block.

The first mode is a 2 by 2 quad of cells, a column pair (stride one column) by two lanes 8 apart; the second mode repeats the quad every 8 columns. The offsets are relative to the block base plus the lane and column the hardware adds for this thread, and the flat order, column within the pair fastest, then lane, then repeat, is the order tcgen05.ld.16x256b fills registers in.

Parameters

• ​num_cols (Int): The columns of the block, a multiple of 8.

TMem16x256bRegisters

comptime TMem16x256bRegisters[num_cols: Int] = Layout(Coord(Coord(), ComptimeInt()), Coord(Coord(), ComptimeInt()))

The register tile layout for a TMem16x256bLayout[num_cols] view.

Same shape as the view, packed in the view's flat order, so copying a fragment into a tile with this layout lets a register be addressed by the coordinate of its cell, and copy_from moves it without reordering.

Parameters

• ​num_cols (Int): The columns of the block, a multiple of 8.

TMEM_NUM_COLS

comptime TMEM_NUM_COLS = 512

The number of 32-bit TMEM columns per lane.

TMEM_NUM_LANES

comptime TMEM_NUM_LANES = 128

The number of TMEM lanes (rows) per CTA.

Also the column stride of a layout over TMemEngine, whose flat indices walk the lane-by-column grid lane-first. Every allocation has this many lanes, whatever its column count, so the stride is the same for every tile.

Structs

• ​TMemAccessShape: Names a supported tcgen05.ld or tcgen05.st data movement shape.
• ​TMemEngine: Implements TensorEngine over Tensor Memory via tcgen05.ld/st.
• ​TMemStorage: Encoded TMEM address: lane in the upper 16 bits, column in the lower.

Functions

• ​tmem_16x256b: Returns this thread's 16x256b fragment view of a (16, N) lane block.
• ​tmem_cell: Returns the lane and column this thread's element at coord of a TMEM view occupies.
• ​tmem_copy_async: Issues the stores that copy src into the TMEM row dst, without waiting.