MoonMath AI 开源 AMD MI300X 的 HIP 注意力内核，在各项指标和舍入模式下均超越 AITER v3

MoonMath AI 团队近日发布了针对 AMD MI300X GPU 的 bf16 前向注意力内核。该内核使用 HIP 编写（而非手写汇编），以 MIT 许可证开源发布。MoonMath.ai 团队声称，该内核在测试的所有形状上均超越了 AMD 自己的优化内核 AITER v3。裸机访问由 AMD 云服务提供商 HotAisle 提供。

注意力机制是每个 Transformer 中的融合 softmax(QKᵀ/√d)·V 操作。MI300X 是 AMD 的 CDNA3 数据中心 GPU，目标 ISA 为 gfx942，该内核仅在此硬件上运行。

内核是一款直接在 GPU 众多核心上运行的小程序，专门用于尽可能快地执行注意力运算。该内核仅在 MI300X 上计算 bf16 前向注意力，支持 BSHD 或 BHSD 布局，无需转置，头维度固定为 128，支持任意序列长度包括交叉注意力。

目前存在一些限制：不支持因果掩码、GQA 或变长批处理。输出为 bf16，仅在 gfx942 硬件上运行。数值控制严格：所有三种舍入模式均与 AITER 的逐模式舍入规则匹配，每个有限输出与 AITER 的偏差在 1 bf16 ULP 以内，NaN 和 Inf 处理完全一致，结果具有确定性。

核心技巧：单指令汇编包装器 内核的核心理念规避了常见的困境：编译器内联函数保持代码整洁但可能允许编译器重排序或重命名操作数；而原始内联汇编提供控制但需手动管理寄存器和地址。MoonMath 团队将单条指令包装在 device forceinline 函数中，使用扩展 asm 约束描述操作数。团队选择操作码，编译器仍分配寄存器并跟踪数据流。

例如，asm mfma 函数中，"+v"(c) 约束将累加器输入和输出绑定到同一 VGPR，无需发出复制指令。这使得内核接近普通的 HIP 代码，但仍能逐条指令控制机器。

架构：八波流水线，两组，两个屏障 一个 CDNA3 计算单元有四个 SIMD 单元。传统方法是四个波前，但 MoonMath 在每个块中使用八个波前，分为两组，每组四个。两组执行相同的 Q*K、softmax、O += P*V 序列，但相位偏移。当一组饱和矩阵核心时，另一组执行 softmax 并发出加载指令，然后交换，确保矩阵核心永不空闲。每次迭代有两个 s_barrier：一个在相位交接处，一个在迭代边界。其余同步由逐计数器等待处理。这类似于 FlashAttention-3 的矩阵乘法和 softmax 交替模式，但未采用其生产者和消费者波前拆分。在 CDNA3 上，所有内存移动已是异步的，因此无需专用的生产者波前。

数据存放位置及 16×16×16 的选择 大多数加速来自内存布局：K 从 HBM 流入 LDS，双缓冲，由所有八个波前共享；V 保持在 L1 中，每次 PV 矩阵乘法时读取；Q 和累加器位于寄存器中。团队选择了 16×16×16 MFMA 而非 32×32×8，两者吞吐量相同，但更小的瓦片每个通道累加 4 个 fp32 元素（而非 16 个），降低了累加器压力，为更深的预取和第三个 Q 瓦片留出空间。

后续两个优化进一步缩小差距：第三个 Q 瓦片（3Q）提高了每个加载的 K 和 V 瓦片的数据复用率；类似 Flash-Decoding 的尾部 KV 拆分解决了 MI300X 的 304 个 CU 上剩余分数的循环问题。这些优化层层递进：将 V 移到 L1 释放了 LDS，然后第三个 Q 瓦片填充了它。

基准测试 测试在 MI300X 上以 bf16 格式执行，头维度 128。每个形状在三种舍入模式下测量：RTNE（舍入到最近偶数）、RTNA（舍入到最近，远离零）、RTZ（向零截断）。结果显示 MoonMath 的内核全面领先：几何平均相对 AITER 为 1.18×（RTNE）、1.15×（RTNA）、1.08×（RTZ）；相对 Modular MAX 的几何平均加速比为 1.44× 到 1.49×，单形状最高达 1.59×。

实际用例 该内核通过 pip 安装，提供简洁的 API，在调用者流上启动，从而与大流水线重叠。具体用例包括视频扩散：团队为 SGLang 扩散添加了 LiteAttention 支持并提交了 PR，在 Wan2.1-T2V-1.3B-Diffusers 上，将注意力从 AITER 切换到 liteattention_rocm，端到端生成性能提升了 1.23×，且无可见质量退化。BSHD 布局直接适用于扩散张量，交叉注意力支持任意 KV 长度且无需填充。

关键要点

• 内核为 MI300X 的 bf16 前向注意力，HIP 编写，MIT 许可。
• 在每种形状和舍入模式下优于 AITER v3，几何平均 1.18×/1.15×/1.08×。
• 单指令汇编包装器在保留操作码控制的同时让编译器分配寄存器。
• 内存布局是主要加速来源：K 在 LDS，V 在 L1，Q 在寄存器。
• 实际 SGLang PR 加速 Wan2.1 视频扩散 1.23× 且无质量损失。

更多技术细节请查看原文。欢迎关注我们的 Twitter、加入 150k+ ML SubReddit 或订阅通讯。