HipKittens: Fast and Furious AMD Kernels — L2 #

William Hu, Drew Wadsworth, Sean Siddens, Stanley Winata, Daniel Y. Fu, Ryan Swann, Muhammad Osama, Christopher Ré, Simran Arora | 2025-11 | https://arxiv.org/abs/2511.08083 Category: kernel | Tags: AMD GPU, tile-based DSL, CDNA4, MFMA, GEMM, attention kernel, chiplet scheduling, register pinning, wave scheduling

---

§1 TL;DR #

HipKittens: C++ tile-based DSL for AMD GPUs. Replaces NVIDIA's wave specialization with 8-wave ping-pong; bypasses HIPCC via pinned registers for AGPR→MFMA input; introduces chiplet-aware XCD grid scheduling for joint L2+LLC reuse. Matches AMD hand-tuned assembly on GEMM/attn fwd, 1.8–2.5× faster on GQA bwd.

---

§2 Q1 / Q2 / Q3 #

Q1 — 痛点 #

AMD MI355X offers competitive peak compute (2.5 PFLOPs BF16, 5.0 PFLOPs FP8, 8 TB/s HBM) versus NVIDIA B200, but peak-performance kernels require hand-tuned raw assembly (AITER library) that cannot scale to the breadth of AI workloads. The "CUDA moat" persists due to three AMD-specific architectural obstacles:

1. Static register allocation kills wave specialization. AMD hardware statically partitions 512 registers/SIMD across co-resident waves. The dominant NVIDIA producer-consumer wave specialization pattern wastes registers on non-compute producers, capping output tile size and arithmetic intensity. Result: wave specialization achieves only 80% of peak BF16 GEMM on MI355X (Tab. 2: 4P/8C = 893 TFLOPS vs 0P/8C = 1610 TFLOPS).
1. Non-compositional MFMA layouts prevent unified swizzles. NVIDIA matrix instructions share a compositional $16 \times 16$ core building block enabling a single swizzle strategy. Each AMD MFMA instruction uses an entirely different thread-element mapping with undocumented phase orderings (e.g., ds_read_b128 uses 4 phases across 64 banks; ds_read_b96 uses 8 phases across 32 banks). No single swizzle pattern eliminates bank conflicts for all co-occurring access patterns — proven formally in Appendix D.1.
1. Chiplet memory hierarchy requires joint L2+LLC optimization. MI355X has 8 XCDs (chiplets), each with private 4 MB L2. Naive row-major grid scheduling yields 36–55% L2 hit rate. Optimizing solely for L2 can degrade LLC from 95% to 24%, worsening overall bandwidth (Tab. 4: L2-only optimization = 991 TFLOPS, worse than naive row-major = 1113 TFLOPS).

Existing compilers fail to close the gap: Triton struggles with register lifetime tracking and intrinsic lowering; HIPCC prevents AGPR usage as MFMA inputs despite hardware support, inserting redundant v_accvgpr_read moves; Mojo's attention kernel suffers 50% performance loss from bank conflicts.

Q2 — 方法 #

HipKittens provides three pillars of AMD-specific primitives wrapped in the same tile-based front-end API as ThunderKittens (NVIDIA):

Pillar 1 — Programmable memory (§3.2). Developer-controlled register tiles bypass HIPCC limitations. Pinned register ranges (e.g., split_many_t>, 4>) let developers assign specific VGPR/AGPR registers to tiles, enabling AGPRs as direct MFMA inputs. Per-instruction swizzle patterns handle heterogeneous LDS banking: XOR-based address rewriting applied at the HBM address level (not LDS address like NVIDIA). Direct HBM→LDS async loads via buffer_load_dword bypass the register file.

Pillar 2 — Scheduling (§3.3). Two patterns replace wave specialization:

• 8-wave ping-pong (balanced workloads): 8 waves per thread block, 2 per SIMD, split into two groups of 4. On each SIMD, paired waves alternate — one issues MFMA compute at priority 3 (s_setprio), the other issues memory loads, then they swap via conditional barrier. Sufficient for GEMM and attention forward.
• 4-wave interleave (imbalanced workloads): 1 wave per SIMD, each wave issues both compute and memory in staggered sequence using compiler scheduling hints (sched_group_barrier). Needed for peak attention backward performance.

Pillar 3 — Non-programmable cache (§3.4). Algorithm 1 remaps grid block IDs in two steps: (a) XCD grouping assigns chunks of $C$ consecutive blocks to the same XCD for L2 locality; (b) windowed traversal of height $W$ folds the grid into rectangular L2 tiles, controlling LLC footprint across XCDs. The cost model $\text{BW} = \text{LLC\_BW} \times \text{LLC\_hit\%} + \text{L2\_BW} \times \text{L2\_hit\%}$ guides the $W$/$C$ tradeoff, leveraging the fact that L2 bandwidth $\approx 3\times$ LLC bandwidth on MI355X.

核心技术壁垒: The 8-wave ping-pong scheduling insight — recognizing that AMD's static register partition makes wave specialization fundamentally counterproductive (every producer wave halves effective register budget per compute wave, reducing achievable output tile size and arithmetic intensity) and that a simple alternating pattern with priority hints and conditional barriers suffices to fully overlap compute and memory at peak utilization, without assembly-level instruction interleaving. Replicating this requires deep understanding of the interaction between AMD's static register file, MFMA scheduling latencies, and memory pipeline occupancy — knowledge not present in public ISA documentation.

