cute-dsl-ref（可爱的 DSL 参考）

CuTe Python DSL Reference

Execution Model

CuTe Python DSL is the Python surface for NVIDIA's CuTe layout algebra. Unlike cuTile's block-level abstraction, CuTe DSL exposes explicit thread/warp/warpgroup control, TMA pipelines, barrier choreography, and shared memory management.

Two-Level Host/Device Pattern

Every CuTe DSL kernel has two functions:

1. @cute.jit host function — runs on CPU, sets up TMA descriptors, computes grid, allocates shared memory, launches the kernel
2. @cute.kernel device function — runs on GPU, contains the actual computation

import cutlass
import cutlass.cute as cute

@cute.kernel
def my_kernel(tiled_mma: cute.TiledMma, ...):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()
    # ... GPU code

@cute.jit
def host_fn(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
    # Setup TMA descriptors, compute grid, allocate SMEM
    my_kernel(...).launch(grid=grid_shape, block=block_shape, smem=smem_bytes)

Compilation Pipeline

Three-stage JIT compilation:

1. Pre-Staging: Python AST rewriting
2. Meta-Stage: Python interpreter executes meta-programs (compile-time constants resolved)
3. Object-Stage: Compiler backend generates PTX → SASS

JIT artifacts are cached automatically. Control with CUTE_DSL_CACHE_DIR and CUTE_DSL_DISABLE_FILE_CACHING.

Key Difference from cuTile

Core API Table

Decorators and Compilation

Thread/Block/Grid Indexing

Layout and Tensor Construction

Layout Algebra

MMA Operations

Copy and TMA Operations

Shared Memory

Tensor Memory (SM100 only)

Synchronization

Math (Element-wise on TensorSSA)

All math functions accept fastmath=True for approximate hardware intrinsics.

Tensor Operations

Debugging

Key Constraints

1. @cute.kernel function name must match module filename — same rule as cuTile
2. TMA descriptors are host-side only — create in @cute.jit, pass to @cute.kernel
3. Register budget: 255 max/thread — validate with docs/devices/ specs
4. SMEM limits vary by device and SM version — SM100 (B200/B300): 228 KB/SM, 227 KB/block opt-in; SM120 (RTX 5090): 96 KB
5. Pipeline stages consume SMEM — each stage needs its own buffer; validate total
6. Barrier sync errors are silent — cause incorrect results, not crashes; always test with --check
7. Cluster dimensions must be compatible with grid — cluster shape must evenly divide grid
8. print() is compile-time only — use cute.printf() for device-side output
9. No early-exit breaks in loops — use predication instead
10. 32-bit layout algebra — shapes/strides limited to 32-bit integers
11. Architecture support: Ampere (SM80) and above; SM100 (B200/B300) for full features including tcgen05 and TMEM
12. Architecture-specific kernels: Some kernels target a specific SM version (e.g., tcgen05 MMA ops require SM100, not just "Blackwell"). The "Blackwell" marketing name spans multiple SM versions with different instruction sets: Always check the device before running, compiling, or profiling: nvidia-smi --query-gpu=name --format=csv,noheader. Then match against the kernel's target SM version — the GPU name alone is not sufficient; you must know which SM version it maps to. If the SM version does not match, skip the run/profiling step rather than attempting execution that will fail or produce misleading results.

- SM100 (B200, B300) — datacenter GPUs, supports tcgen05 MMA, TMEM, full cluster features - SM120 (GeForce RTX 5090, RTX 5080) — consumer Blackwell, uses sm_120a, does not support tcgen05 ops

Control Flow

Blackwell Datacenter MMA Operations (tcgen05) — SM100 only

SM100 required. tcgen05 ops are available on B200/B300 (SM100) only. They are not available on consumer Blackwell GPUs like RTX 5090 (SM120/sm_120a).

from cutlass.cute.nvgpu import tcgen05

# Create MMA operation for SM100 (B200/B300) tensor cores
op = tcgen05.MmaF16BF16Op(
    dtype=cutlass.Float16,           # Input type
    acc_dtype=cutlass.Float32,       # Accumulator type
    shape=(128, 128, 64),            # M x N x K tile shape
    cta_group=tcgen05.CtaGroup.ONE,  # ONE or TWO CTAs
)

tiled_mma = cute.make_tiled_mma(op)

Common Patterns

GEMM with TMA Pipeline (Simplified)

@cute.kernel
def gemm_kernel(tiled_mma, tma_a, tma_b, smem_a, smem_b, gmem_c, ...):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, _ = cute.arch.block_idx()

    # Get thread's MMA partition
    thr_mma = tiled_mma.get_slice(tidx)

    # Allocate accumulator in registers
    acc = tiled_mma.partition_shape_C(tile_shape)
    tCrC = cute.make_fragment(acc, cutlass.Float32)

    # K-loop with TMA pipeline
    for k_tile in range(num_k_tiles):
        # TMA copy: global → shared (async)
        cute.copy(tma_a, gA_tile, sA_tile, tma_bar_ptr=barrier)
        cute.copy(tma_b, gB_tile, sB_tile, tma_bar_ptr=barrier)

        # Wait for TMA to complete
        cute.arch.mbar.wait(barrier, phase)

        # Partition shared memory for this thread
        tCsA = thr_mma.partition_A(sA)
        tCsB = thr_mma.partition_B(sB)

        # MMA: accumulate into registers
        cute.gemm(tiled_mma, tCrC, tCsA, tCsB, tCrC)

    # Epilogue: write back to global memory
    cute.basic_copy(tCrC, gC_partition)

Warpgroup Specialization (Producer/Consumer)

warp_idx = cute.arch.warp_idx()
is_producer = warp_idx < num_producer_warps

if is_producer:
    # Issue TMA copies for upcoming pipeline stages
    for stage in range(num_stages):
        cute.copy(tma_atom, src, dst, tma_bar_ptr=barriers[stage])
        cute.arch.mbar.arrive_and_expect_tx(barriers[stage], bytes_per_stage)
else:
    # Consume data from shared memory, execute MMA
    for stage in range(num_stages):
        cute.arch.mbar.wait(barriers[stage], phase)
        cute.gemm(tiled_mma, acc, sA_stage, sB_stage, acc)

Persistent Kernel Loop

@cute.kernel
def persistent_kernel(tiled_mma, num_tiles, ...):
    bidx, _, _ = cute.arch.block_idx()
    num_blocks = cute.arch.grid_dim()[0]

    tile_id = bidx
    while tile_id < num_tiles:
        # Process tile
        # ... TMA copy, MMA, epilogue ...
        tile_id += num_blocks

Detailed References

• Architecture operations (thread indexing, TMA, MMA atoms, barriers, SMEM/TMEM): See references/architecture-ops.md
• Official documentation index: See references/docs-index.md
• Example kernels: See the /learn-cute-dsl skill for a categorized example source index