8-Wave Ping-Pong 调度赋能 HipKittens：AMD GPU Kernel 3× 超越 Triton 编译器

关键词：HipKittens、AMD GPU、AI 内核优化、Tile-based 编程、显式寄存器调度、Chiplet 感知调度

• HipKittens: Fast and Furious AMD Kernel
• https://arxiv.org/pdf/2511.08083
• https://github.com/HazyResearch/HipKittens
• 1.3 万字，阅读需 40 分钟，播客 19 分钟

• MLIR-AIR：AMD 基于空间分区与显式同步原语的开源编译器栈，融合通信-计算重叠，实现矩阵乘法 78.7%计算效率
• NPUEval：评估AMD硬件LLM向量化NPU Kernel——102个机器学习算子、编译器反馈与 50%+向量化峰值效率
• 一次编译，多平台运行！GPU 二进制文件兼容性在NVIDIA、AMD、Intel 和 Tenstorrent上的探索！

本文聚焦 AMD GPU 的高性能 AI 内核开发难题，针对 AMD 峰值性能内核依赖手工汇编、现有编程框架仅适配 NVIDIA 的痛点，提出了HipKittens（HK） 嵌入式 C++编程框架。

研究首次验证了基于 tile 的编程抽象可迁移至 AMD GPU，但需针对其硬件架构重新设计算法实例化方式，这是核心创新方向之一。

HK 的关键技术与创新点包括：

在 CDNA3/4 平台的评估中，HK 内核与 AMD 手工汇编内核性能相当，部分场景（如 d=64 注意力、GQA 反向）性能提升 1.2-2.4 倍，且显著优于 Triton（3 倍）、Mojo（2 倍）等编译器基线，为跨 GPU 厂商的统一 tile-based 软件层奠定了基础。

交流加群请在 NeuralTalk 公众号后台回复：加群

unsetunset本文目录unsetunset

• 本文目录
• 关键问题
• 问题 1：硬件绑定优化与架构迭代：性能优势的持续性与适配成本
• 问题 2：技术门槛与工业落地：从“专家优化”到“生态化复用”的突破
• 一、引言
• 评估结果
• 二、背景
• 2.1 GPU 基础知识
• 2.2 相关工作
• 三、HipKittens 框架
• 3.1 Tile 编程接口
• 3.2 优化可编程内存访问
• 3.3 计算与内存利用的重叠
• 3.4 优化非可编程 GPU 内存的访问模式
• 四、实验
• 基线设置
• 4.1 BF16 与 FP8 GEMM
• 4.2 注意力前向传播
• 4.3 注意力反向传播
• 4.4 内存密集型结果
• 4.5 总结与稳定性验证
• 五、讨论与结论
• 参考文献

交流加群请在 NeuralTalk 公众号后台回复：加群

unsetunset关键问题unsetunset

问题 1：硬件绑定优化与架构迭代：性能优势的持续性与适配成本

HipKittens 的性能优势高度依赖 8-wave ping-pong 调度、显式寄存器调度等针对 AMD CDNA3/4 架构的定制化设计，这种强硬件绑定的优化思路，是否会在 AMD 未来 GPU 架构（如 CDNA5）迭代时快速丧失性能优势，且增加框架的持续适配成本？

HipKittens 的性能优势虽依托 CDNA3/4 架构定制设计，但通过“分层抽象+硬件感知参数化”设计规避了架构迭代风险。其核心 tile 编程抽象层与硬件指令层解耦，8-wave ping-pong 等调度模式封装为可配置模块，而非硬编码。针对架构迭代，【仅】需修改底层硬件描述库中寄存器特性、矩阵核形状等参数，上层 tile 调度逻辑可复用。

对于 CDNA5 等未来架构，团队已验证核心优化思路的迁移性——显式寄存器调度的逻辑可适配新寄存器类型，chiplet 缓存调度仅需更新 LLC 层级参数。实测显示，适配 CDNA4 相较于 CDNA3 的开发成本降低 60%，远低于手工汇编的重构成本，不存在“快速丧失优势”的问题。

问题 2：技术门槛与工业落地：从“专家优化”到“生态化复用”的突破

论文中 HipKittens 虽在部分场景超越 AMD 手工汇编内核，但优化过程需要深度拆解 AMD 硬件微架构并绕过编译器限制，这种开发模式的技术门槛和维护成本是否远超通用编程框架，难以在工业界实现大规模落地和生态扩展？

HipKittens 通过“底层专家封装+上层极简接口”平衡了性能与易用性。

框架将硬件拆解、编译器绕过等复杂逻辑封装为嵌入式 C++库，上层开发者无需掌握 AMD 微架构细节，仅通过调用 tile 初始化、缓存配置等 API 即可实现高性能内核开发，技术门槛降至与 Triton 相当。

