ThunderKittens Documentation

ThunderKittens is a framework designed for writing high-performance deep learning kernels in CUDA for NVIDIA GPUs. It is built around the principle of manipulating small tiles of data (minimum 16x16) to maximize hardware utilization.

Key features include:

• Tensor Core acceleration: Support for asynchronous WGMMA (H100) and TCGEN05 (B200) calls.
• Memory Management: Optimized shared memory usage to avoid bank conflicts, and support for Distributed Shared Memory.
• Asynchronous Operations: Uses TMA (Tensor Memory Accelerator) for asynchronous loads/stores and hides latencies.
• Work Overlapping: Provides a Load-Store-Compute-Finish template to overlap I/O with computation.
• Multi-GPU Support: Enables data transfer over NVLink and NVSwitch acceleration.

Note on Hardware Support:

• Primary Support: Hopper and Blackwell GPUs.
• Legacy Support: Ampere GPUs are supported but no longer actively maintained.
• AMD Support: Use HipKittens instead.

ThunderKittens operates on two primary data type categories across different NVIDIA GPU scopes:

Data Types

• Register Objects: Exist at the Warp level. Their contents are split among the threads of the warp.
  • Register Tiles (kittens::rt): Parameterized by layout, type, and size. Example: kittens::rt_bf<32,16> is a 32x16 bfloat16 register tile (row-layout is default).
  • Register Vectors: Associated with tiles. Flavors include naive (for compute-heavy tasks like layernorm), column (to reduce/map across tile rows), and row (to reduce/map across tile columns).
• Shared Objects: Exist at the Block level and reside in shared memory. Example: kittens::st_hf<32, 64>.

Scopes

• Warp (Default): Most programming happens here. A single warp performs the work. Use kittens::warpid() to get the current warp ID.
• Warpgroup: A group of 4 warps. Required for H100 asynchronous matrix multiply-accumulate (WGMMA) instructions. Access via kittens::group<4> or kittens::warpgroup.
• Block: A group of warps (often 8) that shares shared memory.
• Grid: A collection of blocks.

To use the provided kernels (e.g., kernels/gemm/bf16_h100) in your PyTorch code, follow these steps:

1. Prerequisites: Ensure your Python environment has PyTorch 2.8+ and PyBind11 installed, with CUDA compatibility.
2. Environment Setup: Set the following environment variables to speed up the build process:

   export PYTHON_VERSION=$(python3 -c "import sysconfig; print(sysconfig.get_config_var('LDVERSION'))")
   export PYTHON_INCLUDES=$(python3 -c "import sysconfig; print('-I', sysconfig.get_path('include'), sep='')")
   export PYBIND_INCLUDES=$(python3 -m pybind11 --includes)
   export PYTORCH_INCLUDES=$(python3 -c "from torch.utils.cpp_extension import include_paths; print(' '.join(['-I' + p for p in include_paths()]))")
   export PYTHON_LIBDIR=$(python3 -c "import sysconfig; print('-L', sysconfig.get_config_var('LIBDIR'), sep='')")
   export PYTORCH_LIBDIR=$(python3 -c "from torch.utils.cpp_extension import library_paths; print(' '.join(['-L' + p for p in library_paths()]))")

3. Build:
   • cd into the specific kernel directory.
   • Edit the Makefile if necessary.
   • Run make.
4. Run: Execute make run to run correctness tests and benchmarks.

# Setup environment
export PYTHON_VERSION=$(python3 -c "import sysconfig; print(sysconfig.get_config_var('LDVERSION'))")
export PYTHON_INCLUDES=$(python3 -c "import sysconfig; print('-I', sysconfig.get_path('include'), sep='')")
export PYBIND_INCLUDES=$(python3 -m pybind11 --includes)
export PYTORCH_INCLUDES=$(python3 -c "from torch.utils.cpp_extension import include_paths; print(' '.join(['-I' + p for p in include_paths()]))")
export PYTHON_LIBDIR=$(python3 -c "import sysconfig; print('-L', sysconfig.get_config_var('LIBDIR'), sep='')")
export PYTORCH_LIBDIR=$(python3 -c "from torch.utils.cpp_extension import library_paths; print(' '.join(['-L' + p for p in library_paths()]))")

# Build and run
cd kernels/gemm/bf16_h100
make
make run

ThunderKittens is a header-only library. There is no installation process required; simply clone the repository and include kittens.cuh in your CUDA projects.

Build Requirements

1. CUDA Environment

Requires CUDA 12.8+. Ensure your environment variables are set correctly:

export CUDA_HOME=/usr/local/cuda-<YOUR-CUDA-VERSION> # ex. cuda-12.8
export PATH=${CUDA_HOME}/bin:${PATH} 
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH

2. C++20 Compiler

ThunderKittens uses C++20 concepts. Ensure you have a modern compiler (e.g., gcc-11 or clang-11). On Ubuntu, you can update using:

sudo apt update
sudo apt install gcc-11 g++-11

sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11

sudo apt update
sudo apt install clang-11

3. Troubleshooting libc10.so errors

If you encounter libc10.so errors, you may need to add the PyTorch library path to your LD_LIBRARY_PATH. Find the path using Python and export it:

# Find the path
python -c "import torch; print(torch.file)"

# Export the path (replace <PRINTED_PATH> with the output from the command above)
export LD_LIBRARY_PATH=<PRINTED_PATH>/lib:$LD_LIBRARY_PATH