Atom Operations

CuTe DSL atom operations for copy and matrix multiply-accumulate

Atoms are the fundamental building blocks in CuTe DSL that define hardware-level operations. They encapsulate the metadata and behavior of copy instructions and MMA (Matrix Multiply-Accumulate) operations.

Overview

An Atom consists of:

• An Operation (pure Python class modeling a specific instruction)
• A Trait (wraps IR value with metadata using CuTe Layouts)

Atoms are composed into Tiled operations that map across thread blocks.

Copy Atoms

make_copy_atom

Creates a copy atom from a copy operation.

copy_atom = cute.make_copy_atom(copy_op, element_type)

copy_op · CopyOp · required

Copy operation (e.g., cute.nvgpu.CopyUniversalOp())

element_type · Type[Numeric] · required

Element data type (e.g., cute.Float32, cute.BFloat16)

Common Copy Operations:

# Universal copy (works across memory spaces)
copy_atom = cute.make_copy_atom(
    cute.nvgpu.CopyUniversalOp(),
    cute.Float32
)

# CP.ASYNC for global -> shared memory
copy_atom = cute.make_copy_atom(
    cute.nvgpu.cpasync.CopyG2SAsync(),
    cute.BFloat16
)

# TMA (Tensor Memory Accelerator) for Hopper+
tma_atom = cute.make_copy_atom(
    cute.nvgpu.tcgen05.TMA_LOAD(),
    cute.Float16
)

make_tiled_copy

Creates a tiled copy operation with explicit tiling and thread-value layout.

tiled_copy = cute.make_tiled_copy(
    copy_atom=copy_atom,
    thr_layout=thread_layout,
    val_layout=value_layout
)

copy_atom · CopyAtom · required

Base copy atom

thr_layout · Layout · required

Thread layout mapping thread IDs to coordinates

val_layout · Layout · required

Value layout defining vectorization per thread

Example:

# Thread layout: 4x32 threads, column-major
thr_layout = cute.make_layout((4, 32), stride=(32, 1))

# Value layout: 4x4 values per thread, row-major for vectorization
val_layout = cute.make_layout((4, 4), stride=(4, 1))

# Create tiled copy
copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cute.Float32)
tiled_copy = cute.make_tiled_copy(copy_atom, thr_layout, val_layout)

make_tiled_copy_tv

Creates a tiled copy with automatic TV (Thread-Value) layout inference.

tiled_copy = cute.make_tiled_copy_tv(
    copy_atom,
    thr_layout,
    val_layout
)

Automatically infers the tiler and TV layout from thread and value layouts.

Specialized Tiled Copy Constructors

# Source-specialized copy
tiled_copy_S = cute.make_tiled_copy_S(copy_atom, thr_layout, val_layout)

# Destination-specialized copy  
tiled_copy_D = cute.make_tiled_copy_D(copy_atom, thr_layout, val_layout)

# Specialized for tensor A/B/C in GEMM
tiled_copy_A = cute.make_tiled_copy_A(copy_atom, ...)
tiled_copy_B = cute.make_tiled_copy_B(copy_atom, ...)
tiled_copy_C = cute.make_tiled_copy_C(copy_atom, ...)

TiledCopy Operations

# Get per-thread slice
thr_copy = tiled_copy.get_slice(thread_idx)

# Partition source tensor
thrA_src = thr_copy.partition_S(block_tensor)

# Partition destination tensor
thrA_dst = thr_copy.partition_D(fragment)

# Execute copy
cute.copy(copy_atom, thrA_src, thrA_dst)

MMA Atoms

make_mma_atom

Creates an MMA atom from an MMA operation.

mma_atom = cute.make_mma_atom(mma_op)

mma_op · MmaOp · required

MMA operation (e.g., warp or warpgroup MMA)

Ampere Warp-Level MMA:

# FP32 accumulator, FP32 inputs
mma_atom = cute.make_mma_atom(
    cute.nvgpu.warp.MMA_F32F32F32_M16N8K8_TN()
)

# FP32 accumulator, BF16 inputs
mma_atom = cute.make_mma_atom(
    cute.nvgpu.warp.MMA_F32BF16BF16_M16N8K16_TN()
)

Hopper Warpgroup MMA:

mma_atom = cute.make_mma_atom(
    cute.nvgpu.warpgroup.MMA_F32F16F16_M64N256K32_TN()
)

Blackwell tcgen05 MMA:

mma_atom = cute.make_mma_atom(
    cute.nvgpu.tcgen05.MMA_F32F16F16_M64N256K32_SS()
)

make_tiled_mma

Creates a tiled MMA operation.

tiled_mma = cute.make_tiled_mma(
    mma_atom,
    thr_layout=None  # Optional: custom thread layout
)

mma_atom · MmaAtom · required

Base MMA atom

thr_layout · Layout · default: "None"

Optional custom thread layout (auto-derived if None)

Example:

# Create MMA atom
mma_atom = cute.make_mma_atom(
    cute.nvgpu.warp.MMA_F32F32F32_M16N8K8_TN()
)

