This article introduces HipKittens, a programming framework designed to unlock the performance of modern AMD GPUs for AI workloads, which is currently hampered by an immature software ecosystem. The authors identify key architectural differences between AMD's CDNA and NVIDIA's hardware, noting AMD's larger register file and chiplet design versus NVIDIA's superior SRAM, register reallocation, and asynchronous memory operations. HipKittens addresses three primary challenges: first, it provides optimized memory access patterns that navigate irregular AMD register layouts and compiler limitations through explicit register scheduling; second, it replaces the ineffective NVIDIA-centric "wave specialization" paradigm with two novel scheduling patterns, "8-wave ping-pong" and "4-wave interleave," which better suit AMD's architecture; third, it introduces a chiplet-aware grid scheduling algorithm to optimize data locality across the disaggregated L2 and shared LLC caches inherent to AMD's multi-chiplet design. The result is a set of primitives that enables developers to write kernels for workloads like GEMM and attention that achieve state-of-the-art performance on AMD hardware.
HipKittens: A Framework for High-Performance AI on AMD GPUs
The Core Problem: AMD's latest GPUs (e.g., MI355X) offer competitive peak compute performance but are hobbled by an immature software stack (HIPCC compiler), preventing AI developers from effectively utilizing the hardware. HipKittens is a library of programming primitives designed to bridge this gap.
Hardware Divergence: Effective optimization requires acknowledging architectural differences. Compared to NVIDIA's B200, AMD's MI355X has a 2x larger register file and more processors but lacks features like register reallocation, tensor memory acceleration (TMA), and robust asynchronous matrix multiplication instructions.
Optimized Memory Access: HipKittens provides solutions for AMD-specific memory challenges. This includes managing irregular register layouts that change with data type and instruction shape, implementing explicit register scheduling to circumvent compiler deficiencies, and reverse-engineering undocumented memory instruction behaviors to avoid bank conflicts.
Failure of a Paradigm: The dominant NVIDIA scheduling pattern, "wave specialization" (producer-consumer waves), underperforms significantly on AMD. The lack of register reallocation means producer waves are allocated registers they do not need, while compute-heavy consumer waves are starved for them, reducing arithmetic intensity.
Effective Scheduling Patterns: The authors propose two alternative, AMD-centric patterns:
8-wave ping-pong: Two waves per SIMD alternate between executing clusters of memory and compute instructions. This is programmable and sufficient for most high-performance kernels.
4-wave interleave: One wave per SIMD finely interleaves memory and compute operations. This offers higher performance at the cost of more complex code.
Chiplet-Aware Scheduling: Modern AMD GPUs use a chiplet architecture (8 "XCDs" on the MI355X), each with a private L2 cache. Naive thread block scheduling leads to poor cache reuse. HipKittens introduces a grid scheduling strategy that groups blocks to maximize data reuse across both the private L2 and the shared Last Level Cache (LLC).
Goal of Multi-Silicon AI: The project's explicit aim is to enable high-performance AI on diverse and open hardware. Making competitive alternatives to the incumbent hardware provider truly accessible is necessary for a healthier and more innovative ecosystem.