HipKittens Framework: AI Kernels on AMD GPUs

• HipKittens framework is a C++ embedded DSL that enables high-performance AI kernel development on AMD GPUs via explicit tile-based abstractions.
• It employs asynchronous memory primitives and specialized wave scheduling to effectively map AI workloads to AMD CDNA architectures.
• The framework achieves assembly-level performance in AI applications, outperforming compiler-based approaches and enabling flexible algorithmic strategies.

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 (Hu et al., 11 Nov 2025).

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 (Hu et al., 11 Nov 2025). 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 (Hu et al., 11 Nov 2025).

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 $C$0 (memory-bound, favors 4-wave) and $C$1 (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 (Hu et al., 11 Nov 2025). Key findings:

Kernel	HK Peak/Best Baseline	HK vs Assembly	HK vs Compiler
BF16 GEMM (8k$C$2)	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 $C$3 demonstrates that HipKittens approaches $C$4 on compute-bound workloads (BF16 GEMM), exploits wave interleaving to mitigate memory bottlenecks in mixed-attention, and leverages grid-level swizzling to boost $C$5 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 (Hu et al., 11 Nov 2025). 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 (Hu et al., 11 Nov 2025). 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.