tilelang.cuda.intrinsics.layout.mma_sm100_layout¶

CuTe TMEM fragment layouts for dense SM100 MMA (tmem_frg family).

A TMEM fragment is a cute.Layout with cute.ScaledBasis strides, mapping logical buffer coordinates to hierarchical (datapath, column) TMEM coordinates, with the column measured in the buffer’s value type. Exactly as in CUTLASS mma_traits_sm100.hpp, an atom in virtual TMEM addressing (datapath stride one, column stride 128) is repeated over the tile counts, then composed with tmem_restride into physical coordinates. Both TMEM allocation and MMA address formation consume the same layout, and a sliced operand is handled by cute.restrict like any other CuTe layout.

The atoms are the PTX “TMEM data path layout organization” variants (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-organization):

PTX

CUTLASS atom

shape:stride (virtual addressing)

Layout A

2SM M256 (M128/CTA)

(128, N):(1, 128)

Layout B

2SM M128 (M64/CTA, “2x2”)

(64, (N/2, 2)):(1, (128, 64))

Layout C

2SM M64 (M32/CTA, “1x4”)

(32, (N/4, 4)):(1, (128, 32))

Layout D

1SM M128

(128, N):(1, 128)

Layout E

1SM M64 .ws C/D

(64, (N/2, 2)):(1, (128, 64))

Layout F

1SM M64 half-datapath

((16, 4), N):((1, 32), 128)

Layout G

1SM M32 .ws C/D

(32, (N/4, 4)):(1, (128, 32))

Only the compiler’s MMA lowering (this module and tcgen05_macro_generator) sees CuTe layouts; everywhere else — the layout map, TMEM allocation, and copy lowering — a fragment travels as its TileLang lambda *indices: [datapath, column] form via cute.Layout.to_tilelang(), and cute.Layout.from_tilelang_hierarchical() recovers the hierarchical CuTe layout losslessly.

As everywhere else in the MMA lowering, the last two modes of a buffer are the matrix modes; any leading modes are batch dimensions (for example a software-pipeline stage) and repeat the fragment along TMEM columns.

Classes¶

TCGEN05TmemAllocMode

Dense TMEM allocation modes from CUTLASS UMMA::TmemAllocMode.

TCGEN05Meta

The TCGEN5MMA instruction atom selected by get_tcgen5_mma_meta.

Functions¶

make_tmem_frg_atom(num_sms, alloc_mode, atom_m, atom_n)

Return CUTLASS tmem_frg's (tmem_atom, outer_tile_stride) pair.

make_tmem_frg_ws_atom(atom_m, atom_n)

Return CUTLASS tmem_frg_ws's dense 1SM weight-stationary C/D atom.

make_tmem_frg(buffer, num_sms, alloc_mode, atom_m, ...)

Tile a legal ordinary CUTLASS atom over the buffer's matrix modes.

make_tmem_frg_a(buffer, meta)

Build the exact dense TS FrgTypeA selected by CUTLASS.

make_tmem_frg_c(buffer, meta, *, is_ts)

Build the dense FrgTypeC paired with the selected instruction.

validate_tcgen05_ts_instruction(meta, a_dtype, transposed)

Validate a dense TS instruction atom against CUTLASS SM100 wrappers.

Module Contents¶

class tilelang.cuda.intrinsics.layout.mma_sm100_layout.TCGEN05TmemAllocMode¶

Bases: enum.Enum

Dense TMEM allocation modes from CUTLASS UMMA::TmemAllocMode.

INTERLEAVED = 'interleaved'¶
NON_INTERLEAVED = 'non_interleaved'¶
DUPLICATED = 'duplicated'¶
class tilelang.cuda.intrinsics.layout.mma_sm100_layout.TCGEN05Meta¶

The TCGEN5MMA instruction atom selected by get_tcgen5_mma_meta.

atom_m: int¶
atom_n: int¶
atom_k: int¶
enable_ws: bool¶
enable_2cta: bool¶
classmethod from_ffi(values)¶
Parameters:

values (collections.abc.Iterable[int])

Return type:

TCGEN05Meta

property num_sms: int¶
Return type:

int

property atom_m_per_cta: int¶
Return type:

int

property b_atom_n_per_cta: int¶

Each 2SM instruction reads an N/2 shard of B from each CTA.

Return type:

int

tilelang.cuda.intrinsics.layout.mma_sm100_layout.make_tmem_frg_atom(num_sms, alloc_mode, atom_m, atom_n)¶

Return CUTLASS tmem_frg’s (tmem_atom, outer_tile_stride) pair.

atom_m is the per-CTA M extent. The 1SM M64 atom is PTX Layout F (half datapaths, 32-datapath halves interleaved); the 2SM atoms fold the peer CTA’s N shard into the upper datapaths (PTX Layouts A/B/C).

Parameters:
Return type:

tuple[tilelang.layout.cute.Layout, int]

tilelang.cuda.intrinsics.layout.mma_sm100_layout.make_tmem_frg_ws_atom(atom_m, atom_n)¶

Return CUTLASS tmem_frg_ws’s dense 1SM weight-stationary C/D atom.

.ws instructions with M below 128 fold N shards into the upper datapaths instead of leaving them idle: PTX Layout E for M64 and Layout G for M32.

Parameters:

  • atom_m (int)

  • atom_n (int)

Return type:

tilelang.layout.cute.Layout

tilelang.cuda.intrinsics.layout.mma_sm100_layout.make_tmem_frg(buffer, num_sms, alloc_mode, atom_m, atom_n, storage_bits)¶

Tile a legal ordinary CUTLASS atom over the buffer’s matrix modes.

Parameters:

  • buffer (tvm.tirx.Buffer)

  • num_sms (int)

  • alloc_mode (TCGEN05TmemAllocMode)

  • atom_m (int)

  • atom_n (int)

  • storage_bits (int)

Return type:

tilelang.layout.cute.Layout

tilelang.cuda.intrinsics.layout.mma_sm100_layout.make_tmem_frg_a(buffer, meta)¶

Build the exact dense TS FrgTypeA selected by CUTLASS.

Parameters:
Return type:

tilelang.layout.cute.Layout

tilelang.cuda.intrinsics.layout.mma_sm100_layout.make_tmem_frg_c(buffer, meta, *, is_ts)¶

Build the dense FrgTypeC paired with the selected instruction.

Ordinary SS accumulators use the interleaved allocation; dense TS traits pin CUTLASS’s NonInterleaved (1SM) / Interleaved (2SM) modes; .ws SS instructions use the weight-stationary fragment (PTX Layouts E/G).

Parameters:
Return type:

tilelang.layout.cute.Layout

tilelang.cuda.intrinsics.layout.mma_sm100_layout.validate_tcgen05_ts_instruction(meta, a_dtype, transposed)¶

Validate a dense TS instruction atom against CUTLASS SM100 wrappers.

Parameters:
Return type:

None