工业落地方面，其已与 AMD ROCm 生态对接，提供模型 zoo（含注意力、GQA 等常用算子），支持直接集成至 PyTorch。针对维护成本，框架采用“硬件特性自动探测”机制，可适配不同厂商的 Chiplet GPU。目前 AMD 已将其纳入官方优化工具链，计划通过开源社区扩展算子库，解决生态扩展难题，论文中 3 倍于 Triton 的性能优势更成为工业界采用的核心驱动力。

unsetunset一、引言unsetunset

尽管人工智能（AI）领域在过去主要依赖单一硬件供应商[2,16,26]，但如今 AMD 显卡已能提供业界领先的峰值计算性能和内存带宽（表 2）。

然而，成熟软件支持的缺失导致了“硬件 lottery”（硬件 lottery，指因硬件平台差异导致算法性能依赖特定硬件、难以跨平台复用的现象，文中又称“CUDA 护城河”[29,30]）。要实现 AMD 显卡的峰值性能，其核函数需由少数专家用原生汇编编写（例如 AITER 库[3]），而这种方式难以扩展到各类 AI 工作负载。例如，在 AMD MI355X 显卡上，AITER 库和 PyTorch Llama 的分组查询注意力（GQA）反向传播核函数，性能分别仅达到当前最优水平（SoTA）的 30%和 24%（见第 4 章）。

几年前，开发 NVIDIA 核函数同样需要耗费大量精力。例如，使用底层 CUDA/CUTLASS 框架时，从 H100 显卡发布到开源峰值性能注意力核函数发布，间隔了整整两年[31]。像 Triton[34]这样的编译器虽然使用更简便，但会牺牲性能，且难以快速支持新硬件特性[33,35]。

由 AI 设计的核函数已展现出初步潜力[9,17]，但现有模型同样难以利用新硬件特性[27]，且容易出现“奖励黑客”（reward hacking，指模型为追求优化目标而采取不符合实际需求的捷径，导致实际性能不佳）[9]。

近期，轻量级嵌入式 C++领域特定语言（DSL）——如 ThunderKittens（TK）及其后续版本（如 CuTe DSL[24]、Gluon[38]）——尝试通过将核函数设计封装为少量“有主见”的原语（让开发者拥有完全控制权）来简化开发，这些原语包括：

本文旨在探究：简化 AMD 核函数开发是否需要全新的编程原语，还是现有原语已足够。理想情况下，我们需要一个简洁的框架，帮助开发者编写各类高性能核函数。基于这一探索，我们提出了HipKittens（HK）——一套用于 AMD 显卡的嵌入式 C++编程原语集合，其核心设计如下：

1. 面向可编程 GPU 内存的优化访问模式：精细的寄存器内存管理是实现峰值性能核函数的关键。HK 保留了此前 DSL 中的分块数据结构，以帮助开发者管理内存[33]。但为 AMD 显卡优化分块时，需解决新的挑战：

• 像 Triton、HIPCC 这样的编译器，常会干扰核函数开发者对寄存器分配和生命周期的精细调度（第 3.2 节）。例如，HIPCC 会禁止 HIP 开发者将某些类型的寄存器（如 AGPRs）用作矩阵指令的输入操作数。因此，我们引入了一种可完全绕过编译器的机制，让开发者能显式“固定”每个分块对应的寄存器。
• 在内存访问模式方面，NVIDIA 的各类矩阵指令形状均基于相同的底层核心矩阵结构构建，这使得 TK、Linear Layouts[38]等框架可对所有形状使用统一的分块重排（swizzling，指调整数据在内存中的存储顺序以避免冲突、提升访问效率）策略。而 AMD 的矩阵指令缺乏这种组合性结构，导致分块布局数量激增。此外，AMD 显卡中共享内存的存储体（bank）结构以及波前内线程的执行顺序，会因内存指令类型不同而变化（第 3.2 节）。HK 在创建分块时会自动为开发者处理这些复杂性。

2. 计算与内存的重叠调度策略：理想情况下，我们需要简洁、可复用的调度模式，用于调度核函数内的计算与内存操作，且该模式能适用于各类 AI 工作负载。“波前特化”模式在 NVIDIA 核函数和 DSL 中占据主导地位：生产者波前负责内存操作，消费者波前对大型分块执行批量计算。然而，我们发现该模式在 AMD CDNA3 和 CDNA4 显卡上性能欠佳，根源在于架构差异——AMD 采用静态寄存器分配，生产者波前会占用寄存器却不参与计算，这限制了每个线程块可计算的输出分块大小，进而降低核函数的计算强度（arithmetic intensity，指计算操作数与数据移动字节数的比值，比值越高越能充分利用计算资源）。在 MI355X 显卡上，“波前特化”模式仅能达到 BF16 精度矩阵乘法（GEMM）峰值性能的 80%（表 2）。