Q3 — 结果 #

Workload	HK best (TFLOPS)	Best baseline	HK speedup
BF16 GEMM ($M{=}N{=}K{=}8192$)	1610	AITER/HipBLASLT ~1600	~1.0×
FP8 GEMM (4-wave)	3327	Assembly baselines	~1.0×
GQA fwd ($D{=}128$, various $S$)	varies	AITER	1.0–2.1×
GQA bwd (causal + non-causal)	up to 1091	All baselines (best: AITER 384)	1.8–2.5×
Fused dropout-residual-LN	varies	AITER / PyTorch compiled	1.1–2.2×
vs Triton (all GEMMs)	—	Triton	1.3–3.0×

Correctness validated by pretraining Llama 1B and BERT 110M on Slim Pajama — matching perplexity of PyTorch/AITER baselines after 10B tokens.

---

§3 架構 / 方法図 #

Architecture overview #



Paper's Figure 1, verbatim (caption: "We study whether existing tile based programming primitives suffice for AMD kernels, or whether entirely new primitives are needed. Our study led to HipKittens...").

The left panel shows tile-based programmable memory management with developer-controlled register allocation (VGPR/AGPR pinning) and per-instruction swizzle patterns. The center panel illustrates the 8-wave ping-pong schedule — the paper's central innovation — where Wave A and Wave B on each SIMD alternate between compute (MFMA, high priority) and memory roles, synchronized by conditional barriers, replacing NVIDIA's producer-consumer wave specialization. The right panel shows chiplet-aware non-programmable memory scheduling (Algorithm 1) that remaps block IDs for joint L2+LLC cache reuse across 8 XCDs. The bottom demonstrates how these primitives compose to produce a suite of AMD AI kernels (GEMM, attention, memory-bound ops).

Swizzle mechanism #



Paper's Figure 4, verbatim (caption: "Swizzle pattern for a 16x32 tile of BF16s. Shared memory on AMD CDNA4 GPUs have different banking behavior depending on the instruction...").

The left panel shows an unswizzled layout suffering from 2-way bank conflicts under ds_read_b128. The right panel applies an XOR-based swizzle: starting from row 8, the first 8 columns and last 8 columns are swapped. This simultaneously eliminates bank conflicts for both ds_read_b128 (row-major, 64 banks × 4 phases) and ds_read_b64_tr_b16 (column-major transpose), allowing the same LDS region to serve both access patterns. This dual-pattern bank-conflict freedom is critical for attention backward where mixed MFMA shapes ($16 \times 16 \times 32$ and $32 \times 32 \times 16$) require both row and column register tile layouts from the same shared memory.

Target operation #

Operation	Input shape	Output shape	Input dtype	Acc dtype	Output dtype
BF16 GEMM	A: $[M, K]$, B: $[K, N]$	C: $[M, N]$	BF16	FP32	BF16
FP8 GEMM	A: $[M, K]$, B: $[K, N]$	C: $[M, N]$	MXFP8	FP32	FP8/BF16
MHA fwd	Q,K,V: $[B, H, S, D]$	O: $[B, H, S, D]$	BF16	FP32	BF16
GQA fwd	Q: $[B, H_q, S, D]$, K,V: $[B, H_{kv}, S, D]$	O: $[B, H_q, S, D]$	BF16	FP32	BF16
MHA/GQA bwd	dO, Q, K, V, O	dQ, dK, dV	BF16	FP32	BF16
Fused dropout-res-LN	$[B, H, D]$	$[B, H, D]$	BF16	FP32	BF16
RoPE	$[B, H, S, D]$	$[B, H, S, D]$	BF16	—	BF16

Evaluated regimes: Batch 16; heads 16 or 64 (query) / 8 (KV for GQA); $D \in \{64, 128\}$; $S \in \{1024, 2048, 4096, 8192\}$; GEMM $M{=}N{=}K \in \{1024, \ldots, 14592\}$. All prefill-regime (large batch × $S^2$ attention, compute-bound). Causal and non-causal masks for attention. No structured sparsity.

Hardware model #

Primary target: AMD MI355X (CDNA4, gfx950)

Resource	Value
CUs	256 (8 XCDs × 32 CUs/XCD)
SIMDs per CU	4
Waves per SIMD (HK)	2 (8-wave) or 1 (4-wave)
Registers per SIMD	512 (split: 256 VGPR + 256 AGPR at 1 wave/SIMD; 128+128 each at 2 waves/SIMD)
BF16 matrix core peak	2.5 PFLOPs
MXFP8 matrix core peak	5.0 PFLOPs
MXFP6 / MXFP4 peak	10.1 PFLOPs
HBM capacity / bandwidth	288 GB / 8.0 TB/s
L2 cache	4 MB per XCD (32 MB total, not shared across XCDs)
LLC	Shared across all XCDs, between L2 and HBM
L2 miss penalty	~300 ns worst case
LLC miss penalty	~500 ns worst case

Secondary target: AMD MI325X (CDNA3, gfx942) — 38 CUs/XCD, 65 KB LDS per CU (smaller than MI355X). 8-wave pattern adapted: double-buffers in register file instead of LDS due to LDS size constraint.

