GEMM Kernel in CuTe DSL

Build a complete dense GEMM kernel using CuTe DSL with tensor core operations

This guide demonstrates how to implement a dense General Matrix Multiply (GEMM) kernel using CUTLASS CuTe DSL. The example shows a complete implementation for both Ampere and Blackwell architectures.

Overview

The GEMM kernel computes C = A * B where:

• Matrix A is MxKxL (L is batch dimension)
• Matrix B is NxKxL
• Matrix C is MxNxL

Key features:

• Utilizes GPU tensor cores for MMA operations
• Multi-stage pipelining to overlap compute and memory access
• Shared memory buffering for efficient data transfer
• Support for various data types (fp16, bf16, fp8, int8, tf32)

Ampere Architecture Example

Blackwell Architecture Example

For Blackwell GPUs, the implementation uses more advanced features:

class DenseGemmKernel:
    def __init__(
        self,
        acc_dtype: Type[cutlass.Numeric],
        use_2cta_instrs: bool,
        mma_tiler_mn: Tuple[int, int],
        cluster_shape_mn: Tuple[int, int],
        use_tma_store: bool,
    ):
        self.acc_dtype = acc_dtype
        self.use_2cta_instrs = use_2cta_instrs
        self.cluster_shape_mn = cluster_shape_mn
        self.mma_tiler_mn = mma_tiler_mn
        self.use_tma_store = use_tma_store
        self.cta_group = (
            tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
        )

Key Blackwell Features

Tensor Memory Access (TMA)

TMA provides efficient memory operations:

# Setup TMA load for A
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
    self.cluster_shape_mn, tiled_mma.thr_id
)
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
    a_op, a, a_smem_layout, self.mma_tiler, tiled_mma,
    self.cluster_layout_vmnk.shape
)

# Use TMA to copy data
cute.copy(
    tma_atom_a,
    tAgA[(None, k_tile_idx)],
    tAsA[(None, producer_handle.index)],
    tma_bar_ptr=producer_handle.barrier,
    mcast_mask=a_full_mcast_mask,
)

Cluster Computing

Blackwell supports CTA clusters for better parallelism:

# Compute cluster layout
self.cluster_layout_vmnk = cute.tiled_divide(
    cute.make_layout((*self.cluster_shape_mn, 1)),
    (tiled_mma.thr_id.shape,),
)

# Launch with cluster
self.kernel(...).launch(
    grid=grid,
    block=[self.threads_per_cta, 1, 1],
    cluster=(*self.cluster_shape_mn, 1),
    stream=stream,
)

Running the Example

Ampere (SM80/SM86)

python examples/python/CuTeDSL/ampere/tensorop_gemm.py \
  --mnkl 8192,8192,8192,1 \
  --atom_layout_mnk 2,2,1 \
  --ab_dtype Float16 \
  --c_dtype Float16 \
  --acc_dtype Float32 \
  --a_major m --b_major n --c_major n

python examples/python/CuTeDSL/blackwell/dense_gemm.py \
  --ab_dtype Float16 \
  --c_dtype Float16 \
  --acc_dtype Float32 \
  --mma_tiler_mn 256,128 \
  --cluster_shape_mn 2,1 \
  --mnkl 8192,8192,8192,1 \
  --use_tma_store \
  --use_2cta_instrs

Performance Profiling

Use NVIDIA Nsight Compute to profile your kernel:

ncu python examples/python/CuTeDSL/ampere/tensorop_gemm.py \
  --mnkl 8192,8192,8192,1 \
  --atom_layout_mnk 2,2,1 \
  --ab_dtype Float16 \
  --c_dtype Float16 \
  --acc_dtype Float32 \
  --a_major m --b_major n --c_major n \
  --skip_ref_check --iterations 2

Key Concepts

Tile Shape Selection

The tile shape determines how the problem is divided:

• M dimension: Affects register usage and occupancy
• N dimension: Should align with MMA instruction shape
• K dimension: Larger values improve arithmetic intensity

Common configurations:

• Ampere: 128x128x32, 128x256x32, 256x128x32
• Blackwell: 128x128x64, 256x128x128

Memory Hierarchy

Understanding the memory hierarchy is crucial:

1. Global Memory (GMEM): Input matrices A, B, output C
2. Shared Memory (SMEM): Tile-level cache, swizzled to avoid bank conflicts
3. Registers (RMEM): Thread-level storage for MMA operands
4. Tensor Memory (TMEM): Blackwell-specific accumulator storage

Pipeline Stages

Multi-stage pipelining overlaps memory and compute:

• Stage 1: Load next tile from GMEM to SMEM
• Stage 2: Compute current tile MMA
• Stage 3: Write previous results to GMEM

More stages provide better overlap but increase SMEM usage.

Source Code

Ampere GEMM

Complete source: examples/python/CuTeDSL/ampere/tensorop_gemm.py

Blackwell GEMM

Complete source: examples/python/CuTeDSL/blackwell/dense_gemm.py