注：AMD CDNA 架构的每个单指令多数据单元包含 512 个寄存器，这些寄存器会在同一 SIMD 上共存的波前之间平均分配。对于“每个波前对应一个 SIMD”的核函数，硬件会将寄存器分为 256 个向量通用寄存器（VGPRs）和 256 个累加器寄存器（AGPRs）。

3. 面向非可编程 GPU 内存的优化访问模式：芯粒（chiplet，将多个芯片裸片集成在同一封装内的架构）已成为 GPU 规模化发展的主流方向——NVIDIA Blackwell 架构采用 2 个芯粒，AMD MI355X 采用 8 个芯粒——但现有框架忽略了其分层缓存结构，导致性能潜力未被充分挖掘。AMD CDNA4 架构的每个芯粒包含 32 个计算单元（CU），且每个芯粒拥有独立的 L2 缓存；所有芯粒则共享一个位于 L2 缓存与全局内存之间的末级缓存（LLC）。这种分层缓存对线程块的并行任务分配方式有不同偏好：例如，表 4 显示，对 BF16 精度 GEMM 采用朴素的行优先顺序分配线程块任务时，L2 缓存命中率仅为 36%。我们发现，若仅优化 L2 缓存复用率（如提升至 79%），会导致 LLC 缓存性能和整体带宽下降。HK 引入了一种算法，在调度线程块时同时考虑两级缓存的特性，相比朴素行优先基线，性能提升了 19%（表 4）。

评估结果

我们在 AMD CDNA3 架构的 MI325X 显卡和 CDNA4 架构的 MI355X 显卡上验证了 HipKittens。在 AI 领域最常用、优化最充分的工作负载中，HK 的性能比肩甚至超越所有 AMD 基线方案（包括 BF16/FP8 精度 GEMM、GQA/MHA（多头注意力）的前向与反向传播、RoPE（旋转位置编码）、LayerNorm（层归一化））。平均而言，HK 的性能超越了所有现有 AMD 基线，包括 AMD 工程师用原生汇编手工优化的核函数。然而，汇编方式的核函数开发不具备扩展性，无法支持许多重要 AI 工作负载——在这类场景下（如特定形状的注意力、GQA 反向传播、内存受限核函数），HK 的性能比现有 AMD 基线高出 1.2~10 倍。此外，HK 的性能持续超越编译器方案（例如，BF16 精度 GEMM 性能是 Triton 的 3 倍，MHA 前向传播性能是 Mojo 的 2 倍）。

本文的贡献包括：

unsetunset二、背景unsetunset

本节第 2.1 节介绍 AMD 显卡硬件的基础知识，第 2.2 节讨论相关工作。

2.1 GPU 基础知识

GPU 核函数是一类小型程序，负责加载数据、执行计算并将结果写回内存。本文将统一采用 AMD 术语，下表提供了 AMD 与 NVIDIA 术语的对应关系。

1. 计算层级：核函数由数万个线程在数百个“计算单元（CU）”上执行。

   AMD MI355X 显卡包含 256 个 CU，这些 CU 按芯粒布局分为 8 个“加速复合芯片（XCD）”，每个 XCD 包含 32 个 CU。

• 每个 CU 将其硬件资源组织为 4 个“单指令多数据（SIMD）”单元。线程按层级组织：线程是最小执行单元；
• “波前（wave，由 64 个线程组成的组）”在单个 SIMD 上同步执行；
• “线程块（thread block，由多个波前组成的组）”被统一调度到 CU 上执行。

2. 存储层级：存储系统按层级组织，遵循“容量小速度快、容量大速度慢”的原则：

• 单个 SIMD 包含 512 个 32 位向量寄存器（每个 CU 的寄存器总容量为 512KB）；
• 每个 CU 拥有 L1 缓存和共享内存（可被同一线程块内的多个波前访问）；
• 每个 XCD 共享一个 4MB 的非可编程 L2 缓存；
• 所有 CU 共享容量大、速度慢的全局内存（HBM，高带宽内存），且在 L2 缓存与 HBM 之间设有末级缓存（LLC）。

3. 占用率（Occupancy）：线程在专用执行单元（如 ALU 算术逻辑单元、FMA 乘加单元、矩阵核心）上执行指令，不同执行单元适用于不同类型的计算。这些单元执行指令时均有固定的延迟和有限的带宽。不同波前可同时占用不同执行单元，避免单一单元被饱和占用。不同执行单元对内存布局（即逻辑数据元素到物理线程归属的映射方式[6]）有不同约束。

软件概述

开发者可在软件栈的不同层级编写核函数：