MFMA instructions used:

• $16 \times 16 \times 32$ (BF16): default for GEMM and attention — finest scheduling granularity
• $32 \times 32 \times 16$ (BF16): attention backward only — larger output per instruction, reduces register pressure on dK/dV accumulation
• $16 \times 16 \times 128$ (FP6): FP6 GEMM (Appendix F, preliminary)

Launch config (BF16 GEMM): Output tile $256 \times 256$, $K_{\text{step}} = 64$, grid $\lceil M/256 \rceil \times \lceil N/256 \rceil$, 8 waves × 64 threads = 512 threads/block, WARPS_M=2, WARPS_N=4, DOT_SLICE=32.

---

§4 作者証明 #

符号表 #

Symbol	Meaning
$M, N, K$	GEMM dimensions
$B, H, S, D$	Batch, heads, sequence length, head dimension
$H_q, H_{kv}$	Query heads, KV heads (GQA)
$W$	Window height in Algorithm 1 (controls L2 tile shape)
$C$	Chunk size in Algorithm 1 (consecutive blocks per XCD)
$P / C_w$	Producer / consumer wave count in wave specialization
VGPR / AGPR	Vector general-purpose / accumulator registers
XCD	Accelerator complex die (AMD chiplet)
CU / SIMD	Compute unit / SIMD execution unit within CU
LDS	Local data share (AMD shared memory)

方程物理意義 #

The paper has one explicit equation — the cache bandwidth cost model:

$$\text{BW} = \text{LLC\_BW} \times \text{LLC\_hit\%} + \text{L2\_BW} \times \text{L2\_hit\%}$$

Effective memory bandwidth is the weighted sum of L2 contributions (fast, per-XCD, ~3× higher bandwidth) and LLC contributions (slower, shared across chiplets). The model reveals a non-obvious coupling: L2 and LLC hit rates are not independent. Grouping many blocks per XCD (large $C$) improves L2 hit rate but causes cross-XCD LLC thrashing when different XCDs access disjoint regions. The parameters $W$ (window height) and $C$ (chunk size) in Algorithm 1 provide the control knobs: $W$ determines the shape of the L2 tile (tall window → more row reuse), while $C$ determines how many blocks share an XCD (larger $C$ → better L2, risk of worse LLC).

Roofline placement #

Operation	AI (FLOPs/byte)	Roofline knee (MI355X)	Regime
BF16 GEMM ($8192^3$)	$\approx 2731$	~312	Compute-bound
FP8 GEMM ($8192^3$)	$\approx 5461$	~625	Compute-bound
Attn fwd ($S{=}4096, D{=}128$)	$\approx 32$	~312	Compute-bound
LayerNorm / RoPE	$< 10$	~312	Memory-bound

Roofline knee: $\text{peak FLOPs} / \text{peak BW} = 2500 \times 10^{12} / 8.0 \times 10^{12} \approx 312$ FLOPs/byte (BF16). GEMM and attention sit deeply in the compute-bound regime. Memory-bound kernels (LayerNorm, RoPE) are bandwidth-limited — here HK's advantage comes from correct buffer load instruction selection (HBM→LDS direct, better L2 hit rates) rather than compute scheduling.

% peak derivation #

BF16 GEMM ($M{=}N{=}K{=}8192$): 1610 TFLOPS / 2500 TFLOPS = 64.4% of theoretical peak.

From launch config: 8 waves × 64 threads = 512 threads/block, output tile $256 \times 256$, $K_{\text{step}} = 64$. Each step performs $2 \times 256 \times 256 \times 64 = 8.39 \times 10^6$ FLOPs. With MFMA $16 \times 16 \times 32$, each wave performs $\text{WARPS\_M} \times \text{WARPS\_N} = 8$ MFMA instructions per DOT_SLICE step. Pipeline utilization depends on MFMA latency vs memory latency; at 2 waves/SIMD, the ping-pong hides memory behind partner wave's compute. The 64.4% figure aligns with NVIDIA TK on B200 (1538/2200 = 69.9%), suggesting tile-based DSLs achieve similar hardware efficiency across vendors.

FP8 GEMM: 3327 / 5000 = 66.5% (4-wave interleave); 3222 / 5000 = 64.4% (8-wave).

Tile / launch optimality #

Why $256 \times 256$ and not larger? At 2 waves/SIMD: 128 AGPRs per wave. With WARPS_M=2, WARPS_N=4, each wave accumulates a $(256/2) \times (256/4) = 128 \times 64$ sub-tile = 8192 BF16 elements. In the MFMA accumulator (FP32), that requires $8192 \times 2 = 16384$ bytes → 256 32-bit AGPRs — exactly exhausting the 128 AGPR budget at 2 waves/SIMD (since accumulator elements pack 2 per register in the accumulation layout). A $320 \times 256$ tile would require $10240 \times 2 / 4 = 5120$ bytes = 160 AGPRs per wave → overflow.

Why not smaller? Arithmetic intensity scales with tile dimension. Tab. 2: $128 \times 256$ (4P/8C) = 893 TFLOPS; $192 \times 256$ (0P/8C) = 1281 TFLOPS; $256 \times 256$ (0P/8C) = 1610 TFLOPS. The 25% tile dimension reduction from $256$ to $192$ costs 20% performance.

