GPU Glossary
/
Host Software·3 min read

CuTe

A header-only CUDA C++ template library inside CUTLASS for describing and manipulating tensors of data and threads through composable Layouts.

CuTe (CUDA Templates) is a header-only CUDA C++ library within CUTLASS for describing and manipulating tensors of data and threads.

It is built on C++ templates, C++'s implementation of parametric polymorphism (often called generics in other languages): a polymorphic function is written once but works on inputs of different types. Don't confuse CuTe with CuTe DSL, which exposes CuTe/CUTLASS through a Python Domain-Specific Language (DSL).

Layouts and Tensors#

CuTe's type system centers on two types:

  • A Layout describes a regular pattern of access to a tensor. It is built from Shape and Stride tuples, which describe memory extents and how to traverse them.
  • A Tensor combines a Layout with a pointer to memory.

Critically, Layouts are composable. They form a category with a rich algebra, so they offer both expressiveness and structure.

Correctness and performance from the type system#

CuTe encodes key program metadata (memory organization, strided accesses, tiling) in types. That lets the compiler check many aspects of correctness and preserve invariants while it optimizes. The result is very high-level metaprogramming of kernels without giving up performance; for example, one template can compile into highly optimized kernels for several Streaming Multiprocessor architectures.

Key idea Layouts are resolved at compile time, so memory accesses carry zero extra runtime overhead, overhead that could otherwise kill performance in memory-bound workloads.

See NVIDIA's CuTe documentation for more.

Example: matrix transpose#

This transpose kernel shows CuTe's core features: templating, shapes, layouts and tensors. It is based on the initial "naive" implementation from this Colfax International article. You can run it on an H100 via this Modal Notebook.

// one CuTe trick: transpose a row-major matrix just using Layouts
template <typename T>
__global__ void transpose_kernel(const T* __restrict__ d_S,
                                 T* __restrict__ d_D,
                                 int M, int N)
{
    // define the Shape of tiles worked on by thread blocks
    using b = Int<32>;
    auto block_shape = make_shape(b{}, b{});

    // define the Shape of input/output Tensors
    auto tensor_shape = make_shape(M, N);

    // define the Layout of the input and output Tensors in global memory
    auto gmemLayoutS  = make_layout(tensor_shape, GenRowMajor{}); // input:  row-major
    auto gmemLayoutDT = make_layout(tensor_shape, GenColMajor{}); // output: col-major

    // construct the Tensors
    auto tensor_S  = make_tensor(make_gmem_ptr(d_S), gmemLayoutS);
    auto tensor_DT = make_tensor(make_gmem_ptr(d_D), gmemLayoutDT);

    // define a tile-ing of the Tensors (as a "Tensor of Tensors")
    auto tiled_tensor_S  = tiled_divide(tensor_S,  block_shape);
    auto tiled_tensor_DT = tiled_divide(tensor_DT, block_shape);

    // pull out the tiles this thread block will be working on
    auto tile_S  = tiled_tensor_S (make_coord(_, _), blockIdx.x, blockIdx.y);
    auto tile_DT = tiled_tensor_DT(make_coord(_, _), blockIdx.x, blockIdx.y);

    // create a Layout for threads in the thread block
    auto thr_layout = make_layout(
        make_shape(Int<8>{}, Int<32>{}),
        GenRowMajor{}
    );

    // pull out the tile this thread will work on
    auto thr_tile_S  = local_partition(tile_S,  thr_layout, threadIdx.x);
    auto thr_tile_DT = local_partition(tile_DT, thr_layout, threadIdx.x);

    // define a "Tensor" in register memory
    auto rmem = make_tensor_like<T>(thr_tile_S);

    // copy tile into registers
    copy(thr_tile_S, rmem);
    // copy tile out of registers as though it were column-major
    copy(rmem, thr_tile_DT);
}

The transpose happens purely through layouts: the input is read with a row-major Layout and written back with a column-major one.