CUTLASS, in depth: the GEMM hierarchy, CuTe layouts, Hopper kernel schedules and the Python DSL

The problem CUTLASS solves: feeding tensor cores

A GEMM computes D = αAB + βC with A of shape M×K and B of shape K×N. It performs 2MNK floating-point operations but only needs to move (MK + KN + 2MN) elements, so it has enough reuse to be compute-bound. The catch is that the reuse only materialises if each loaded element is used many times while it sits in fast memory. A naive kernel that computes each output element by streaming a row of A and a column of B from global memory reads 2K elements per 2K flops: one flop per element, far below what a modern GPU needs.

Tiling fixes that. A thread block (CTA) owns an output tile of size BM×BN, walks the K dimension in steps of BK, and at each step stages a BM×BK slice of A and a BK×BN slice of B in shared memory. For a 128×128×64 tile in FP16 that is (128·64 + 64·128) · 2 bytes = 32 KB loaded per step, against 2·128·128·64 ≈ 2.1 million flops: 64 flops per byte from global memory, and more once L2 serves the tiles shared between neighbouring CTAs. Inside the CTA the same trick repeats: warps take sub-tiles into registers and issue tensor-core MMA instructions on fragments. See parallel matrix multiplication for the algorithmic background and shared memory and conflict-free tiling for the staging layer.

The copy engine, shared-memory layout, register fragments and MMA instruction must all agree on which element lives where, and the answer changes each generation: Ampere uses cp.async and mma.sync, Hopper adds TMA, thread-block clusters and warpgroup MMA (wgmma), and datacenter Blackwell moves accumulators into dedicated tensor memory. CUTLASS encodes those choices once.

The hierarchy: device, kernel, collective, atom

CUTLASS 3.x organises a GEMM into four layers, each a template that takes the one below it as a parameter.

• Device layer (GemmUniversalAdapter): the host-side handle; validates arguments, sizes workspace, launches.
• Kernel layer (GemmUniversal): owns the grid, persistence and the tile scheduler.
• Collective layer: the mainloop pipelines A and B tiles through shared memory and issues MMAs over K; the epilogue applies α, β and fused operations and writes D.
• Atoms: an MMA atom wraps one tensor-core instruction; a copy atom wraps one copy instruction such as a TMA load.

CuTe: layouts as functions

CuTe is the layout library underneath CUTLASS 3.x and 4.x. A layout is a pair (shape : stride) that maps a logical coordinate to a linear offset. A 4×8 column-major matrix is (4,8):(1,4): coordinate (i, j) maps to i·1 + j·4. The same matrix row-major is (4,8):(8,1). Shapes and strides can nest, so ((2,2),8):((1,16),2) describes a matrix whose rows are interleaved in pairs, which is exactly the kind of arrangement a tensor-core fragment or a bank-conflict-free shared-memory swizzle needs.

A tensor is a pointer (to global, shared or register memory) plus a layout. Tiling a tensor, partitioning it among threads and matching it to an MMA atom are all operations on layouts: composition, logical division and products. Functions such as local_tile cut a CTA's tile out of a global tensor, and a tiled MMA or tiled copy object partitions that tile so each thread knows which fragment it owns. Because these are compile-time layout transformations, a mismatch between a copy and an MMA is usually a compile error, not silent garbage.

auto layout = cute::make_layout(cute::make_shape(4, 8), cute::make_stride(1, 4));
cute::print_layout(layout);   // prints the 4x8 grid of offsets 0..31, column-major

When a kernel misbehaves, printing thread 0's partitioned tensors tells you more than any profiler.

A complete Hopper GEMM with the builder API

The following is a full type definition for an FP16 GEMM with FP32 accumulation on Hopper. It follows the official example 48, with FP16 inputs. The epilogue is built first because its shared memory is carved out of the mainloop's pipeline budget.

// includes: cutlass/gemm/{device,kernel,collective} and cutlass/epilogue/collective headers
using namespace cute;

// D = alpha * A @ B + beta * C, fp16 inputs, fp32 accumulate, Hopper (sm_90a)
using ElementA = cutlass::half_t;  using LayoutA = cutlass::layout::RowMajor;
using ElementB = cutlass::half_t;  using LayoutB = cutlass::layout::ColumnMajor;
using ElementC = cutlass::half_t;  using LayoutC = cutlass::layout::ColumnMajor;
using ElementAcc = float;
constexpr int AlignA = 128 / cutlass::sizeof_bits<ElementA>::value;  // 8 elements = 16 bytes
constexpr int AlignB = 128 / cutlass::sizeof_bits<ElementB>::value;
constexpr int AlignC = 128 / cutlass::sizeof_bits<ElementC>::value;