Why $16 \times 16 \times 32$ MFMA default? Finer granularity → more MFMA instructions per K-step → better interleaving with memory ops within compute clusters. Attention backward uses mixed shapes ($16 \times 16 \times 32$ + $32 \times 32 \times 16$) because the $Q \cdot K^T$ and $(dS) \cdot V$ matmuls have different register pressure profiles.

Design space & constraint derivation #

Axis	Value chosen	Rejected alternatives	Blocking constraint
Schedule (GEMM, attn fwd)	8-wave ping-pong	Wave specialization	Static register partition: 4P/8C = 893 vs 0P/8C = 1610 TFLOPS
Schedule (attn bwd)	4-wave interleave	8-wave ping-pong	Compute/memory imbalance: 4-wave = 1091 vs 8-wave = 894 TFLOPS
MFMA shape (default)	$16 \times 16 \times 32$	$32 \times 32 \times 16$ alone	Finer scheduling granularity; $32 \times 32$ used only in bwd for register pressure
Output tile	$256 \times 256$	Larger	AGPR budget at 2 waves/SIMD: exactly saturated at $256 \times 256$
Register mgmt (attn bwd)	Pinned registers	HIPCC-managed	HIPCC inserts v_accvgpr_read: pinned 1024 vs unpinned 855 TFLOPS
Swizzle	Per-instruction XOR	Single unified swizzle	Formally impossible: ds_write_b64 64-bit chunks conflict with ds_read_b128 128-bit contiguous
Cache scheduling	Joint L2+LLC via $W$, $C$	L2-only	L2-only (79% L2, 24% LLC) = 991 < row-major (55% L2, 95% LLC) = 1113 TFLOPS
Sync mechanism	LDS atomics	mbarrier	AMD lacks mbarrier HW; atomics overhead negligible (1281 vs 1278 TFLOPS)
Pipeline depth	2-deep (LDS double buffer)	Deeper	LDS ~40% smaller than NVIDIA SMEM per-processor
Global→LDS	buffer_load_dword (async)	Register-mediated	Avoids register pressure; analogous to NVIDIA TMA

6 minimum checks #

1. Hidden assumptions. 8-wave ping-pong assumes balanced compute and memory durations per tile step. When imbalanced (attention backward), it underperforms and 4-wave interleave is needed. Explicitly acknowledged in §3.3.2 with quantified comparison (Tab. 3).
1. Edge cases. Algorithm 1 has worst-case behavior when grid tiles mod 8 is problematic (57 tiles for $M{=}N{=}K{=}14592$). The paper documents this case (Tab. 4) and shows the algorithm still provides 19% improvement via $W$/$C$ tuning.
1. Bandwidth model linearity. Equation 1 assumes L2 and LLC contributions are additive. In practice L2 misses become LLC accesses, creating a sequential dependency. The model holds approximately because L2 and LLC latencies dominate different pipeline stages — validated empirically (Tab. 4 predictions match measured TFLOPS ordering).
1. Compiler interaction. Pinned registers bypass HIPCC's optimization passes for affected code regions, losing automatic instruction reordering. Mitigated via explicit sched_barrier and sched_group_barrier hints (§D.4). Potential fragility across ROCm versions acknowledged implicitly.
1. Generational portability. Phase-bank table (Tab. 5) is reverse-engineered for CDNA3/CDNA4 via custom solver (§D.2). No guarantee these orderings persist in future CDNA generations. Validated on two generations, but per-generation adaptation of swizzle patterns and LDS double-buffering strategies is required.
1. Correctness validation. End-to-end via pretraining perplexity match (Llama 1B + BERT 110M, 10B tokens). No formal kernel-level numerical error bounds (max abs diff, mean rel diff) reported. No discussion of overflow/underflow/subnormal handling.

---

§5 実験与数据 #

GEMM #



Paper's Figure 6 (left), verbatim (caption: "GEMM. We compare HK BF16 and FP8 GEMMs to the strongest available baselines.").

HK competes with AITER/HipBLASLT (hand-optimized assembly) across all square GEMM sizes on MI355X. At $M{=}N{=}K{=}8192$, HK achieves ~1610 TFLOPS, matching assembly. The gap with Triton is dramatic: 1.3–3.0× across all sizes, demonstrating that LLVM-based compilation on AMD cannot close the performance gap without framework-level scheduling and memory abstractions. A single 8-wave kernel schedule generalizes across all problem sizes without per-size tuning.



Paper's Figure 6 (right), verbatim (caption: "GEMM. We compare HK BF16 and FP8 GEMMs to the strongest available baselines.").

FP8 GEMM shows a similar pattern: HK achieves 3222 TFLOPS (8-wave) and 3327 TFLOPS (4-wave), competitive with assembly baselines. The 4-wave interleave provides a modest 3% improvement over 8-wave at the cost of 3.8× more code (183 vs 48 LoC, Tab. 3), suggesting the simpler 8-wave pattern is the better default.

Attention forward #



Paper's Figure 7, verbatim (caption: "Attention forwards. We compare HipKittens GQA ... to the strongest available baselines. We use batch 16, query heads 64, key value heads 8, head dim 64 and 128.").