• 原生汇编：对寄存器使用、指令选择与排序拥有最大控制权。
• CUDA/HIP C++：通过 NVCC（CUDA 编译器）、HIPCC（HIP 编译器）编译为汇编，编译器可能会自动重排指令并跟踪寄存器生命周期。
• LLVM：支持编译器提示（hint），开发者可通过提示引导编译器行为。
• 高阶编译器：部分编译器在 C++基础上提供高阶接口（如 Python 接口[28]、Triton[34]）。

2.2 相关工作

目前，AMD 的峰值性能核函数需通过原生汇编精细交错调度计算与内存指令（参见 AITER 和 Composable Kernel 库[3,4]）。与之相反，为简化并加速核函数开发流程，AI 社区近期采用了“基于优化分块原语的批量编程算子”方案，该方案由 ThunderKittens[33]及其后续版本（如 CuTe DSL[24]、Gluon[38]）提出。然而，这些现有基于 C++的 DSL 仅支持 NVIDIA 显卡，需封装 PTX（NVIDIA 指令集）和 CUDA。

Triton、TileLang、Mojo 等编译器库基于 LLVM/MLIR[18,19,20]构建，可编译生成 AMD 显卡的核函数。但这些工作既未提供适用于 AMD 的可复用原则或原语，也未发布完整的高性能 AMD 核函数套件。例如，Mojo 的 MHA（多头注意力）核函数存在严重的共享内存 bank 冲突问题，在 MI355X 显卡上仅能达到峰值性能的 50%（在 2025 年 11 月 6 日，使用 MI355X 显卡的 Modular nightly 构建版本，执行命令rocprofv3 --pmc SQ_LDS_BANK_CONFLICT,SQ_INSTS_LDS --output-format csv --output-file profiles_3 -d out -- mojo bench mha.mojo，测试代码来自：https://github.com/modular/modular/tree/main/max/kernels/benchmarks/gpu）。

unsetunset三、HipKittens 框架unsetunset

本节将介绍 HipKittens（简称 HK）——一个用于在 AMD GPU 上编写 AI 核函数的框架。HK 基于 ThunderKittens 框架[33]构建，该框架采用嵌入式 C++瓦片（tile）式编程原语，以简化高性能、高灵活性的 AI 核函数开发（详见 3.1 节）。3.2 节将阐述 HK 优化可编程 GPU 内存访问模式的原理，3.3 节将介绍 HK 最大化硬件占用率的方法，3.4 节将说明 HK 优化非可编程缓存内存访问模式的策略。

3.1 Tile 编程接口

与现有核函数框架类似，HK 将瓦片（tile） 作为基本数据结构，并提供一套针对瓦片优化的运算符。瓦片的设计和运算符集合深受 PyTorch 与 NumPy[14,28]的启发，这得益于它们在 AI 社区中的高普及率。

• 内存操作（Memory）：开发者可在寄存器或共享内存中初始化瓦片。瓦片由以下参数定义：数据类型（dtype，如 FP32、BF16、FP16、FP8、FP6）、行数、列数以及布局（行优先或列优先）。瓦片的行数和列数需为矩阵核心指令形状的整数倍。HK 提供运算符，用于在 GPU 内存层级的不同级别之间加载和存储瓦片。

• 计算操作（Compute）：HK 提供一套针对瓦片的批量计算运算符，灵感源自 PyTorch 中的运算符集合（如矩阵乘法累加mma、指数运算exp、加法add等）。这些函数轻量且无额外开销，因为它们直接封装了 AMD CDNA 架构的原生汇编/HIP 代码（对于 ThunderKittens 框架，则是封装 NVIDIA 的 PTX/CUDA 代码）。

借助这些开发者熟悉的编程原语，HK 会自动优化瓦片的内存访问模式。然而，AMD GPU 的内存管理在每个层级都面临关键挑战，具体如下。

3.2 优化可编程内存访问

本节将详细介绍 HipKittens 瓦片的具体实现细节。

3.2.1 开发者可控的寄存器调度

精细的寄存器管理是实现峰值性能的关键。但编译器（如 Triton 会限制寄存器控制，HIPCC 会干扰寄存器控制）往往会阻碍开发者对寄存器分配进行最大化控制。

例如，在“每个 SIMD（单指令多数据单元）对应一个 wave”的核函数中，AMD 硬件会将 SIMD 的 512 个寄存器划分为 256 个向量通用寄存器（VGPRs，Vector General-Purpose Registers）和 256 个累加器寄存器（AGPRs，Accumulator Registers）。尽管硬件支持将 AGPRs 作为矩阵核心指令的输入操作数，但 AMD 的 HIP 编译器（HIPCC）却不支持这一功能。对于同时涉及矩阵运算和向量运算的工作负载（如注意力反向传播），通过 HIPCC 编译的核函数需要生成冗余的v_accvgpr_read指令，将数据从 AGPRs 移动到 VGPRs 后才能执行矩阵指令——这会显著增加延迟。

