---
title: 'HipKittens Framework: AI Kernels on AMD GPUs'
url: https://www.emergentmind.com/topics/hipkittens-hk-framework
type: topic
---

# HipKittens Framework: AI Kernels on AMD GPUs

The HipKittens (HK) framework is a C++-embedded domain-specific programming model designed to enable high-performance AI kernel development for AMD GPUs. Developed in the context of increasing demand for efficiently mapping AI workloads to AMD architectures, HipKittens generalizes tile-based abstractions previously popularized on NVIDIA hardware via DSLs such as ThunderKittens (TK) and Triton, while introducing AMD-specific algorithmic and architectural refinements. By leveraging explicit tile-level programming, asynchronous memory primitives, and specialized wave/block scheduling patterns, HipKittens enables assembly-level performance for a suite of matrix and attention kernels across AMD CDNA3 and CDNA4 platforms, outperforming existing compiler approaches and achieving competitiveness with hand-optimized assembly [2511.08083].

## 1. Motivation and Architectural Context

AMD’s CDNA3 and CDNA4 GPUs provide high raw FLOPs and memory bandwidth, but the lack of mature, maintainable, and high-level AI programming tools—unlike NVIDIA’s CUDA-ecosystem—has forced the reliance on labor-intensive, brittle assembly for peak kernel performance (e.g., AIITER). This leaves significant portions of the AI software stack unoptimized (e.g., PyTorch on AMD <30% of peak). Existing high-level compiler-based solutions (Triton, LLVM IR hints) fail to provide adequate register allocation control (e.g., AGPR vs VGPR), timely hardware intrinsic exploitation, and generally lag in performance, yielding up to 2–3× lower throughput compared to assembly [2511.08083]. The HipKittens framework targets these gaps by providing primitives and abstractions specifically matched to AMD’s compute (CUs with 4 SIMDs per unit), memory (on-chip LDS, multi-chiplet L2, global HBM), and matrix core MFMA instruction characteristics.

## 2. Programming Primitives and Memory Model

HipKittens adopts and customizes tile-level data structures and memory manipulation routines to match the hierarchical architecture of AMD GPUs. The fundamental programming abstractions are:

- **Register-Tile (rt):** Defined as `rt<dtype, T_m, T_k, layout, mfma_shape, [register_ranges]>`, this structure specifies a tile in registers, supporting explicit pinning to VGPR/AGPR registers and mapping to MFMA shapes. Types supported include fp32, bf16, fp16, fp8, and fp6. Explicit register pinning bypasses HIPCC limitations (e.g., AGPR as inputs), avoiding costly register movements.
- **Shared-Tile (st):** Analogous shared-memory tiles, `st<dtype, T_m, T_k, smem_layout>`, are declared with memory swizzling tailored to avoid bank conflicts. Unlike NVIDIA, AMD’s MFMA shapes each require distinct swizzle patterns; for example, a simple 8-column swap via a bitwise offset avoids LDS banking conflicts for 16×32 tiles.
- **Asynchronous Primitives:** Global to LDS memory transfers use `buffer_load_dwordx4` (bypassing VGPRs for AGPRs), while LDS to VGPR employs phase-aware `ds_read_b128` et al., with explicit software-managed synchronization via `vmcnt(x)`/`lgkmcnt(x)`.

The explicit memory model and register assignment are necessary due to the static partitioning of VGPRs/AGPRs per wave and the fine-grained control required to exploit AMD matrix core instructions.

## 3. Wave and Grid Scheduling Techniques

Performance portability and peak throughput on AMD platforms require distinct scheduling policies, both at the wave (intra-block) and grid (inter-block/chiplet) levels:

- **8-Wave Ping-Pong (Preferred):** Assigns 8 waves per block (2 waves per SIMD), partitioned into compute and memory groups. Compute waves issue MFMA instructions, while memory waves prefetch input tiles; roles alternate after barrier synchronization. This promotes large register tiles, low code complexity (~48 LoC hot loop), and achieves parity with hand-coded assembly for BF16/FP8 GEMMs and forward attention.
- **4-Wave Interleave:** Used for highly imbalanced (compute or memory) workloads, assigns 1 wave per SIMD, with fine-grained interleaving of load and compute, resulting in larger code size (~183 LoC) but up to 1.2× higher peak throughput in memory-bound scenarios (e.g., GQA backward).
- **Grid/Chiplet Scheduling:** HipKittens introduces a 2-step block-to-chiplet swizzle. Blocks are chunked for grouped L2 locality on XCDs, and tiles are traversed within chiplets to maximize L2 and LLC bandwidth, controlled by parameters such as window size $W$ and chunk size $C$. This raises $BW_{effective}$ by 15–20% vs. naive row-major layouts [2511.08083].

