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 the warp's base lane in the address and hand thread t lane base + t, so a thread can only touch its own lane. A copy therefore means "these columns of the lane this thread owns", and the TMEM operand must be a warp base, never base + lane. The per-thread view of a warp's (32, N) tile is its row 0 (tile.tile[1, N](Coord(Idx[0], Idx[0]))), whose storage is the unchanged warp base.

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

comptime values

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

• ​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_copy_async: Issues the stores that copy src into the TMEM row dst, without waiting.