### Setup Python Environment Source: https://github.com/hazyresearch/thunderkittens/blob/main/demos/based/README.md Installs necessary Python packages for the project, including PyTorch, Transformers, and Flash Attention. ```bash conda create -n dev python=3.11 pip3 install torch torchvision torchaudio pip install transformers pip install einops pip install hydra-core pip install flash-attn ``` -------------------------------- ### PyBind11 Kernel Integration Setup Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Shell commands to install PyBind11 and PyTorch, set necessary environment variables for Python and PyTorch includes/libraries, and build the kernel. ```bash # Install dependencies pip install pybind11 torch # Navigate to kernel directory with PyBind11 support cd kernels/gemm/bf16_h100 # Set environment variables 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 make # Test from Python python -c "import torch; from kernel import my_function; print(my_function)" ``` -------------------------------- ### Example: Producer/Consumer Register Configuration Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Example demonstrating conditional configuration of producer and consumer registers based on group ID. ```cuda __global__ void my_kernel(...) { if (kittens::warpgroup::groupid() == 0) { kittens::warpgroup::producer_registers(); } else { kittens::warpgroup::consumer_registers<4>(); } } ``` -------------------------------- ### Install Flash Linear Attention Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Install the flash-linear-attention package. This is a prerequisite for using the Based architecture demos. ```bash pip install -U git+https://github.com/sustcsonglin/flash-linear-attention ``` -------------------------------- ### Install Benchmarking Baselines Source: https://github.com/hazyresearch/thunderkittens/blob/main/demos/based/README.md Installs optional baseline kernels for benchmarking, including Flash Linear Attention and Mamba kernels. ```bash git clone https://github.com/sustcsonglin/flash-linear-attention.git pip install -U git+https://github.com/sustcsonglin/flash-linear-attention pip install mamba_ssm ``` -------------------------------- ### Clone CUTLASS Repository Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md Initial setup involves cloning the CUTLASS repository and navigating into its build directory. ```bash git clone https://github.com/NVIDIA/cutlass.git cd cutlass mkdir build cd build ``` -------------------------------- ### Customize Kernel Build Options Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Example Makefile snippet showing how to customize tile sizes, data types, and enable profiling before building. ```makefile # Example: Change tile sizes TILE_M ?= 128 TILE_N ?= 128 TILE_K ?= 32 # Example: Change data types DATA_TYPE ?= bf16 # Example: Enable profiling PROFILE ?= 1 ``` -------------------------------- ### Set Up CUDA Environment Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Configure your environment variables to point to your CUDA installation. Ensure CUDA 12.8 or newer is used. ```bash export CUDA_HOME=/usr/local/cuda- # ex. cuda-12.8 export PATH=${CUDA_HOME}/bin:${PATH} export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH ``` -------------------------------- ### CUDA Example for Global to Register Load Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/03-memory-operations.md This example demonstrates how to collaboratively load a tile from global memory into a register tile using the `kittens::group<4>::memory::load` function. Ensure all warps in the group are synchronized. ```cuda kittens::rt_bf<32, 64> C; kittens::group<4>::memory::load<2>(C, globals.C, {0, 0}); ``` -------------------------------- ### Install ThunderKittens Kernels Source: https://github.com/hazyresearch/thunderkittens/blob/main/demos/based/README.md Installs the ThunderKittens kernels after selecting the 'based' configuration in the project's config file. ```bash python setup.py install ``` -------------------------------- ### Check GPU Availability and CUDA Setup Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Verify that your NVIDIA GPU is recognized by the system and that the CUDA toolkit is functioning correctly. ```bash # Check GPU nvidia-smi # Test CUDA /usr/local/cuda/bin/deviceQuery # Should show your GPU ``` -------------------------------- ### Login to Hugging Face CLI Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Authenticate with the Hugging Face CLI to download models. Ensure you have the CLI installed. ```bash huggingface-cli login ``` -------------------------------- ### Shared Memory Allocation Examples Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Demonstrates declaring shared memory tiles statically or allocating them dynamically via kernel configuration. ```cuda extern __shared__ char shared_memory[]; __shared__ kittens::st_bf<32, 64> tile1; __shared__ kittens::st_bf<32, 64> tile2; ``` ```cuda __global__ void kernel() { __shared__ kittens::st_bf<32, 64> S; // ... } // Launch with dynamic shared memory kernel<<>>(args); ``` -------------------------------- ### FP8 Conversion Example Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/06-advanced-features.md Demonstrates converting a float to FP8 format and then converting FP8 back to float within a kernel. Requires specific conversion utilities. ```cuda float f = 3.14f; fp8e4m3 fp8 = kittens::base_types::convertor::convert(f); // In compute kittens::rt_fp8e4m3<32, 64> A; kittens::rt_fl<32, 64> result; kittens::warp::convert(result, A); // Convert to float ``` -------------------------------- ### Verify ThunderKittens Installation with a CUDA Kernel Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Create a simple CUDA file to test compilation and runtime of ThunderKittens. This kernel initializes a tile and prints a success message. ```cuda #include "kittens.cuh" #include __global__ void test_kernel() { kittens::rt_bf<32, 64> tile; kittens::warp::zero(tile); if (kittens::laneid() == 0) { printf("ThunderKittens test passed!\n"); } } int main() { test_kernel<<<1, 32>>>(); cudaDeviceSynchronize(); return 0; } ``` -------------------------------- ### Producer-Consumer Synchronization with Semaphores Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Example of using semaphores to coordinate producer and consumer threads for data loading and processing. ```cuda __shared__ cuda::counting_semaphore inputs_ready, results_ready; // Producer thread if (is_producer) { // ... load data ... inputs_ready.release(); // Signal ready } // Consumer thread if (is_consumer) { inputs_ready.acquire(); // Wait for data // ... compute ... results_ready.release(); // Signal done } ``` -------------------------------- ### Initiate and Wait for Async MMA Operations Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/04-compute-operations.md This example demonstrates initiating multiple asynchronous MMA operations and then waiting for a specific number of them to complete using `mma_async_wait`. Ensure the correct count is provided to `mma_async_wait` to avoid deadlocks or race conditions. ```cuda // Initiate async MMAs kittens::warpgroup::mma_async(C, A, B); kittens::warpgroup::mma_async(C, A2, B2); // Do other work // ... // Wait for all to complete kittens::warpgroup::mma_async_wait<2>(); // Wait for 2 pending ops ``` -------------------------------- ### Update GCC and G++ to Version 11 Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Install and configure GCC and G++ version 11 for C++20 compatibility. This ensures your compiler meets ThunderKittens' requirements. ```bash 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 ``` -------------------------------- ### Perform Reduction Operation Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Example of performing a reduction operation, specifically summing elements within a tile. ```cuda // Reduction: result = sum(tile) float result = kittens::warp::sum(tile); ``` -------------------------------- ### Check CUDA Version Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Verify your installed CUDA version using the nvcc command. Ensure it is 12.8 or newer. ```bash nvcc --version ``` ```text # Output should show 12.8 or higher: # nvcc: NVIDIA (R) Cuda compiler driver # Copyright (c) 2005-2024 NVIDIA Corporation # Built on Tue_Feb_27_16:32:37_PST_2024 # Cuda compilation tools, release 12.8, V12.8.89 ``` -------------------------------- ### Set CUDA Home and Path Environment Variables Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md If CUDA is not in the default path, set CUDA_HOME, PATH, and LD_LIBRARY_PATH environment variables to point to your CUDA installation. ```bash # Identify your CUDA installation ls /usr/local/ | grep cuda # Set environment variables export CUDA_HOME=/usr/local/cuda-12.8 export PATH=${CUDA_HOME}/bin:${PATH} export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH # Verify nvcc --version ``` -------------------------------- ### Check CUDA Compiler Version Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Verify that the installed NVIDIA CUDA compiler (nvcc) meets the minimum version requirement. ```bash # Check CUDA nvcc --version # Should be 12.8+ ``` -------------------------------- ### Profile BF16 GEMM Kernels for Various Shapes Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md Run the CUTLASS profiler to find the best kernel configuration for BF16 GEMM operations across different matrix dimensions (N). This example iterates through several sizes of N, M, and K. ```bash # BF16_BF16_FP32_void_BF16 GEMM for N in 1024 2048 4096 8192 16384; do ./tools/profiler/cutlass_profiler \ --operation=gemm \ --m=${N} \ --n=${N} \ --k=${N} \ --output=bf16_${N} \ --kernels="cutlass3x_sm100_tensorop_gemm_bf16_bf16_f32_void_bf16" \ --warmup-iterations=500 \ --profiling-iterations=100 \ --verification-enabled=false \ --enable-best-kernel-for-fixed-shape=true \ --use_pdl=1 \ --dist="uniform,min:-1,max:1,scale:-1" done ``` -------------------------------- ### CMake Configuration for Multiple GEMM Kernel Types Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md Configure CMake to build multiple types of GEMM kernels simultaneously. This example includes BF16, MXFP8, and NVFP4 GEMM kernels. ```bash # All at once cmake .. \ -DCUTLASS_NVCC_ARCHS=100a \ -DCUTLASS_UNITY_BUILD_ENABLED=OFF \ -DCUTLASS_ENABLE_TESTS=OFF \ -DCUTLASS_ENABLE_EXAMPLES=OFF \ -DCUTLASS_LIBRARY_KERNELS="cutlass3x_sm100_tensorop_gemm_bf16_bf16_f32_void_bf16*_tnt_*,cutlass3x_sm100_bstensorop_gemm_ue8m0xe4m3_ue8m0xe4m3_f32_void_bf16*_tnt_*,cutlass3x_sm100_bstensorop_gemm_ue4m3xe2m1_ue4m3xe2m1_f32_void_f32*_tnt_*" ``` -------------------------------- ### Example: Adding a Row Vector in CUDA Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/04-compute-operations.md Demonstrates how to use the `add_row` function to add a row vector to a CUDA tensor. Ensure the row vector is correctly defined and compatible with the tensor's layout. ```cuda kittens::rt_bf<32, 64> A, B; auto row_vec = typename kittens::rt_bf<32, 64>::row_vec; kittens::warp::add_row(B, A, row_vec); ``` -------------------------------- ### Complex Arithmetic Example Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/06-advanced-features.md Demonstrates the usage of complex multiply-accumulate (mma) and magnitude extraction (abs) operations with complex register tiles (crt_fl). ```cuda kittens::crt_fl<32, 64> A, B, C; // Complex multiply-accumulate kittens::warpgroup::mma(C, A, B); // C += A * B (complex) // Get magnitude kittens::rt_fl<32, 64> magnitude; kittens::warp::abs(magnitude, C); ``` -------------------------------- ### CMake Configuration for BF16 GEMM Kernels Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md Configure CMake to build specific GEMM kernels for BF16 data types with Tensor Core operations. This example targets a specific kernel pattern. ```bash # BF16_BF16_FP32_void_BF16 GEMM cmake .. \ -DCUTLASS_NVCC_ARCHS=100a \ -DCUTLASS_UNITY_BUILD_ENABLED=OFF \ -DCUTLASS_ENABLE_TESTS=OFF \ -DCUTLASS_ENABLE_EXAMPLES=OFF \ -DCUTLASS_LIBRARY_KERNELS=cutlass3x_sm100_tensorop_gemm_bf16_bf16_f32_void_bf16*_tnt_* ``` -------------------------------- ### CMake Configuration for MXFP8 GEMM Kernels Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md Configure CMake to build specific GEMM kernels for MXFP8 data types with Tensor Core operations. This example targets a specific kernel pattern. ```bash # MXFP8_MXFP8_FP32_void_BF16 GEMM cmake .. \ -DCUTLASS_NVCC_ARCHS=100a \ -DCUTLASS_UNITY_BUILD_ENABLED=OFF \ -DCUTLASS_ENABLE_TESTS=OFF \ -DCUTLASS_ENABLE_EXAMPLES=OFF \ -DCUTLASS_LIBRARY_KERNELS=cutlass3x_sm100_bstensorop_gemm_ue8m0xe4m3_ue8m0xe4m3_f32_void_bf16*_tnt_* ``` -------------------------------- ### Build and Run a ThunderKittens Kernel Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md After setting up the environment and navigating to the kernel directory, use 'make' to build and 'make run' to execute the kernel. Ensure the Makefile is configured for your specific needs. ```bash make ``` ```bash make run ``` -------------------------------- ### Build and Run Unit Tests Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Builds the project with default settings and runs the unit tests. After execution, it shows how to list and view details of failed tests. ```bash cd tests # Build with default settings (SM90) make -j # Or specify architecture make ARCH=SM90 -j make ARCH=SM100 -j # Run tests mkdir outputs ./unit_tests printout # View failed test details ls outputs/ cat outputs/failed_test_name.txt ``` -------------------------------- ### Python setuptools for CUDA Extension Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Configures Python setuptools to build a CUDA extension using nvcc, integrating with PyBind11 and PyTorch. ```python from setuptools import setup, Extension from setuptools.command.build_ext import build_ext import pybind11 import torch class CudaExtension(Extension): def __init__(self, name, sources, **kwargs): super().__init__(name, sources, **kwargs) class CudaBuildExt(build_ext): def build_extension(self, ext): import subprocess cuda_flags = [ '-DKITTENS_SM90', '-std=c++20', '-O3', f'-I{pybind11.get_include()}', f'-I{torch.utils.cpp_extension.include_paths()[0]}', ] compile_cmd = [ 'nvcc', *cuda_flags, '-c', ext.sources[0], '-o', self.get_ext_filename(ext.name) ] subprocess.check_call(compile_cmd) setup( name='my_kernel', ext_modules=[CudaExtension('my_kernel', ['kernel.cu'])], cmdclass={'build_ext': CudaBuildExt}, ) ``` -------------------------------- ### Launch Nsight Systems GUI Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Open the Nsight Systems graphical user interface to analyze profiling traces. This tool provides a visual representation of application performance. ```bash # nsight-systems GUI nsight-systems ``` -------------------------------- ### Run Custom Generation Demo Source: https://github.com/hazyresearch/thunderkittens/blob/main/demos/based/README.md Launches the generation script to run models with user-provided prompts. ```bash python generate_based.py ``` -------------------------------- ### Tile Initialization Utilities Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/API-SUMMARY.md Provides utility functions to initialize tiles with zero values or fill them with a specified scalar value. Use these for setting up tiles before computation. ```cuda // Tile utilities void zero(T &tile); void fill(T &tile, typename T::T value); ``` -------------------------------- ### Build and Test a CUDA Kernel Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Commands to navigate to a kernel directory, view its Makefile, build, test, benchmark, and clean the kernel. ```bash # Navigate to kernel directory cd kernels/gemm/bf16_h100 # View Makefile configuration cat Makefile # Build make # Run tests make test # Run benchmarks make benchmark # Clean make clean ``` -------------------------------- ### Compile and Run CUDA Test File Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Compile the test CUDA file using nvcc, specifying necessary flags and include paths. Then, execute the compiled binary. ```bash nvcc -DKITTENS_SM90 -std=c++20 -I/path/to/thunderkittens/include \ test_tk.cu -o test_tk ./test_tk ``` -------------------------------- ### Get Warp Group ID within Block Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md In group-scope operations, returns which warp group (0, 1, ...) the thread belongs to. ```cuda __device__ static inline int groupid(); ``` -------------------------------- ### Get Lane ID within Warp Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Returns the ID of the current thread within its warp (0-31) or group (0 to GROUP_THREADS-1 in group scope). ```cuda __device__ static inline int laneid(); ``` -------------------------------- ### Creating Register Tiles Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/API-SUMMARY.md Demonstrates the creation of register tiles in row-major and column-major layouts. Use `kittens::col_l` for column-major. ```cuda kittens::rt_bf<32, 64> A; kittens::rt_bf<32, 64, kittens::col_l> B; ``` -------------------------------- ### Store Data to Global Memory Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Examples for storing data from registers or shared memory to global memory. Use `store_async` for shared memory to overlap memory transfers. ```cuda // Register to global kittens::group<4>::memory::store(global_layout, reg_tile, tile_coord); // Shared to global (async) kittens::group<4>::memory::store_async(global_layout, shared_tile, tile_coord, barrier); ``` -------------------------------- ### Load Data from Global Memory Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Examples for loading data into registers or shared memory from global memory. Use `load_async` for shared memory to overlap memory transfers with computation. ```cuda // Register from global kittens::group<4>::memory::load(reg_tile, global_layout, tile_coord); // Shared from global (async) kittens::group<4>::memory::load_async(shared_tile, global_layout, tile_coord, barrier); // Register from shared kittens::group<4>::memory::load(reg_tile, shared_tile); ``` -------------------------------- ### Basic Compilation Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/API-SUMMARY.md Compiles a CUDA kernel with specific architecture and optimization flags. ```bash # Basic compilation nvcc -DKITTENS_SM90 -std=c++20 -O3 kernel.cu -o kernel ``` -------------------------------- ### Initial CMake Configuration for Kernel Discovery Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md Configure CMake to discover available kernels. This step helps in listing all potential kernels before selecting specific ones for compilation. ```bash cmake .. \ -DCUTLASS_NVCC_ARCHS=100a \ -DCUTLASS_UNITY_BUILD_ENABLED=OFF \ -DCUTLASS_ENABLE_TESTS=OFF \ -DCUTLASS_ENABLE_EXAMPLES=OFF \ -DCUTLASS_LIBRARY_KERNELS=cutlass3x_sm100_* ``` -------------------------------- ### Register Pressure Reduction via Sequential Computation Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/06-advanced-features.md Provides an example of reducing register pressure in kernels with tight register budgets by performing computations on smaller tiles sequentially and reusing registers. ```cuda // Do computation on smaller tiles sequentially kittens::rt_bf<16, 32> A, B, C; kittens::warp::mma(C, A, B); kittens::warp::zero(A); // Reuse registers kittens::warp::load(A, other_src); kittens::warp::mma(C, A, B); // Second compute ``` -------------------------------- ### Run LoLCATS-Llama 3.1 8B Demo Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Execute the demo script for the LoLCATS-Llama 3.1 8B model with TK attention. Navigate to the demo directory and run the provided shell script. ```bash cd lolcats_demo/ bash demo_8b.sh ``` -------------------------------- ### Run Document Information Extraction Demo Source: https://github.com/hazyresearch/thunderkittens/blob/main/demos/based/README.md Executes the document information extraction demo script for different models (based, attn, mamba) to evaluate in-context learning ability and task completion time. ```bash cd ThunderKittens/demos/based_demo/ python document_ie_based.py -- model_name based python document_ie_based.py -- model_name attn python document_ie_based.py -- model_name mamba ``` -------------------------------- ### Benchmark Standalone Kernel Source: https://github.com/hazyresearch/thunderkittens/blob/main/demos/based/README.md Runs the script to benchmark the standalone kernel performance, aiming to reproduce the plot shown in the README. ```bash python benchmark_kernel.py ``` -------------------------------- ### Shared Tile (st) Subtile Access (Rectangular Region) Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/02-types-reference.md Provides a device function to get a view into a rectangular sub-region of a Shared Tile. Requires specifying subtile dimensions and a row/column index. ```cuda // Get a view into a rectangular region template __device__ inline st_subtile subtile(int2 rowcol); ``` -------------------------------- ### Profile-Guided Compilation Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Compiles CUDA code with profiling enabled, useful for performance analysis. ```bash nvcc -DKITTENS_SM90 -std=c++20 -pg kernel.cu -o kernel ``` -------------------------------- ### Handle libc10 Issues Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md If you encounter 'libc10.so' errors, you may need to set the LD_LIBRARY_PATH to include the PyTorch library directory. ```bash # Get PyTorch library path python -c "import torch; print(torch.file)" # Set library path (example) export LD_LIBRARY_PATH=/path/to/pytorch/lib:$LD_LIBRARY_PATH ``` -------------------------------- ### Get Warp ID within Block Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Returns which warp (0, 1, 2, ...) the current thread belongs to. Warp 0 = threads 0-31, Warp 1 = threads 32-63, etc. ```cuda __device__ static inline int warpid(); ``` -------------------------------- ### Shared Tile (st) Subtile Access (Column Subdivision) Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/02-types-reference.md Provides a device function to get a true subtile, which is a column subdivision of a Shared Tile. Requires specifying the subtile column count and an index. ```cuda // Get a true subtile (column subdivision only) template __device__ inline st& subtile(int idx); ``` -------------------------------- ### Run End-to-End Pretrained Model Benchmarking Source: https://github.com/hazyresearch/thunderkittens/blob/main/demos/based/README.md Executes the benchmarking script to compare different attention approaches on end-to-end pretrained models and generate a comparison plot. ```bash python benchmark/benchmark.py ``` -------------------------------- ### Get Type Name and Check Tile Properties Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/06-advanced-features.md Use `kittens::get_rt_type_name` to retrieve the runtime type name for printing. Static assertions can verify tile properties like the number of elements and elements per thread. ```cuda // Get type name for printing constexpr const char* name = kittens::get_rt_type_name(); // Check tile properties static_assert(kittens::rt_bf<32, 64>::num_elements == 2048); static_assert(kittens::rt_bf<32, 64>::elements_per_thread == 16); ``` -------------------------------- ### Run NVIDIA System Profiler Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Use 'nsys profile' to profile the execution of your kernel. This command generates a trace that can be analyzed with Nsight Systems. ```bash # nsys (system profiler) nsys profile ./kernel ``` -------------------------------- ### Run Specific Unit Test Sections Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Allows running specific functionalities or all tests within the suite. Useful for targeted testing or quick checks. ```bash # Test specific functionality make TEST_WARP_MEMORY -j make TEST_WARP_MEMORY_VEC -j make TEST_WARP_COMPUTE -j # Run all tests make TEST_ALL -j ``` -------------------------------- ### Run Llama 3 8B Demo Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Execute the demo script for the Llama 3 8B model using TK GQA attention. Navigate to the demo directory and run the provided shell script. ```bash cd llama_demo/ bash demo_8b.sh ``` -------------------------------- ### CMake Configuration for CUDA Project Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Sets up a CMake project to build a CUDA executable, including specifying C++ standard, CUDA architectures, and include paths. ```cmake cmake_minimum_required(VERSION 3.18) project(MyKernel CUDA) set(CUDA_STANDARD 20) set(CUDA_ARCHITECTURES 90) # ThunderKittens include path set(THUNDERKITTENS_INCLUDE /path/to/thunderkittens/include) add_executable(my_kernel kernel.cu) target_include_directories(my_kernel PRIVATE ${THUNDERKITTENS_INCLUDE}) # Add architecture macro target_compile_options(my_kernel PRIVATE -DKITTENS_SM90) # CUDA-specific settings set_target_properties(my_kernel PROPERTIES CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON ) ``` -------------------------------- ### Asynchronous GEMM Pipeline Implementation Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/06-advanced-features.md An example of an asynchronous GEMM pipeline within a CUDA kernel, demonstrating how to overlap memory loads, computations, and stores using TMA (Tensor Memory Access) and barriers. This pattern aims to hide memory latency by processing data in stages. ```cuda struct PipelineState { kittens::st_bf<32, 64> A_stage1, A_stage2; kittens::st_bf<64, 64> B; kittens::rt_fl<32, 64> C; }; __global__ void gemm_pipeline(const float *A, const float *B, float *C, int M, int N, int K) { // ... setup ... PipelineState state; cuda::barrier load_barrier, compute_barrier; // Prefetch first tile kittens::group<4>::tma::load_async(state.A_stage1, globals.A, {0, 0}, load_barrier); for (int k = 1; k < K_tiles; k++) { // Load next tile kittens::group<4>::tma::load_async(state.A_stage2, globals.A, {k, 0}, load_barrier); // Compute with current tile (while next loads) cuda::arrive_and_wait(load_barrier); kittens::group<4>::memory::load(state.C, state.A_stage1); kittens::warpgroup::mma(state.C, state.A_stage1, state.B); // Swap stages std::swap(state.A_stage1, state.A_stage2); } } ``` -------------------------------- ### Simple Matrix Multiplication Kernel in Cuda Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md This Cuda kernel implements matrix multiplication using ThunderKittens primitives. It defines layouts for data and computation, and includes producer and consumer logic for asynchronous data loading and computation. This example is designed for high performance on GPUs like the H100. ```Cuda #include "kittens.cuh" #include "prototype.cuh" using namespace kittens; using namespace kittens::prototype; using namespace kittens::prototype::lcf; template struct matmul_layout { using base_tile = st_bf<64, 64>; using global_layout = gl; struct globals { global_layout A, B, C; }; struct input_block { base_tile a[M_BLOCK], b[N_BLOCK]; }; struct finish_block { base_tile c[M_BLOCK][N_BLOCK]; }; struct common_state { int2 coord; }; struct consumer_state { rt_fl<16, N_BLOCK*base_tile::cols> accum; }; }; template struct matmul_template { static constexpr int M_BLOCK = _M_BLOCK, N_BLOCK = _N_BLOCK, SUPER_M = _SUPER_M; using layout = matmul_layout; using wide_tile = st_bf<64, 64*N_BLOCK>; static constexpr int NUM_CONSUMER_WARPS=M_BLOCK*4, INPUT_PIPE_STAGES=4, PRODUCER_BARRIER_ARRIVALS=1; // Helper functions template __host__ static inline dim3 grid(int M, int N, int K) { return dim3(PERISISTENT_GRID ? 132 : M*N/(M_BLOCK*N_BLOCK*layout::base_tile::num_elements)); } // ThunderKittens template functions __device__ static inline void common_setup(common_setup_args args) { int Rblocks = args.globals.C.rows() / (M_BLOCK*64), Cblocks = args.globals.C.cols() / (N_BLOCK*64); int super_rows = (Rblocks/SUPER_M)*SUPER_M, final_rows = Rblocks - super_rows, super_repeat = SUPER_M*Cblocks; int task_id = args.task_iter*gridDim.x + blockIdx.x; if (task_id < super_rows * Cblocks) args.common.coord = { SUPER_M*(task_id/super_repeat) + task_id%SUPER_M, (task_id%super_repeat)/SUPER_M }; else if (task_id < Rblocks*Cblocks) { int remainder_id = task_id - super_rows*Cblocks; args.common.coord = { super_rows + (remainder_id%final_rows), remainder_id/final_rows }; } else { // Id is too high, no more work to do args.num_iters = -1; return; } args.num_iters = args.globals.A.cols()/64; int id = warpgroup::groupid() == NUM_CONSUMER_WARPS/4 ? 0 : warpgroup::groupid(); // producer sets as 0 args.common.coord = { args.common.coord.x*M_BLOCK + id, args.common.coord.y*N_BLOCK }; } struct producer { __device__ static void setup(producer_setup_args args) { warpgroup::decrease_registers<40>(); // decrease registers for producers } __device__ static void load(producer_load_args args) { if (warpgroup::laneid() == 0) { tma::expect(args.inputs_arrived, args.input); for(int i = 0; i < M_BLOCK; i++) tma::load_async(args.input.a[i], args.globals.A, {args.common.coord.x+i, args.iter}, args.inputs_arrived); for(int i = 0; i < N_BLOCK; i++) tma::load_async(args.input.b[i], args.globals.B, {args.iter, args.common.coord.y+i}, args.inputs_arrived); } } }; struct consumer { __device__ static void setup(consumer_setup_args args) { warpgroup::increase_registers<232>(); // increase registers for consumers kittens::warp::zero(args.state.accum); } __device__ static void compute(consumer_compute_args args) { warpgroup::mma_AB( args.state.accum, // dest registers args.input.a[warpgroup::groupid()], // A matrix reinterpret_cast(args.input.b) // B matrix ); warpgroup::mma_async_wait(); if (warp::laneid() == 0) arrive(args.inputs_finished); } __device__ static void finish(consumer_finish_args args) { warpgroup::store(reinterpret_cast(args.finish.c[warpgroup::groupid()]), args.state.accum); warpgroup::sync(warpgroup::groupid()+4); if (warpgroup::laneid() == 0) for(int i = 0; i < N_BLOCK; i++) { tma::store_async(args.globals.C, args.finish.c[warpgroup::groupid()][i], {args.common.coord.x, args.common.coord.y+i}); tma::store_async_read_wait(); // wait that store is finished before reusing finish memory } kittens::warp::zero(args.state.accum); if (warp::laneid() == 0) arrive(args.finish_finished); } }; }; ``` -------------------------------- ### Conversions and Layouts Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/API-SUMMARY.md Handles type and layout conversions for tiles, and transposing. ```APIDOC ## Conversions and Layouts ### Description Handles type and layout conversions for tiles, and transposing operations. ### Namespace `kittens::warp` ### Methods - `void convert(D &dst, const A &src)`: Converts the type of the source tile to the destination tile. - `void layout_convert(D &dst, const A &src)`: Converts the layout of the source tile to the destination tile. - `void transpose_sep(D &dst, const A &src)`: Transposes the source tile into a separate destination buffer. ### Parameters - `dst`: Destination tile or buffer. - `src`: Source tile. ``` -------------------------------- ### Creating Shared Tiles Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/API-SUMMARY.md Shows how to declare shared memory tiles, with options for default swizzle or explicit swizzle configuration. The second template argument specifies the swizzle size. ```cuda __shared__ kittens::st_bf<32, 64> S; __shared__ kittens::st_bf<32, 64, true, 64> S_swizzle64; ``` -------------------------------- ### Compile and Run Operator Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/parallel/README.md Navigate to the desired operator directory and use 'make run' to compile and execute the kernel with an 8-GPU torchrun configuration. Alternatively, use 'make' to compile only. ```bash cd # navigate to the desired operator directory make run # compiles and executes with an 8-GPU torchrun configuration ``` ```bash make ``` -------------------------------- ### kittens::sync() Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Combines the functionality of `kittens::arrive()` and `kittens::wait()` for cluster synchronization. ```APIDOC ## kittens::sync() ### Description Combined arrive and wait for cluster synchronization. Signals arrival at a cluster barrier and then waits for all threads in the cluster to arrive. ### Method `__device__ static inline void sync()` ``` -------------------------------- ### Pipeline Pattern in CUDA Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Demonstrates a three-stage pipeline pattern using CUDA barriers for synchronization between loading, computation, and storing stages within a block. ```cuda __global__ void pipeline_kernel(...) { __shared__ st_bf<32, 64> S_in, S_out; rt_bf<32, 64> R; cuda::barrier in_ready, out_ready; // Stage 1: Load group<4>::memory::load_async(S_in, globals.A, {0, 0}, in_ready); // Stage 2: Compute (while loading) cuda::arrive_and_wait(in_ready); group<4>::load(R, S_in); // ... compute on R ... // Stage 3: Store (results can be computed while storing previous) group<4>::memory::store(S_out, R); group<4>::memory::store_async(globals.C, S_out, {0, 0}, out_ready); } ``` -------------------------------- ### Rebuild Kernel with Custom Options Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Command to rebuild a kernel after modifying the Makefile, specifying custom tile sizes. ```bash make TILE_M=64 TILE_N=64 ``` -------------------------------- ### Perform Matrix Multiply and Element-wise Addition Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Demonstrates matrix multiplication (D += A * B) using `mma` and element-wise addition (C = A + B) using `add`. ```cuda // Matrix multiply: D += A * B kittens::warpgroup::mma(D, A, B); // Element-wise: C = A + B kittens::warp::add(C, A, B); ``` -------------------------------- ### Register Tile Packing for Parallel Operations Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/06-advanced-features.md Illustrates packing multiple smaller tiles into registers for parallel execution. Operations on packed tiles can run concurrently. ```cuda // Two small tiles in one register space kittens::rt_bf<16, 16> tile1, tile2; // Operations on tile1 and tile2 can execute in parallel kittens::warp::mma(result1, A1, B1); kittens::warp::mma(result2, A2, B2); ``` -------------------------------- ### Run ThunderKittens Unit Tests Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Execute compiled unit tests and direct failed test outputs to the 'outputs/' folder. Note that MMA tests may occasionally fail due to floating-point arithmetic differences. ```bash mkdir outputs ./unit_tests printout ``` -------------------------------- ### load / load_async Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/03-memory-operations.md Loads data from global memory to shared memory. Supports both synchronous and asynchronous operations, with the latter requiring a barrier for synchronization. ```APIDOC ## load / load_async ### Description Loads data from global memory to shared memory. Supports both synchronous and asynchronous operations, with the latter requiring a barrier for synchronization. ### Method `load`: Synchronous `load_async`: Asynchronous ### Signature ```cuda template __device__ static inline void load(ST &dst, const GL &src, const COORD &idx); template __device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, cuda::barrier &bar); ``` ### Parameters - **dst** (ST) - Destination shared tile - **src** (GL) - Source global memory layout - **idx** (COORD) - Tile coordinate - **bar** (cuda::barrier) - (async only) Block-level barrier to signal completion - **axis** (int) - (template) Dimension axis ``` -------------------------------- ### Run NVIDIA Compute Unit Profiler Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Use 'ncu' to profile specific kernels, targeting them with a regular expression. This is useful for detailed analysis of kernel performance. ```bash # ncu (compute unit profiler) ncu --kernel_regex gemm kernel ``` -------------------------------- ### Thunderkittens Repository File Structure Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/00-index.md Overview of the directory and file organization within the Thunderkittens repository, including header files, kernels, and tests. ```text thunderkittens/ ├── include/ │ ├── kittens.cuh # Main entry point │ ├── common/ │ │ ├── common.cuh │ │ ├── base_types.cuh │ │ ├── base_ops.cuh │ │ └── util.cuh │ ├── types/ │ │ ├── types.cuh │ │ ├── register/ # rt, rv │ │ ├── shared/ # st, sv │ │ ├── global/ # gl, TMA │ │ ├── system/ # IPC, PGL, VMM │ │ └── tensor/ # tt │ ├── ops/ │ │ ├── ops.cuh │ │ ├── thread/ │ │ │ ├── memory/ │ │ │ ├── mma/ │ │ │ └── util/ │ │ └── group/ │ │ ├── memory/ # load/store │ │ ├── mma/ # matrix multiply │ │ ├── register/ # arithmetic │ │ ├── shared/ # shared ops │ │ └── util/ # sync, tma │ └── pyutils/ # Python integration ├── kernels/ # Pre-implemented examples │ ├── gemm/ │ ├── attention/ │ ├── based/ │ ├── lolcats/ │ └── flux/ ├── tests/ # Unit test suite ├── README.md └── LICENSE ``` -------------------------------- ### Fix libc10.so Error Source: https://github.com/hazyresearch/thunderkittens/blob/main/README.md Resolve 'libc10.so' errors by setting the LD_LIBRARY_PATH to the correct directory. First, find the path to torch.file, then append '/lib' to it. ```bash # Take the from below python -c "import torch; print(torch.file)" # And run the command below export LD_LIBRARY_PATH=/lib:$LD_LIBRARY_PATH ``` -------------------------------- ### Python Wrapper for CUDA Kernel Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Python code demonstrating how to create tensors and call a custom CUDA kernel (gemm) using PyBind11. ```python import torch import my_kernel_module # Create tensors A = torch.randn(1024, 1024, dtype=torch.bfloat16, device='cuda') B = torch.randn(1024, 1024, dtype=torch.bfloat16, device='cuda') C = torch.zeros(1024, 1024, dtype=torch.float32, device='cuda') # Call kernel my_kernel_module.gemm(A, B, C) # Check results print(C) ``` -------------------------------- ### Benchmark NVFP4 GEMM Operations Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md This command benchmarks NVFP4 GEMM operations for various matrix dimensions (N). It configures the profiler for block scaled GEMM with specific kernel names and distribution settings. ```bash # NVFP4_NVFP4_FP32_void_FP32 GEMM for N in 1024 2048 4096 8192 16384; do ./tools/profiler/cutlass_profiler \ --operation=block_scaled_gemm \ --m=${N} \ --n=${N} \ --k=${N} \ --output=nvfp4_${N} \ --kernels="cutlass3x_sm100_bstensorop_gemm_ue4m3xe2m1_ue4m3xe2m1_f32_void_f32" \ --warmup-iterations=500 \ --profiling-iterations=100 \ --verification-enabled=false \ --enable-best-kernel-for-fixed-shape=true \ --use_pdl=1 \ --dist="uniform,min:-6,max:6,scale:-1" done ``` -------------------------------- ### Minimal Compilation Test for ThunderKittens Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Attempt the simplest possible compilation command for a CUDA file using ThunderKittens. This helps isolate build issues. ```bash # Minimal test nvcc -DKITTENS_SM90 -std=c++20 test_tk.cu \ -I/path/to/thunderkittens/include -o test_tk # If this fails, check the error carefully ``` -------------------------------- ### Semaphore Operations Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/05-utilities-synchronization.md Illustrates the basic operations for initializing, acquiring, and releasing a CUDA counting semaphore. ```cuda // Initialize cuda::counting_semaphore sem; // Acquire (decrement) sem.acquire(); // Release (increment) sem.release(); // Release with count sem.release(count); ``` -------------------------------- ### Clone ThunderKittens Repository Source: https://github.com/hazyresearch/thunderkittens/blob/main/_autodocs/08-compilation-building.md Clone the ThunderKittens repository to your local machine. ```bash git clone https://github.com/HazyResearch/ThunderKittens.git cd ThunderKittens ``` -------------------------------- ### Benchmark MXFP8 GEMM Operations Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md This command benchmarks MXFP8 GEMM operations for various matrix dimensions (N). It configures the profiler for block scaled GEMM with specific kernel names and distribution settings. ```bash # MXFP8_MXFP8_FP32_void_BF16 GEMM for N in 1024 2048 4096 8192 16384; do ./tools/profiler/cutlass_profiler \ --operation=block_scaled_gemm \ --m=${N} \ --n=${N} \ --k=${N} \ --output=mxfp8_${N} \ --kernels="cutlass3x_sm100_bstensorop_gemm_ue8m0xe4m3_ue8m0xe4m3_f32_void_bf16" \ --warmup-iterations=500 \ --profiling-iterations=100 \ --verification-enabled=false \ --enable-best-kernel-for-fixed-shape=true \ --use_pdl=1 \ --dist="uniform,min:-448,max:448,scale:-1" done ``` -------------------------------- ### Analyze CUTLASS Benchmark Results Source: https://github.com/hazyresearch/thunderkittens/blob/main/kernels/gemm/baselines/cutlass/README.md This Python script analyzes CSV benchmark results to find the best performing kernel for each shape based on GFLOPs. It filters for successful GEMM operations and prints key performance metrics. ```python python3 <