CuTe DSL

CuTe DSL is a Python-based Domain-Specific Language (DSL) for writing and dynamically compiling kernels, aiming for both high performance and high developer productivity.

It is part of CUTLASS, a collection of CUDA C++ templates and DSLs. Unlike cuBLAS or cuDNN, which offer ready-to-call kernels for common operations, the CUTLASS stack gives you tools for composably defining your own high-performance kernels.

Core abstractions#

• Layouts describe how data is organized in memory and across threads.
• Tensors combine data pointers or iterators with layout metadata.
• Atoms represent fundamental hardware operations, such as matrix multiply-accumulate (MMA) or memory copy.
• Tiled operations describe how atoms are applied across thread blocks and warps.

For the underlying details, see CuTe.

@cute.jit and @cute.kernel#

Launching a CuTe DSL kernel from Python takes two layers: Python calls a @cute.jit function, which launches a @cute.kernel function.

• @cute.jit declares a JIT-compiled function callable from Python or from other CuTe DSL functions.
• @cute.kernel defines a GPU kernel that can be launched from a @cute.jit function. Python code cannot call it directly.

Example: elementwise addition#

Below is a naive (unoptimized) kernel that adds two one-dimensional tensors elementwise. This is the "hello world" of GPU programming, going back to Ian Buck's Brook framework, which preceded and inspired CUDA. You can edit and run it on a B200 GPU with this Modal Notebook.

import cutlass.cute as cute
import torch

Tensor = cute.Tensor | torch.Tensor


@cute.kernel
def elem_add_kernel(a: cute.Tensor, b: cute.Tensor, out: cute.Tensor):
    block_x, _, _ = cute.arch.block_idx()
    block_dim_x, _, _ = cute.arch.block_dim()
    thread_x, _, _ = cute.arch.thread_idx()

    i = block_x * block_dim_x + thread_x

    if i < out.shape[0]:
        out[i] = a[i] + b[i]


@cute.jit
def elem_add(a: Tensor, b: Tensor, out: Tensor):
    n = out.shape[0]
    threads_per_block = 128
    blocks = (n + threads_per_block - 1) // threads_per_block

    elem_add_kernel(a, b, out).launch(
        grid=(blocks, 1, 1),
        block=(threads_per_block, 1, 1),
    )

elem_add_kernel is the kernel. Each thread computes one output element, with its global index i derived from the thread block index, the number of threads per block, and the thread's index within the block:

i = block_x * block_dim_x + thread_x

elem_add computes how many thread blocks are needed to cover the output tensor and launches the kernel on a one-dimensional thread block grid.

The example is for teaching, not speed. Still, it shows a good basic access pattern: adjacent threads read adjacent elements of a and b and write adjacent elements of out. That's what coalesced access to global memory needs; see memory coalescing.

Why it suits high-performance kernels#

Performance engineering is hard because kernels must map closely onto hardware: which threads handle which data, how memory is accessed, how work is tiled, and which hardware operations the generated code uses. Layout concerns like these are one reason CuTe DSL is useful. It lets you express these mappings explicitly while reusing much of the same kernel code across shapes and Streaming Multiprocessor architectures.

How Python competes with compiled languages#

Performance engineers from other domains may wonder how an interpreted language like Python can compete with compiled ones.

Key idea CuTe DSL kernels are compiled Just-In-Time (JIT). The Python source is converted to an abstract syntax tree (AST), traced with proxy arguments, and then compiled.

Only a subset of Python semantics is supported in JIT-compiled code.

At time of writing, in CUTLASS 4.x, compilation goes through Multi-Level Intermediate Representation (MLIR), then the PTX IR, then device-specific SASS, which is what executes.

In practice: FlashAttention-4#

The FlashAttention-4 kernels reach state-of-the-art performance directly from CuTe DSL. Modal's writeup of the open source code shows how they use pipelined warp specialization, Tensor Core operations, and Tensor Memory and Tensor Memory Accelerator operations.

For more, see NVIDIA's CuTe DSL documentation and CuTe DSL overview blog post.

Related#

• CuTe: the C++ library whose abstractions CuTe DSL exposes
• CUTLASS: the larger stack CuTe DSL belongs to
• Memory coalescing: the access pattern the example demonstrates
• Thread hierarchy: how threads, blocks and grids fit together