HazyResearch/HipKittens

Fast and Furious AMD Kernels

What it solves

HipKittens provides a set of low-level C++ programming primitives designed to help developers write high-performance AI kernels specifically for AMD GPUs. It addresses the challenge of creating a multi-silicon future where AI software can be efficiently ported across different hardware platforms by optimizing for the specific architectural nuances of AMD's CDNA3 and CDNA4 chiplet architectures.

How it works

The library is built from the hardware up, focusing on how the silicon actually operates. It implements several key primitives:

  • Tile Primitives: Memory operations sized for tensor core units that are bank-conflict free and coalesced to minimize address computation costs.
  • Python-inspired Functions: Lightweight bulk compute functions that wrap assembly and HIP.
  • Asynchronous Loads/Stores: Direct buffer loads to shared memory used to hide latency and address generation.
  • ** uma Scheduling and Overlapping**: Specific patterns (such as 8-wave ping pong and 4-wave interleave) to overlap compute and memory operations.

Who it’s for

It is intended for developers and researchers building high-performance AI kernels (such as GEMM, Attention, and Layernorm) who are targeting AMD hardware (MI300X, MI325X, MI350X, MI355X).

Highlights

  • Hardware-Centric Design: Specifically optimized for CDNA3 and CDNA4 architectures.
  • Broad Kernel Support: Includes implementations for BF16 GEMM, various Attention forwards/backwards (MHA, GQA, Causal), Rotary, and fused Layernorm.
  • AITER Integration: Officially serves as a backend for AITER.
  • Performance Benchmarking: Includes scripts to compare performance against baselines like Triton, CK, HipBLASLT, and Mojo.

Related

  • Project
    HazyResearch/ThunderKittens

    ThunderKittens is a header‑only CUDA library that provides a tiny, tile‑oriented DSL for writing fast deep‑learning kernels (e.g., GEMM, FlashAttention). It abstracts tensor‑core calls, async memory moves, and shared‑memory banking while keeping the code compact and type‑safe. Designed for Hopper/Blackwell (and Vera Rubin) GPUs, it lets researchers and engineers produce near‑peak‑performance kernels that can be called from C++ or wrapped for PyTorch.

  • 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
    ROCm/aiter

    A high-performance AI operator library for AMD ROCm that provides optimized GPU kernels for inference and training workloads across various AI frameworks.

  • Project
    stillwater-sc/universal

    A header-only C++ template library providing custom arithmetic plug-in types to optimize AI and HPC algorithms through mixed-precision and alternative number formats.

  • Project
    NVIDIA/cutlass

    A collection of CUDA C++ and Python DSL abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and the related linear algebra operations on NVIDIA GPUs.