using TileShape    = Shape<_128, _128, _64>;   // CTA tile M x N x K
using ClusterShape = Shape<_2, _1, _1>;        // 2 CTAs share B via TMA multicast

using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
    cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
    TileShape, ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
    ElementAcc, ElementAcc,
    ElementC, LayoutC, AlignC,
    ElementC, LayoutC, AlignC,
    cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;

using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
    cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
    ElementA, LayoutA, AlignA,
    ElementB, LayoutB, AlignB,
    ElementAcc, TileShape, ClusterShape,
    cutlass::gemm::collective::StageCountAutoCarveout<
        static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
    cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;

using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
    Shape<int, int, int>, CollectiveMainloop, CollectiveEpilogue>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

Alignment means each row or column starts on a 16-byte boundary, which TMA requires. The cluster shape groups two CTAs along M so they share loads of the same B tile through TMA multicast (see the Tensor Memory Accelerator). StageCountAutoCarveout fits as many pipeline stages as shared memory allows after the epilogue's share, and KernelScheduleAuto lets it choose a warp-specialized schedule. Host code then looks like this:

auto hw = cutlass::KernelHardwareInfo::make_kernel_hardware_info<Gemm::GemmKernel>(0);
typename Gemm::Arguments args{
    cutlass::gemm::GemmUniversalMode::kGemm,
    {M, N, K},
    {ptrA, strideA, ptrB, strideB},
    {{alpha, beta}, ptrC, strideC, ptrD, strideD},
    hw};
args.scheduler.max_swizzle_size = 4;   // L2-friendly tile order

Gemm gemm;
CUTLASS_CHECK(gemm.can_implement(args));          // rejects bad alignment or shapes
size_t ws = Gemm::get_workspace_size(args);
cutlass::device_memory::allocation<uint8_t> workspace(ws);
CUTLASS_CHECK(gemm.initialize(args, workspace.get()));
CUTLASS_CHECK(gemm.run());

At run time a producer warp TMA-loads A and B tiles into a ring of shared-memory stages and signals barriers; consumer warpgroups wait, issue wgmma into register accumulators and release each stage; after the last K step the epilogue applies α and β and stores D.

Kernel schedules: who does what inside a CTA

On Hopper the mainloop is warp-specialized: some warps only move data and others only compute, a design covered in warp specialization architecture. CUTLASS exposes the choice as a kernel schedule type in cutlass::gemm, which you can pass instead of KernelScheduleAuto.

Persistent kernels launch roughly one CTA per SM and loop over output tiles supplied by a tile scheduler. Rasterisation and swizzling (raster_order, max_swizzle_size) order tiles so those sharing A rows or B columns are in L2 together. Split-K and stream-K divide K across CTAs when there are too few output tiles to fill the GPU, at the cost of a reduction and workspace.

Worked example: tile choice and wave quantisation

Take an H100 SXM, which has 132 SMs, and two shapes from a transformer layer. For a square 4096×4096×4096 GEMM with 128×128 tiles there are 32·32 = 1,024 output tiles. With one CTA per SM that is 1,024 / 132 ≈ 7.76 waves: seven full waves, then a final wave in which only 100 of 132 SMs have work. With 128×256 tiles there are 512 tiles and 3.88 waves: more reuse per tile, but twice the accumulator registers per CTA.

Now a decode-time shape: M = 4096 output features, N = 16 tokens, K = 4096. With 128-wide tiles in N, seven-eighths of every tile is padding, and there are only 32 tiles for 132 SMs, so three quarters of the GPU idles. The fixes are a narrow N tile, split-K or stream-K for parallelism along K, or treating the operation as memory-bound. No tile shape suits every problem, which is why libraries ship many instantiations and pick per shape.

Epilogue fusion: the reason many teams adopt CUTLASS