# Create tiled MMA
tiled_mma = cute.make_tiled_mma(mma_atom)

# Get per-thread slice
thr_mma = tiled_mma.get_slice(thread_idx)

TiledMma Operations

# Partition A tensor
thrA = thr_mma.partition_A(block_A)

# Partition B tensor
thrB = thr_mma.partition_B(block_B)

# Partition C (accumulator) tensor
thrC = thr_mma.partition_C(block_C)

# Execute GEMM
cute.gemm(tiled_mma.atom, thrC, thrA, thrB, thrC)

Runtime State Management

Some atoms have runtime state (e.g., tcgen05 MMA accumulate flag).

set

Sets runtime fields of an atom.

tiled_mma.set(field, value)

Example:

# Set accumulate mode for tcgen05 MMA
tiled_mma.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True)

get

Gets runtime fields of an atom.

value = tiled_mma.get(field)

Example:

# Get accumulate mode
accum = tiled_mma.get(cute.nvgpu.tcgen05.Field.ACCUMULATE)

with_

Returns a new atom with modified runtime state.

new_atom = atom.with_(field=value)

Architecture-Specific Atoms

Ampere (SM80)

# Warp-level tensor core MMA
mma = cute.make_mma_atom(cute.nvgpu.warp.MMA_F32F16F16_M16N8K16_TN())

# Async copy (cp.async)
copy = cute.make_copy_atom(cute.nvgpu.cpasync.CopyG2SAsync(), cute.Float16)

Hopper (SM90)

# Warpgroup MMA
mma = cute.make_mma_atom(cute.nvgpu.warpgroup.MMA_F32F16F16_M64N256K32_TN())

# TMA load
tma_load = cute.make_copy_atom(cute.nvgpu.tcgen05.TMA_LOAD(), cute.Float16)

# TMA store
tma_store = cute.make_copy_atom(cute.nvgpu.tcgen05.TMA_STORE(), cute.Float16)

Blackwell (SM100)

# tcgen05 MMA with scaled inputs
mma = cute.make_mma_atom(cute.nvgpu.tcgen05.MMA_F32F16F16_M64N256K32_SS())

# Configure accumulate mode
mma.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True)

Complete Example: Tiled Copy

import cutlass.cute as cute

@cute.kernel
def copy_kernel(
    gSrc: cute.Tensor,
    gDst: cute.Tensor,
    thr_layout: cute.Layout,
    val_layout: cute.Layout,
):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    
    # Create copy atom
    copy_atom = cute.make_copy_atom(
        cute.nvgpu.CopyUniversalOp(),
        gSrc.element_type
    )
    
    # Create tiled copy
    tiled_copy = cute.make_tiled_copy_tv(
        copy_atom,
        thr_layout,
        val_layout
    )
    
    # Get thread slice
    thr_copy = tiled_copy.get_slice(tidx)
    
    # Partition tensors
    blkSrc = gSrc[((None, None), bidx)]
    blkDst = gDst[((None, None), bidx)]
    
    thrSrc = thr_copy.partition_S(blkSrc)
    thrDst = thr_copy.partition_D(blkDst)
    
    # Allocate fragment
    frag = cute.make_fragment_like(thrSrc)
    
    # Execute copy
    cute.copy(copy_atom, thrSrc, frag)
    cute.copy(copy_atom, frag, thrDst)

Complete Example: Tiled MMA

import cutlass.cute as cute

@cute.kernel
def mma_kernel(
    mA: cute.Tensor,
    mB: cute.Tensor,
    mC: cute.Tensor,
):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    
    # Create MMA atom
    mma_atom = cute.make_mma_atom(
        cute.nvgpu.warp.MMA_F32F16F16_M16N8K16_TN()
    )
    
    # Create tiled MMA
    tiled_mma = cute.make_tiled_mma(mma_atom)
    
    # Get thread slice
    thr_mma = tiled_mma.get_slice(tidx)
    
    # Partition block tensors
    blkA = mA[((None, None, None), bidx)]
    blkB = mB[((None, None, None), bidx)]
    blkC = mC[((None, None), bidx)]
    
    thrA = thr_mma.partition_A(blkA)
    thrB = thr_mma.partition_B(blkB)
    thrC = thr_mma.partition_C(blkC)
    
    # Allocate fragments
    fragA = cute.make_fragment_like(thrA[None, None, 0])
    fragB = cute.make_fragment_like(thrB[None, None, 0])
    fragC = cute.make_fragment_like(thrC)
    
    # Initialize accumulator
    cute.basic_copy(thrC, fragC)
    
    # GEMM loop
    K_tiles = cute.size(thrA.shape[2])
    for k in range(K_tiles):
        # Load tiles
        cute.basic_copy(thrA[None, None, k], fragA)
        cute.basic_copy(thrB[None, None, k], fragB)
        
        # Compute
        cute.gemm(tiled_mma.atom, fragC, fragA, fragB, fragC)
    
    # Store result
    cute.basic_copy(fragC, thrC)

See Also

• Operations Reference - High-level DSL operations
• Memory Operations - Memory management
• Kernel Definition - Defining GPU kernels