NVlabs/CuTe

Reference implementation and examples of the CuTe Layout representation and algebra.

What it solves

PyCuTe provides a pure-Python reference implementation of the hierarchical layout-and-tensor algebra used in NVIDIA's CUTLASS 3.x. It allows developers to learn the complex algebra, prototype new data transformations, and generate test vectors without needing a GPU or the C++ environment required by the original CuTe library.

How it works

At its core, PyCuTe defines a Layout as a function that maps coordinates to offsets using a Shape (the domain) and a Stride (the mapping logic). It implements a comprehensive set of algebraic operations—such as coalesce, composition, logical_divide (tiling), and logical_product (repetition)—to manipulate these layouts. It also supports F2 (XOR-swizzle) strides to handle shared-memory bank conflicts. A Tensor is then created by pairing a Layout with an Accessor (like a pointer) to read or write data.

Who it’s for

It is designed for developers and researchers working with GPU tensor computations, CUTLASS users who want to prototype layout transformations in Python, and those studying the mathematical foundations of hierarchical tensor layouts.

Highlights

  • Pure Python: No GPU or CUDA dependencies required for the core algebra.
  • Layout Algebra: Full implementation of operations like composition, complement, and nullspace for complex tiling and partitioning.
  • Visualization Tools: Built-in support for rendering layouts as ASCII tables, colored SVGs, or LaTeX/PDFs to visualize data movement and bank conflicts.
  • Reference Implementation: Serves as a faithful Python counterpart to the C++ CuTe library and the CuTe DSL.

Related

  • Project
    NVIDIA/cutlass

    CUTLASS 4.8.1 is NVIDIA’s open‑source C++ template library (plus a new Python CuTe DSL) for building ultra‑fast GEMM, convolution, and attention kernels on CUDA GPUs. It supports a wide range of precisions (FP64‑FP8, BF16, integer, binary), all recent NVIDIA architectures, and provides high‑level operator APIs and extensible primitives for AI researchers and framework developers.

  • Project
    rapidsai/cugraph

    cuGraph is an open‑source RAPIDS library that provides GPU‑accelerated graph algorithms with Python (NetworkX‑style) and low‑level C/C++ APIs, enabling fast, end‑to‑end graph analytics on the GPU.

  • Project
    facebookincubator/AITemplate

    A Python framework that transforms deep neural networks into optimized CUDA/HIP C++ code for high-performance inference serving on NVIDIA and AMD GPUs.

  • Project
    gau-nernst/learn-cuda

    A collection of CUDA kernel implementations and educational resources for PyTorch users to learn GPU programming and hardware optimization.

  • Project
    Dao-AILab/quack

    A collection of high-performance GPU kernels for deep learning operations like RMSNorm and Softmax, implemented using CuTe-DSL for modern NVIDIA GPUs.