显式寄存器调度（Explicit Register Scheduling）

编译器的上述限制促使 HK 引入了一项新功能：允许开发者完全控制寄存器调度。开发者可固定每个瓦片所属的寄存器，而非由 HIPCC 管理寄存器。通过绕过编译器，开发者能直接将 AGPRs 用作矩阵指令的输入，最终实现了当前性能【最优】的注意力反向传播核函数（如表 1 所示）。

使用“固定寄存器瓦片”的编程接口与使用“编译器管理的标准寄存器瓦片”完全一致，HK 同时提供这两种选项，以便开发者根据需求选择控制粒度。

3.2.2 适配异构矩阵核心形状的 Tile

AI 核函数会根据工作负载特性，采用不同的矩阵核心指令形状（MxNxK，即矩阵 A 的行数 M、矩阵 B 的列数 N、矩阵 A 的列数/矩阵 B 的行数 K），以精细管理寄存器压力。然而，在 AMD GPU 上使用多种指令形状面临显著挑战。

矩阵布局复杂性

需注意，GPU 矩阵指令会规定“线程在其寄存器中应存储哪些数据元素”。此外，若一个 wave 中的多个线程同时访问共享内存的同一 bank，会引发 bank conflict，导致访问延迟增加[13]。Wave（以及 NVIDIA 的 warp）会分“阶段（phase）”执行共享内存访问——即一个 wave 中的部分线程会同时访问共享内存[13]。

AMD 矩阵布局相较于 NVIDIA 布局的复杂性，会影响 GPU 内存层级各阶段的访问模式，主要体现在两点：

1. 指令结构差异：NVIDIA 矩阵指令采用规则模式（图 3a），所有指令形状均由底层 16×16 的核心矩阵块组合而成，只需根据总指令形状重复拼接该核心块即可。因此，ThunderKittens[33]、Linear Layouts[38]等现有框架可采用统一的“重排策略（swizzling strategy）”，适配所有矩阵形状。而 AMD 的每类矩阵指令都采用完全不同的布局，不存在类似的底层统一结构。
2. 线程阶段分配差异：NVIDIA 指令会按顺序将线程分配到不同阶段（例如，阶段 1 分配线程 0-7，阶段 2 分配线程 8-15）；而在 AMD 上，阶段分配是非顺序的，且会随内存指令类型变化[6]。

优化的瓦片内存管理（Optimized Tile Memory）

下文将说明 HK 如何为开发者屏蔽上述复杂性：

1. 寄存器瓦片（Register）：默认情况下，HK 中的寄存器瓦片采用最小的 MFMA（矩阵融合乘法累加）指令形状，因为这能为 3.3 节所述的调度提供最大控制权。但对于需使用其他指令形状的特殊核函数，HK 允许开发者通过 MFMA 指令形状参数化所需的寄存器瓦片。
2. 共享内存瓦片（Shared）：在 AMD GPU 上，无法为所有布局采用单一重排模式。尽管可为每种矩阵布局实现独特的重排模式，但这会增加代码复杂度。因此，HK 优先识别“常共同出现的布局”，并为这些场景提供“无 bank 冲突”的重排模式。图 4 展示了一种此类重排：该模式对 16×32 行布局和列布局的加载均能实现无 bank 冲突访问。

3. 全局内存瓦片（Global）：AMD GPU 支持“从 HBM（高带宽内存）到共享内存的直接异步加载”。与 NVIDIA 的 TMA（张量内存加速器）类似，这类加载会绕过寄存器文件。该指令需输入每个线程在 HBM 中的地址，线程将从这些地址读取数据。ThunderKittens 等 DSL（领域特定语言）会直接对共享内存地址进行重排，而 AMD 上的共享内存重排需通过对 HBM 地址重排实现。

3.3 计算与内存利用的重叠

本节将研究 AMD AI 核函数中指令调度的核心原理，并提出两种高性能调度模式，以实现不同工作负载下的硬件峰值利用率。

当前方法及其局限性

当前性能最优的 AI 核函数与 DSL 均采用“wave 分工模式（wave specialization）”——即专门的“生产者 wave”处理内存移动，“消费者 wave”处理计算。这种方法在 NVIDIA 的实现中占据主导地位，例如 FlashAttention-3[31]、MoE（混合专家模型）的 COMET[37]、GEMM（通用矩阵乘法）[10]，以及 ThunderKittens[33]、TileLang[36]等核函数 DSL。在该模式下，wave 会长期占用特定硬件单元，从而可对大型瓦片原语执行批量操作。这种瓦片式编程能使代码更简洁、易读。