GQA forward (batch 16, 64 query heads, 8 KV heads) shows HK outperforming AITER assembly by 1.0–2.1×, PyTorch SDPA by 1.3–4.5×, CK by 1.0–1.4×, and Triton by 1.2–4.5×. The 8-wave ping-pong schedule is sufficient — online-softmax vector operations (col_max, exp2, col_sum) are interleaved with MFMA instructions within the compute cluster, filling pipeline bubbles. At $D{=}64$, the advantage is largest because no existing baseline is tuned for this head dimension.

Attention backward #



Paper's Figure 8, verbatim (caption: "Attention backwards. We compare HipKittens GQA ... to the strongest available baselines. We use batch 16, query heads 64, key value heads 8, and head dim 128.").

GQA backward is HK's strongest result: 1.8–2.5× over all baselines including AITER assembly. All three pillars contribute simultaneously: mixed MFMA shapes ($16 \times 16 \times 32$ + $32 \times 32 \times 16$) require per-instruction swizzle; explicit register pinning enables AGPRs as MFMA inputs; and the 4-wave interleave pattern handles the imbalanced compute/memory ratio. AITER achieves only 272–384 TFLOPS (causal/non-causal, seq 8192) while HK reaches 1091 TFLOPS — a gap that demonstrates raw assembly does not scale to complex kernels with multiple MFMA shapes and register pressure requirements.

Memory-bound kernels #



Paper's Figure 9, verbatim (caption: "Memory bound. We compare HipKittens fused dropout-residual-layernorm and rotary kernels to the strongest available baselines at batch 16, heads 16, and head dim 128.").

Fused dropout-residual-LayerNorm and RoPE are memory-bandwidth limited. HK outperforms AITER and PyTorch compiled kernels by 1.1–2.2× across settings. The advantage comes from HK's correct use of buffer_load instructions (which PyTorch's torch.compile does not emit on AMD) and better L2 hit rates — torch-compiled LayerNorm exhibits 23% lower L2 hit rate than HK.

Key tables #

Table 2 — Wave specialization analysis (BF16 GEMM, $M{=}N{=}K{=}8192$, MI355X):

Config	MFMA Shape	Output Tile	TFLOPS
HK 4P/8C	$16 \times 16 \times 32$	$128 \times 256$	893
HK 4P/12C	$16 \times 16 \times 32$	$192 \times 256$	1278
HK 0P/8C	$16 \times 16 \times 32$	$192 \times 256$	1281
HK 0P/8C	$16 \times 16 \times 32$	$256 \times 256$	1610
TK (B200)	$256 \times 256 \times 16$	$256 \times 256$	1538
CUTLASS (B200)	$256 \times 256 \times 16$	$256 \times 256$	1570

Zero producers + maximal output tile is optimal. 4P/8C wastes 4/8 of register file on non-compute waves, shrinking the achievable tile from $256 \times 256$ to $128 \times 256$ and halving arithmetic intensity.

Table 1 — Register pinning impact (MHA bwd, batch 16, heads 16, $D{=}128$, MI355X):

Method	Seq Length	TFLOPS
HK (compiler-managed)	4096	855
HK + pinned registers	4096	1024
AITER (assembly)	4096	1018
HK (compiler-managed)	8192	909
HK + pinned registers	8192	1091
AITER (assembly)	8192	1169

Pinned registers yield 20% improvement by enabling AGPRs as MFMA A/B inputs, matching or approaching assembly performance.

Table 4 — Chiplet cache scheduling (BF16 GEMM, MI355X):

Block Order	L2%	LLC%	Mem BW	TFLOPS
Row-major ($M{=}N{=}K{=}9216$)	55%	95%	15.1 TB/s	1113
XCD W7/C216	79%	24%	14.9 TB/s	991
XCD W5/C25	75%	93%	18.3 TB/s	1145
Row-major ($M{=}N{=}K{=}14592$)	36%	76%	10.7 TB/s	900
XCD W8/C542	79%	7%	13.9 TB/s	980
XCD W8/C64	78%	55%	16.6 TB/s	1068

The $M{=}N{=}K{=}14592$ case is especially sensitive: 57 tiles across 8 XCDs causes worst-case L2 reuse under row-major. Algorithm 1 provides up to 19% improvement via joint L2+LLC optimization.

Optimization techniques inventory #

Technique	Target bottleneck	HW primitive	Measured contribution
8-wave ping-pong	Compute-memory overlap	s_setprio, conditional barrier	0P/8C = 1610 vs 4P/8C = 893 TFLOPS (Tab. 2)
4-wave interleave	Imbalanced pipeline	sched_group_barrier	4w = 1091 vs 8w = 894 TFLOPS on MHA bwd (Tab. 3)
Pinned registers	HIPCC AGPR restriction	Inline asm register ranges	1024 vs 855 TFLOPS on MHA bwd (Tab. 1)
Per-instruction swizzle	LDS bank conflicts	ds_read_b128 / ds_read_b64_tr_b16	0 conflicts vs 2-way conflicts (Fig. 4)
XCD swizzle	Cache thrashing across chiplets	Block ID integer remapping	1145 vs 1113 TFLOPS (Tab. 4, +19% on 14592)
Direct HBM→LDS loads	Register pressure from staging	buffer_load_dword variants	Qualitative: removes register-file bottleneck
Compiler sched hints	Instruction ordering within clusters	sched_barrier, sched_group_barrier	Qualitative: enables cluster-level control (§D.4)
Online softmax interleaving	Softmax-MFMA pipeline bubble	col_max/exp2 interleaved with mma_AtB	Qualitative: VALU ops fill MFMA latency gaps

