MoonMath AI 开源 HIP 注意力内核，AMD MI300X 性能全面超越 AITER v3

一句话看懂

MoonMath AI 开源了一个 HIP 编写的 bf16 前向注意力内核，在 AMD MI300X 上全面超越官方 AITER v3，几何平均加速 1.08×-1.18×，最高 1.26×。

详细发生了什么

MoonMath AI 团队发布了一个针对 AMD MI300X GPU 的 bf16 前向注意力内核，完全用 HIP 编写（非手写汇编），采用 MIT 开源许可。该内核在测试的所有形状和舍入模式下均击败了 AMD 官方优化的 AITER v3 内核。

核心技巧是“单指令内联汇编包装器”：将每条指令封装在 __device__ __forceinline__ 函数中，通过 asm volatile 约束让编译器负责寄存器分配，同时保留对操作码的精确控制。例如 v_mfma_f32_16x16x16_bf16 指令通过 "+v"(c) 约束将累加器输入输出绑定到同一 VGPR，避免冗余拷贝。

架构上，每个计算单元运行 8 个波（分两组，每组 4 个），两组交替执行矩阵乘和 softmax/加载，确保矩阵核心始终忙碌。内存布局是关键：K 流经 HBM 进入 LDS（双缓冲，8 波共享），V 常驻 L1 缓存，Q 和累加器在寄存器中。选择 16×16×16 MFMA 形状以降低 VGPR 压力。

基准测试在 MI300X 上进行，头维度 128，测试了多种形状和三种舍入模式（RTNE、RTNA、RTZ）。几何平均加速比：vs AITER 为 1.18×/1.15×/1.08×，vs Modular MAX 为 1.44×-1.49×。具体形状如 (2,24,16384,128) RTNE 达 1.26×。

实际应用：团队向 SGLang 提交了 PR，在 Wan2.1-T2V-1.3B-Diffusers 视频扩散模型中将注意力替换为 MoonMath 内核，端到端生成速度提升 1.23×，无质量退化。

中文圈视角

对国内 AMD 用户和 AI 开发者来说，这是个好消息。目前国产 GPU（如华为昇腾、寒武纪）生态仍以 CUDA 兼容为主，AMD ROCm 生态在国内相对小众，但 MI300X 已通过部分云厂商（如 HotAisle）提供裸金属访问。该内核的开源意味着：

平替 NVIDIA H100 的潜力：MI300X 在显存带宽（5.2 TB/s）和容量（192 GB HBM3）上优于 H100，但软件生态是短板。MoonMath 的内核证明，通过精细优化，AMD 硬件在注意力计算上可以超越官方实现，缩小与 NVIDIA 的差距。
对国内模型推理/训练的影响：国内大模型厂商（如 DeepSeek、智谱、百川）若使用 AMD 硬件，可直接采用此内核加速推理。特别是长序列场景（如 131K 上下文），加速比达 1.16×。但需注意：内核仅支持 bf16 前向、头维度 128、无因果掩码/GQA/varlen，适用场景有限。
与国产 GPU 的对比：华为昇腾 910B 的注意力优化主要依赖 CCE 自定义算子，性能数据不透明。MoonMath 的开源实现为国产 GPU 厂商提供了参考：通过内存布局优化（K 在 LDS、V 在 L1）和流水线设计，可在类似架构上复现加速。
中文社区盲点：国内对 AMD ROCm 的关注度远低于 CUDA，但 MI300X 的性价比（相比 H100 约 1/3 价格）可能吸引中小团队。此内核的开源降低了 AMD 生态的使用门槛，但需要用户自行编译 ROCm 环境，且仅支持 gfx942 架构。

几条值得记住的细节

内核仅支持 bf16 前向注意力，头维度固定 128，无因果掩码、GQA 或 varlen batching。
三种舍入模式（RTNE/RTNA/RTZ）均与 AITER 逐位一致，有限输出误差在 1 bf16 ULP 内。
内存布局：K 在 LDS（32 KiB 双缓冲）、V 在 L1（常驻）、Q 和累加器在 VGPR。
实际应用：SGLang 中 Wan2.1 视频扩散加速 1.23×，无质量退化。
安装简单：pip install moonmath_attention，API 兼容 PyTorch，支持 BSHD/BHSD 布局。

一句话总结

如果你在用 AMD MI300X 跑 Transformer 模型，这个开源内核能免费获得 8%-26% 的注意力加速，且无需改代码。