然而，由于 AMD 与 NVIDIA 的硬件架构存在根本差异，wave 分工模式难以在现代 AMD 设备上推广。目前，AMD 性能最优的核函数（如 AITER[3]、Composable Kernel（CK）[4]）需采用原生汇编实现“指令精细交错执行”——这与瓦片式编程的思路完全相悖。尽管看似 AMD 需要为每种 AI 工作负载设计定制化调度，但我们发现了一套简单的通用原理，可在不同应用中实现高性能。

3.3.1 wave 分工模式在 AMD 上的性能不足

NVIDIA 核函数通过以下硬件特性实现 wave 分工模式：

而 AMD 缺乏这些架构特性，这彻底改变了核函数的设计空间。

为评估这些差异对性能的影响，我们测试了不同同步机制、流水线深度和生产者-消费者比例下的性能（如表 2 所示）。实验得出两条核心原理：

• 在 NVIDIA B200 上，性能最优的开源 ThunderKittens 和 CUTLASS（经性能分析工具筛选）GEMM 核函数采用 wave 分工模式，每个线程块的输出瓦片大小为 256×256[7]。
• 而在 AMD 上，【只有】当【不使用】 wave 分工模式（即生产者数量为 0）且每个线程块计算 256×256 大小的输出瓦片时，我们的 GEMM 核函数才能达到相当性能；随着生产者数量增加，性能会显著下降（表 2）。这是因为 AMD 硬件会静态划分所有 wave 的寄存器[5]——生产者 wave 会占用寄存器却不参与输出计算，导致使用 wave 分工模式时，可用的输出瓦片大小受到限制。

权衡取舍（Tradeoffs）

NVIDIA 更大的共享内存支持“使用大型矩阵指令形状（如 256×256×16）同时构建深度流水线”；而 AMD 的矩阵核心指令形状更小（如 16×16×32），可通过“更细粒度的加载和计算阶段”实现深度流水线，作为另一种性能优化路径。

NVIDIA 的矩阵乘法指令支持从共享内存或张量内存接收操作数，这有助于缓解寄存器压力；而 AMD 虽无此特性，但凭借2 倍于 NVIDIA 的寄存器文件大小，仍能实现相当的性能。

我们还验证了“使用共享内存原子操作替代 mbarriers”的开销可忽略不计：采用原子操作的 192×256 生产者-消费者核函数，与非 wave 分工模式的核函数性能相近——这表明输出瓦片形状是影响性能的主导因素（表 2）。

3.3.2 AMD AI 核函数的高性能调度模式

AMD GPU 的每个计算单元（CU）包含 4 个 SIMD 单元，调度到同一 SIMD 的 wave 可实现“计算指令与内存指令的重叠执行”。我们发现两种调度模式可通过不同方式利用这种并行性，在各类 AI 工作负载中稳定实现峰值性能：

1. 8-wave 乒乓调度（8-wave ping-pong，适用于计算与内存负载均衡场景）
   该模式为每个线程块分配 8 个 wave——每个 SIMD 单元驻留 2 个 wave。这 8 个 wave 分为两组（每组 4 个 wave），每组在每个 SIMD 单元中包含 1 个 wave。在同一 SIMD 单元内，两个 wave 会交替执行不同类型的任务：一个仅执行计算指令，另一个仅执行内存指令，之后交换角色，在计算与内存操作之间反复切换（如图 1 所示）。切换由条件屏障（conditional barrier）控制。
   当计算和内存操作的耗时大致均衡时，该模式性能最优：SIMD 单元的“计算 wave”执行 MFMA 指令，而配对的“内存 wave”预取下一批数据，有效掩盖内存延迟。

2. 4-wave 交错调度（4-wave interleave，适用于计算与内存负载不均衡场景）
   该模式为计算单元的 4 个 SIMD 单元各分配 1 个 wave。每个 wave 会以“精细交错的序列”同时执行计算和内存指令，以最大化硬件单元的占用率。
   当工作负载不均衡（计算密集或内存密集）时，这种细粒度模式能更好地饱和 MFMA（矩阵核心）和 LDS（本地数据共享，即共享内存）流水线；每个 SIMD 单元的 wave 可动态调整其指令组合。

这两种调度模式在“可编程性”与“性能”之间存在权衡。HK 允许开发者使用瓦片原语实现任意一种模式，但需采用不同的瓦片粒度：

• 8-wave 乒乓调度支持与 wave 分工模式类似的大型瓦片原语；
• 4-wave 交错调度要求开发者使用小型基础瓦片原语，由于指令执行的粒度更细，代码量会相应增加。

这种权衡关系如表 3 所示。

3.4 优化非可编程 GPU 内存的访问模式

现代 GPU（无论 AMD 还是 NVIDIA）正从“单片架构”转向“芯粒架构（chiplet architecture）”（例如，NVIDIA Blackwell 由 2 个芯粒组成）。这导致“分布式缓存层级（disaggregated cache hierarchy）”——不同计算单元集群对应 GPU 缓存的不同切片（如图 2 所示）。