Numerical considerations #

• Accumulator precision: FP32 for all BF16 and FP8 kernels — standard practice for numerical stability.
• Online softmax: Max-subtract → $\exp_2$ → accumulate with proper rescaling per tile step. Temperature: $\sqrt{1/D} \cdot \log_2(e)$ (= 0.08839 × 1.44270 for $D{=}128$).
• End-to-end validation: Pretraining convergence match (Llama 1B + BERT 110M, 10B tokens). Statistical equivalence suffices; bit-exact agreement not required.
• No formal error bounds: Max absolute difference and mean relative difference against reference implementations are not reported. No overflow/underflow/subnormal discussion.
• FP6 (preliminary): ds_read_b96 loads 12 bytes per instruction. HIPCC-compiled FP6 GEMM at size 16384 spills 54 registers to scratch, producing slow and incorrect output. HK's pinned registers eliminate all spills. Manual 8-cycle v_nop required between v_mov_b32_e32 and dependent MFMA to satisfy hardware timing constraint.

---

§6 論証鎖 #

Step	Claim	Evidence	Depends on
1	AMD's static register partition makes wave specialization counterproductive	Tab. 2: 4P/8C = 893 TFLOPS (tile $128 \times 256$) vs 0P/8C = 1610 TFLOPS (tile $256 \times 256$). Each producer wave consumes 1/8 of SIMD registers without contributing compute output, forcing smaller tiles and lower arithmetic intensity	Arch fact: AMD statically divides 512 regs across co-resident waves
2	8-wave ping-pong achieves compute-memory overlap without dedicated producers	Tab. 2: 0P/8C at 1610 TFLOPS; Tab. 3: 8-wave FP8 GEMM = 3222 TFLOPS. Two waves on same SIMD alternate roles via conditional barrier — MFMA units never idle during partner's memory phase	Step 1 (producers are wasteful → remove them)
3	AMD's heterogeneous MFMA layouts require per-instruction swizzle patterns	Appendix D.1: formal proof that ds_write_b64 (64-bit chunk granularity) and ds_read_b128 (128-bit contiguous requirement) cannot share a single XOR swizzle. Fig. 4: demonstrates dual-pattern swizzle for $16 \times 32$ tile	Arch fact: each MFMA instruction has unique thread-element mapping
4	Pinned registers close the gap with assembly on attention backward	Tab. 1: pinned = 1024 vs unpinned = 855 TFLOPS (seq 4096), matching AITER assembly (1018 TFLOPS). Bypass enables direct AGPR→MFMA input, eliminating v_accvgpr_read overhead	Step 3 (mixed MFMA shapes in bwd amplify register pressure)
5	Chiplet-aware grid scheduling requires joint L2+LLC optimization	Tab. 4: L2-only optimization (79% L2, 24% LLC) = 991 TFLOPS < row-major (55% L2, 95% LLC) = 1113 TFLOPS. Joint (75% L2, 93% LLC) = 1145 TFLOPS. Eq. 1 explains: L2 BW $\approx 3\times$ LLC BW makes LLC drops costly	Arch fact: 8-XCD chiplet with per-XCD L2 and shared LLC
6	All three pillars compose to produce kernels matching or exceeding hand-optimized assembly	§4 full results: GEMMs match AITER; attn fwd 1.0–2.1× AITER; GQA bwd 1.8–2.5× all baselines; memory-bound 1.1–2.2×. Correctness validated via pretraining convergence (Llama 1B, BERT 110M)	Steps 2 + 4 + 5

---

§7 実現 cross-reference #

Repository: https://github.com/HazyResearch/HipKittens (public, C++/HIP + Python bindings, Apache-2.0)

Code citations (from paper Appendix E listings) #

1. GEMM kernel launcher (Paper Fig. 21, lines 1–63): Defines output tile BLOCK_SIZE=256, K_STEP=64, WARPS_M=2, WARPS_N=4. Prologue pre-loads A/B tiles from HBM to LDS via buffer_load_dword. Conditional barrier (if (warpId < 4) barrier.wait()) staggers wave group B behind group A, establishing initial ping-pong offset.
1. GEMM inner loop (Paper Fig. 22, lines 1–50): 8-wave ping-pong hot loop. Compute cluster: s_setprio(3) → mma_ABt(acc, a_tile, b_tile) × DOT_SLICE iterations → s_setprio(0). Memory cluster: load(smem_a, gmem_a_next) → vmcnt(0) fence. Barrier swap between groups. Cache scheduling via chiplet_transform_chunked for XCD-aware block remapping. MI325X variant uses register-file double buffering instead of LDS double buffering.
1. Attention forward inner loop (Paper Fig. 23): 8-wave ping-pong with online-softmax interleaving. Compute cluster chains: mma_AtB ($Q \cdot K^T$) → col_max → sub_col → exp2 → col_sum → mma_ABt (softmax $\cdot V$). Uses sched_barrier_pairs and sched_barrier_exp_pairs template helpers for VALU-MFMA instruction ordering. Epilogue: transpose output tile, store to HBM, compute LSE for backward pass.
1. Register pinning API (Appendix D.3): rt with Q_ranges = split_many_t>, 4> pins registers v[24:39] to the Q tile. Enables AGPR as MFMA A/B input without v_accvgpr_read.
1. XCD swizzle (Algorithm 1, §3.4): 24-line pseudocode. Flatten 2D grid → cycle-aligned boundaries ($\text{nXCD} \times C$) → de-interleave round-robin → windowed traversal of height $W$ → remapped $(b.x', b.y', b.z)$. Purely integer arithmetic per thread block, negligible overhead.