## 4. Algorithmic Instantiations

HipKittens encodes several reference AI algorithms as case studies for its primitives and scheduling policies:

- **BF16 GEMM:** Tiles matrices with per-block $T_m \times T_n = 256 \times 256$, per-wave $r_m \times r_n = 128 \times 64$ (8-wave) and $T_k = 64$. The main computational loop prefetches tiles with asynchronous loads, synchronizes, and multiplies-accumulates via MFMA. Arithmetic intensity is modeled as $AI \approx 2MNK/[(M/T_m)T_mT_k + (N/T_n)T_nT_k + (M/T_m)(N/T_n)T_mT_n] \cdot B$.
- **Attention Forward (GQA/MHA):** Implements $O = \text{SoftMax}(QK^T)V$ with L tiled into blocks of $L_b=128,256$. Waves process $D \times L_b$ output tiles, splitting memory and compute via 8-wave or 4-wave depending on memory/compute dominance. Trade-offs between $D=64$ (memory-bound, favors 4-wave) and $D=128$ (balanced, favors 8-wave) are observed.

These instantiations illustrate the framework’s flexibility and the hardware-tuned tradeoffs inherent to AMD platforms.

## 5. Performance Evaluation and Modeling

Extensive empirical evaluation across AMD CDNA3 (MI325X) and CDNA4 (MI355X) platforms compares HipKittens to baselines including AITER hand-tuned assembly, Composable Kernel (CK), Triton, HipBLASLt, PyTorch SDPA, and rocBLAS. All tests employ 500 warm-up and 100 measured runs with random-normal inputs [2511.08083]. Key findings:

| Kernel                | HK Peak/Best Baseline | HK vs Assembly | HK vs Compiler |
|-----------------------|----------------------|----------------|---------------|
| BF16 GEMM (8k$^3$)    | 1024/1018 TFLOP/s    | 1.01×          | 3× (Triton)   |
| GQA Attention Fwd     | 1.0–2.1× (AITER)     | 2.1× (d=64)    | 1.3–4.5×      |
| GQA Attention Bwd     | 1.8× (AITER)         | 2.5× (best)    | —             |
| Memory-bound Kernels  | 1.1–2.2× (PyTorch)   | —              | —             |

Roofline modeling with $P_{achieved} \leq \min(BW_{effective} \times AI, P_{peak})$ demonstrates that HipKittens approaches $P_{peak}$ on compute-bound workloads (BF16 GEMM), exploits wave interleaving to mitigate memory bottlenecks in mixed-attention, and leverages grid-level swizzling to boost $BW_{effective}$ in memory-bound tasks.

## 6. Cross-Vendor Generalization, Limitations, and Future Directions

HipKittens empirically validates that tile-based abstractions (tiles, async loads, bulk ops, barriers), essential in NVIDIA DSLs, are also fundamental for AMD—pointing to a common front-end API for high-performance AI kernels. However, lower-level algorithmic choices (tiling, swizzling, scheduling, register partitioning, MFMA shape) are necessarily hardware-specialized [2511.08083]. Practical limitations include developer burden of explicit pinning, scheduling, and the need for autotuners for new chiplet/geometries. Current efforts to integrate HK primitives into compiler IR (LLVM, Triton/HIPCC) and support for new datatypes (FP6) are ongoing. Extending the paradigm to additional operator families (MoE, convolutions, SSMs) remains an open avenue.

## 7. Impact and Implications

By abstracting away many architectural idiosyncrasies, HipKittens offers a minimal, C++-based DSL enabling portable, assembly-competitive AI kernels for AMD GPUs. This facilitates broader adoption of AMD hardware for deep learning workloads, reduces the software-hardware co-design gap, and sets the stage for unified programming models across heterogeneous acceleration platforms [2511.08083]. A plausible implication is the acceleration of AI ecosystem maturation on AMD, previously gated by software inertia (“CUDA moat”), and the potential for future cross-vendor domain-specific languages leveraging HipKittens’ architectural insights.

Source: https://www.emergentmind.com/topics/hipkittens-hk-framework