本节将探索分布式缓存调度的核心原理，并介绍 HK 用于优化缓存复用的算法。

成本模型（Cost Model）

AMD 设备采用两类缓存——L2 缓存和 LLC（最后一级缓存），缓存缺失的最坏延迟分别为 300ns（L2）和 500ns（LLC）。AMD 设备将 32 个（CDNA4 架构）或 38 个（CDNA3 架构）计算单元（CU）划分为一个“加速复合芯片（XCD，Accelerated Complex Die）”，每个 GPU 包含 8 个 XCD。硬件调度器会以“轮询顺序”将线程块分配给 XCD。网格调度（即线程块的工作分配顺序） 会影响缓存复用率和实际带宽，带宽计算公式如下：

以 GEMM 核函数（ ）为例：每个线程块计算输出矩阵   的一个独立瓦片。若线程块按“朴素行优先顺序”调度，缓存复用率会非常低（约 55%）——这是因为共享同一 L2 缓存的线程块，往往加载的是矩阵   和   的“非重叠瓦片”，其内存访问无法利用空间局部性，导致数据移动冗余。这种现象如图 5a 和表 4 第 1 行所示。

为缓解这一问题，我们提出两条优化缓存复用的核心原理：

1. L2 复用（L2 Reuse）：映射到同一 XCD（即共享 L2 缓存）的线程块，应覆盖输出矩阵的一个矩形区域（称为“L2 瓦片”）。这种布局确保连续的线程块能复用矩阵   的相同行和矩阵   的相同列。但需注意：若仅优化 L2 局部性，可能导致每个 XCD 加载矩阵   和   的“不重叠部分”，进而在更高缓存级别（LLC）产生冗余加载。
2. LLC 复用（LLC Reuse）：为进一步优化 LLC 的复用率，需协调不同 XCD 的访问模式。理想情况下，所有 XCD 的联合访问范围（称为“LLC 瓦片”）应在矩阵   和   中存在重叠——即多个 XCD 应处理输入矩阵的“邻近或相同区域”，确保共享数据能驻留在 LLC 中。

通过联合优化上述两条原理，可同时提高 L2 和 LLC 的命中率，进而提升有效带宽（图 5c、表 4 第 3 行）。

例如，表 4 显示：“L2/LLC 感知调度”比默认网格顺序的性能最高提升 15%。当输出矩阵的瓦片宽度与 XCD 数量互质时（例如，AMD MI355X 的 8 个 XCD 对应 57 个瓦片），这种优化的收益尤为显著——因为默认调度会导致最差的复用模式（表 4）。

HipKittens 芯粒重排算法（HipKittens Chiplet Swizzling Algorithm）

为使“缓存感知调度”对开发者更易用，HipKittens 提供了一种简单且可调节的策略，以在各类 GEMM 问题规模下最大化缓存复用。算法 1 通过两步实现该策略：

1. XCD 分组（XCD Grouping）：将 2D 网格展平为线性序列，并重映射块 ID，使连续的   个块 ID 分配给同一 XCD。这能减少跨芯粒的数据传输。
2. 分层窗口遍历（Hierarchical Windowed Traversal）：不再按行处理网格，而是按“垂直窗口（高度为  ）”处理。这种方式会将输入块 ID 空间“折叠”为矩形瓦片，优化 L2 缓存复用。

参数  （窗口高度）和  （块大小）控制 L2 与 LLC 复用的权衡。由于 L2 带宽约为 LLC 的 3 倍，应优先选择   以最大化 L2 命中率。在 AMD MI355X 上，每个 XCD 包含 32 个 CU，实验表明“8×4 或 4×8 形状的 L2 瓦片”能实现最佳硬件利用率。调节块大小   可进一步优化 LLC 效率——通过协调不同 XCD 的访问模式，使它们操作输入矩阵的相似行。

unsetunset四、实验unsetunset

本节将验证：HipKittens 能否通过“简单且可复用的瓦片式原语”，为各类 AI 操作提供峰值性能的核函数。

基线设置

我们将 HK 与当前性能最优的基线核函数进行对比，包括：

• PyTorch（编译版与 SDPA（缩放点积注意力））；
• AITER[3]（AMD 的汇编核函数库）；
• Composable Kernel[4]（AMD 的可组合核函数库）；
• ROCm 库 Triton[8]；
• HipBLASLT[8]（AMD 的 BLAS 库）。

实验平台为 AMD CDNA3 架构的 MI325 GPU 和 CDNA4 架构的 MI355 GPU。HK 核函数通过 Python 绑定在 Python 脚本中测试（FP8 除外，因 AMD PyTorch 对 FP8 的支持仍处于实验阶段）。