関鍵実現細節 #

1. Shared memory phase orderings are undocumented in AMD's ISA. The thread-to-phase assignment for LDS instructions (ds_read_b128: 4 phases, ds_read_b96: 8 phases, ds_write_b64: 4 phases, ds_read_b64: 2 phases) is absent from all public CDNA references. The authors created a solver (§D.2) that iterates over thread pairs, performs shared memory instructions on the same bank, and infers phase membership from observed bank conflicts. Without this reverse-engineering, designing bank-conflict-free swizzles for AMD is impossible. Tab. 5 documents the complete phase-bank table — the first public documentation of this behavior.
1. HIPCC silently miscompiles FP6 register shuffles. At GEMM size $16384^3$, HIPCC-compiled FP6 kernels spill 54 registers to scratch memory, producing both slow and incorrect output (§F). Root cause: ds_read_b96 places data at non-contiguous register positions, requiring three v_mov_b32_e32 shuffles that HIPCC cannot schedule correctly. HK's pinned registers eliminate all spills. An additional 8-cycle latency between v_mov_b32_e32 and dependent MFMA must be enforced via manual v_nop insertion — another timing constraint the compiler fails to handle.

Portability analysis #

Dimension	Assessment
CDNA3 → CDNA4	Validated on MI325X and MI355X. Different LDS sizes force different double-buffering strategies (register-file on MI325X, LDS on MI355X). Different CU counts per XCD (38 vs 32) require Algorithm 1 re-tuning. Phase orderings appear consistent across generations but are not guaranteed
AMD → NVIDIA	Front-end tile API mirrors ThunderKittens intentionally. Backend is entirely different: MFMA vs wgmma, LDS vs SMEM, buffer_load vs TMA, sched_barrier vs mbarrier, s_setprio vs no equivalent. A unified front-end with vendor-specific backends is the stated architectural vision
HIP C++ → Triton	Triton lacks register lifetime control and does not default to buffer_load on AMD. Inline assembly in Triton is possible but defeats the abstraction. HK is complementary to, not a replacement for, compiler approaches
HIP C++ → CK/CUTLASS	CK uses deeply nested C++ templates with high complexity. HK is explicitly simpler (48 LoC FP8 GEMM vs CK's template hierarchy) while matching performance
gfx950 → future CDNA	Phase orderings may change. Register file size/partitioning may evolve. Algorithm 1 parameters need per-generation re-tuning. Pinned register approach bypasses compiler → fragile across ROCm versions

---

§8 Deployment context #

Serving stack integration. HK kernels provide attention (MHA/GQA forward and backward) and GEMM operators with Python bindings for PyTorch integration. They can replace the existing SDPA backend on AMD (259 TFLOPS for Llama GQA backward on MI355X — 24% of HK's throughput) or the AITER assembly kernels. Not yet integrated into vLLM, SGLang, or DeepSpeed-Inference, but the operator interface (C++ kernel + Python binding) is compatible with custom operator registration paths in these frameworks.

Regime specialization. All benchmarked configurations are prefill-regime: large batch ($B{=}16$) × full sequence attention. The paper does not evaluate decode-regime (batch=1, single-token, memory-BW-bound) workloads. The tuned configurations target $(S \in [1024, 8192], B{=}16, D \in \{64, 128\})$. Production inference at lower batch sizes, longer sequences, or with paged KV cache is not tested.

Prefill / decode split. All attention kernels are prefill kernels ($B \times S^2$ compute). No specialized decode kernel is provided. The 8-wave ping-pong pattern assumes compute-memory balance, which breaks down in decode where memory bandwidth dominates and the roofline regime shifts.

Fusion scope. Attention forward fuses $Q K^T$ + softmax + $O V$ into one kernel. Memory-bound kernel fuses dropout + residual add + LayerNorm. No cross-layer fusion (e.g., attention + MLP), no RoPE fusion into attention, no QKV projection fusion.

Launch shape. GEMM: $\lceil M/256 \rceil \times \lceil N/256 \rceil$ grid × 512 threads. Attention: one thread block per $(batch, head)$ pair. These match the tuned configurations.

---

§9 Source walkthrough #

Code is public at https://github.com/HazyResearch/HipKittens. The paper provides detailed kernel listings in Appendix E (Figures 21–23). Analysis of the critical code paths:

9.1 GEMM kernel entry (Paper Fig. 21) #

The launcher sets up the tile decomposition and memory layout. The prologue has two phases:

1. Collaborative load: All 8 waves cooperatively load the first A and B tiles from HBM to LDS using buffer_load_dword (direct async, bypasses register file). Each wave loads a portion proportional to its share of the $256 \times 256$ output tile.
1. Staggered start: A conditional barrier (if (warpId < 4) barrier.wait()) delays wave group B by one iteration behind wave group A. This establishes the initial ping-pong offset — when group A enters its first compute cluster, group B is still completing its first memory cluster, and the SIMD hardware overlaps their instruction streams.

9.2 GEMM inner loop (Paper Fig. 22) #

The hot loop is the 8-wave ping-pong at its simplest — 48 lines for FP8 GEMM (Tab. 3). Each iteration:

1. Compute cluster: s_setprio(3) raises wave priority → load(a_reg, smem_a) + load(b_reg, smem_b) from LDS to registers → mma_ABt(acc, a_reg, b_reg) × DOT_SLICE iterations → s_setprio(0) lowers priority. The s_setprio(3) ensures the computing wave group wins arbitration for shared execution units over the memory-issuing partner wave.
1. Memory cluster: load(smem_a_next, gmem_a) + load(smem_b_next, gmem_b) issues HBM→LDS async loads → __builtin_amdgcn_s_waitcnt(vmcnt(0)) drains all pending global loads.
1. Barrier swap: barrier.arrive_and_wait() synchronizes both wave groups, then they swap roles. Pointer swap alternates between double-buffered LDS regions.

XCD-aware scheduling: chiplet_transform_chunked remaps block coordinates before the loop begins, using Algorithm 1's integer arithmetic to assign consecutive chunks to the same XCD.

9.3 Attention forward (Paper Fig. 23) #

The attention kernel extends the GEMM pattern with online-softmax interleaving:

1. Prologue: 8 waves collaboratively load first K/V tiles to shared memory. Each wave loads its personal Q tile ($32 \times 128$ per wave per head/batch). Initial $Q \cdot K^T$ matmul + first half of softmax (col_max, subtract, exp2).
1. Compute cluster: Chains MFMA with VALU: mma_AtB ($Q \cdot K^T$) → col_max → sub_col → exp2 → col_sum → rescale previous accumulator → mma_ABt (softmax weights $\cdot V$). The template helpers sched_barrier_pairs and sched_barrier_exp_pairs interleave VALU softmax ops into MFMA latency gaps.
1. Epilogue: Final online-softmax rescaling (multiply $O$ by $1/\ell$), transpose from MFMA accumulator layout to row-major, store to HBM, compute and store LSE for backward pass.

---

§10 Software → Hardware reverse implication #

What HipKittens' optimizations reveal about desirable future AMD hardware capabilities:

1. Dynamic register partitioning. The entire 8-wave ping-pong exists to work around AMD's static register division across co-resident waves. If future CDNA supported dynamic register reallocation — allowing producer waves to cede unused registers to consumer waves at runtime — wave specialization would become viable on AMD, eliminating the current tradeoff between scheduling simplicity and register efficiency. The 8-wave pattern + the 4-wave interleave variant together represent ~50% of HK's novelty; dynamic partitioning would make both unnecessary.

2. Compiler-accessible AGPR paths. HIPCC prevents AGPRs as MFMA A/B operands despite hardware support. The pinned register workaround bypasses the compiler entirely, losing all optimization passes for affected code regions. If HIPCC (or LLVM's AMDGPU backend) exposed AGPR→MFMA input paths as a first-class compiler feature, attention backward performance would improve without requiring dangerous inline assembly. Current situation forces a choice between compiler safety and hardware utilization.

3. Compositional matrix instruction layouts. The per-instruction swizzle explosion stems from each MFMA shape using a unique thread-element mapping. If future CDNA defined a compositional core matrix building block (analogous to NVIDIA's $16 \times 16$ structure), a single swizzle strategy would suffice for all instruction shapes — reducing the swizzle codepath from $O(n_{\text{instructions}})$ to $O(1)$ and eliminating the need for reverse-engineered phase tables.

4. Documented LDS phase orderings. The phase-bank behavior of LDS instructions is undocumented, requiring reverse-engineering via a custom solver (§D.2). If AMD published phase orderings in the ISA manual, bank-conflict-free swizzle design would be a straightforward lookup rather than an empirical discovery. This is a documentation debt, not a hardware limitation — but it creates a significant barrier to kernel development.

5. Larger per-CU scratchpad. MI355X has ~40% less per-processor shared memory than NVIDIA B200. This limits pipeline depth to double buffering and forces MI325X to double-buffer in the register file (consuming half the register budget for non-compute storage). Larger LDS would enable deeper software pipelines without register budget sacrifice.

6. Hardware synchronization primitives. AMD lacks mbarrier-equivalent hardware. Shared memory atomics work with negligible overhead for the 2-stage ping-pong (Tab. 2: 1281 vs 1278 TFLOPS), but hardware barriers could enable more complex multi-stage pipelines without atomic contention risk.

Concrete proposal: If CDNA5 added (a) dynamic register partitioning and (b) compositional MFMA layouts, HipKittens could replace its 8-wave ping-pong with standard wave specialization and its per-instruction swizzle infrastructure with a single unified pattern — shedding ~40% of framework complexity while potentially improving performance through deeper pipelines with dedicated producer waves.