Writing D and then launching a second kernel to add a bias, apply GELU and cast to FP8 costs a full round trip of the output through HBM. CUTLASS epilogues can do that work on the accumulators while they are still in registers. CUTLASS 3.x expresses these as an Epilogue Visitor Tree: a small graph of nodes such as load C, load a bias vector, add, activate, compute an absolute maximum for FP8 scaling, and store. Tensor core dtype paths explains why FP8 outputs need per-tensor or per-block scale factors, which are exactly what these epilogues compute.

The Python side: CuTe DSL

CUTLASS 4 introduced the CuTe DSL, the same layout and atom concepts embedded in Python and JIT-compiled. It ships as the nvidia-cutlass-dsl package and began as a beta, so check release notes before depending on specific APIs. @cute.kernel functions run on the device, @cute.jit functions run host-side logic and launch kernels, and cute.compile produces a callable you reuse. A first program:

import torch
import cutlass
import cutlass.cute as cute

@cute.kernel
def add_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    i = bidx * bdim + tidx
    m, n = gA.shape
    gC[i // n, i % n] = gA[i // n, i % n] + gB[i // n, i % n]

@cute.jit
def add(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
    threads = 256
    m, n = mA.shape
    add_kernel(mA, mB, mC).launch(grid=((m * n) // threads, 1, 1), block=(threads, 1, 1))

a = torch.randn(2048, 2048, device="cuda", dtype=torch.float16)
b = torch.randn_like(a); c = torch.zeros_like(a)
tensors = [cute.runtime.from_dlpack(t) for t in (a, b, c)]
compiled = cute.compile(add, *tensors)   # JIT once, reuse many times
compiled(*tensors)

This naive version assumes the element count divides by 256. The official notebooks rewrite it with tiled layouts for vectorised loads, then build tensor-core GEMMs for Hopper and Blackwell. The appeal is iteration speed: compiles in seconds, and PyTorch tensors passed directly through DLPack.

Choosing and profiling kernels

CUTLASS ships a profiler that runs every compiled candidate on your shape and reports throughput. Restrict the build with CMake filters; compiling everything takes hours.

cmake .. -DCUTLASS_NVCC_ARCHS=90a -DCUTLASS_LIBRARY_KERNELS="cutlass3x_sm90_tensorop_*gemm_f16_f16_f32*"
make cutlass_profiler -j

./tools/profiler/cutlass_profiler --kernels="cutlass3x_sm90_tensorop_*gemm_f16*" \
    --m=4096 --n=16 --k=4096 --output=decode_shape.csv

Profile the shapes and batch sizes you actually serve, record the winner per shape and instantiate only those. Check the winner in Nsight Compute; the occupancy guide explains why a tensor-core kernel at low occupancy can still be at peak.

Failure modes

• Wrong architecture target. Hopper wgmma and TMA kernels need sm_90a, not plain sm_90. Built for the wrong target, the SM90 kernel paths compile to an error stub and the kernel reports an error or produces nothing useful at run time rather than failing the build.
• Alignment and stride rejections. can_implement returns an error when a leading dimension is not a multiple of the alignment, for example an FP16 tensor with 4,095 columns. Pad the tensor or instantiate a lower-alignment kernel, which will be slower. Always check the return status.
• Shared-memory overflow. Large tiles, many stages and a heavy epilogue can exceed the per-CTA limit; hand-set stage counts are the usual culprit.
• Numerics. Accumulating in FP16 to save registers loses precision on long K; FP8 without correct scale factors overflows silently. Verify against a reference on real activations.
• Bloat and porting. Every instantiation is a separate kernel, so keep a curated list; and a Hopper-tuned kernel needs re-profiling on Blackwell.

Trade-offs: CUTLASS, cuBLAS or Triton

Start with cuBLAS. Move to CUTLASS when profiling shows a hot kernel needing a fusion, type or shape cuBLAS serves badly.

What to do next

1. Profile your model and list the GEMM shapes that dominate runtime, with their dtypes and the operations that follow them.
2. Clone CUTLASS, build example 48 for your architecture and confirm it runs and verifies.
3. Read and print a few CuTe layouts until (shape : stride) notation feels natural.
4. Build the profiler with a narrow kernel filter and sweep your real shapes; record the best kernel per shape and compare it with cuBLAS.
5. If a follow-on elementwise kernel is significant, prototype an Epilogue Visitor Tree fusion and measure the saved HBM traffic.
6. Try the CuTe DSL notebooks for fast iteration, then decide whether the production kernel lives in C++ or Python.