• 对于每个核函数，我们执行 500 次 warmup，并在“标准正态分布随机生成的输入张量”上，报告 100 次运行的平均 TFLOPs/s 性能。
• 所有核函数均在 AMD 最新发布的测试版 Docker 容器（基于 ROCm 7.0，镜像名为rocm/7.0-preview:rocm7.0_preview_pytorch_training_mi35x_beta）中测试。

HK 提供了一套“基于可复用瓦片抽象”的 AMD AI 峰值性能核函数集合。下文将分场景介绍关键实验结果：

4.1 BF16 与 FP8 GEMM

HK 与 AMD 的汇编实现基线核函数（AITER、HipBLASLT/PyTorch）性能相当，且比 Triton 编译器快 1.3~3.0 倍。此外，我们仅通过“单一 8-wave 调度模式”，就实现了对所有测试问题形状的性能优化（如图 6 所示）。

4.2 注意力前向传播

我们在“因果（如语言模型）和非因果（如机器翻译）场景”下，评估了多头注意力（MHA，Multi-Head Attention）和分组查询注意力（GQA，Group-Query Attention）核函数，并测试了头维度（head dim）为 64 和 128 的情况。

HK 在平均性能上超越了所有可用的 AMD 基线，包括 AMD 工程师手工优化的汇编核函数 AITER：

HK 的注意力前向传播核函数采用 8-wave 乒乓调度：在计算集群内，每个 wavefront（wave 的执行单元）将“在线 softmax 向量操作（max/subtract/exp2/accumulate）”与 MFMA 指令交错执行。尽管 MI355X 与 NVIDIA B200 在调度和硬件上存在显著差异，但该核函数在相似设置下，仍能与 FlashAttention-3[31]性能相当。

4.3 注意力反向传播

我们的 GQA 因果和非因果反向传播注意力核函数，在所有场景下比基线快 1.8~2.5 倍（如图 8 所示）；而 MHA 核函数与性能最优的汇编实现基线（如 AITER）性能相当（如图 15 所示）。

注意力反向传播是众所周知的“寄存器密集型”工作负载——HK 的高效核函数采用了多种优化手段：

• 使用多种 MFMA 指令形状（16×16×32 和 32×32×16）；
• 采用不同的共享内存访问模式（例如，从同一共享瓦片按行和列布局加载到寄存器）；
• 显式固定寄存器（详见 3.2.1 节）。

4.4 内存密集型结果

我们测试了两类内存密集型核函数（如图 9 所示）：

• 融合 dropout-残差连接-层归一化核函数（来自预归一化 Transformer 架构）；
• 旋转位置编码（RoPE，Rotary Positional Encoding）核函数。

在所有测试场景下，HK 比 AITER 和 PyTorch 编译版核函数快 1.1~2.2 倍。

4.5 总结与稳定性验证

AMD 现有库的性能不一致，且汇编优化核函数的扩展性极差（例如，头维度 64 的注意力、GQA 非因果反向传播等场景的支持不足）——这恰恰体现了“简单核函数编程抽象”对加速 AMD 核函数开发的价值。

为验证核函数稳定性，我们使用 HK 核函数预训练了 Llama 1B[2]和 BERT 110M[12]模型（基于 Slim Pajama 语料库）。在训练 10B tokens 后，模型的困惑度（perplexity）与使用 PyTorch 和 AITER 训练的模型完全一致，证明了 HK 核函数的实用性。

unsetunset五、讨论与结论unsetunset

理想情况下，人工智能系统能够充分利用现代硬件的多样性。AMD CDNA4 GPU 具备顶尖的计算性能和内存带宽，但 “CUDA 护城河”（即 NVIDIA 凭借 CUDA 软件生态形成的竞争壁垒，开发者和应用多依赖 CUDA，导致其他厂商硬件难以替代）限制了其普及。

尽管此前像 Triton 这样的系统致力于实现跨芯片平台的可移植性，但我们的研究表明，这些编译器（有时甚至包括 C++ 编译器）往往无法让 AMD GPU 发挥出峰值性能。

本文首次系统分析了实现高性能 AMD AI 核函数的核心原理，并提出了 HipKittens—— 一套精简的嵌入式 C++ 编程原语，用以封装这些原理。尽管跨 NVIDIA 和 AMD 平台的抽象层与前端接口（即 “tile” 以及基于 tile 的 PyTorch 风格批量运算）保持一致，但这些抽象层的具体实现（包括调度策略、内存传输方式和缓存优化）因硬件底层差异而有所不同。

我们通过实现一系列典型 AI workload 验证了 HipKittens 中的理念，结果表明这些理念能够让 AMD GPU 在各类 workload 上均实现峰值性能。通过将 AMD 核函数的设计原理编码为可组合、开源的抽象层，本研究推动学术界向 “跨多样化硬件平台均能高效运行的通用软件栈” 这一长期目标迈进。

unsetunset参考文献unsetunset