Zyphra 发布了 ZYFA(Zyphra FlashAttention on AMD)，一套 FP8 和 BF16 prefill attention 内核，用于 AMD Instinct MI355X GPU 的 FP8 和 BF16 prefill 注意力内核。ZYFA 是 Zyphra 用于长上下文 prefill 的 FP8 注意力引擎,代码以 Apache 2.0 许可开源。在真实模型激活上，其 FP8 内核运行速度快 1.16-1.76x 于 AITER(AMD 的默认内核库)快 1.16-1.76 倍。在 head 维度 128 下,ZYFA 的 FP8 误差比 AITER 低 62-78%。
TL;DR
- 框架: 它是一个内核库。支持 FP8 e4m3 和 BF16 的因果前向(prefill)注意力,head 维度为 128 以及 MLA 形状(192/128)。
- 运行环境: 仅支持 AMD Instinct MI355X(CDNA4,gfx950),ROCm 7.2.3,PyTorch 2.12。即插即用的
flash_attn_func接口。 - 性能: 在真实激活上达到 AITER FP8 吞吐量的 1.16-1.76 倍;BF16 大致持平(1.01-1.02 倍)。
- 最佳: 在 32K tokens 下达到 2,672 TFLOP/s FP8,为 FP8 数据手册峰值的 53%(AITER 为 39%)。
- 最差: 在 Mistral-7B 激活上仅为 1.16 倍;FP8 时钟降至 1,628 MHz,而 AITER 为 2,108 MHz。
什么是 ZYFA?
ZYFA 以技术 论文 的形式发布,并附带 Zyphra 的代码。它为 Zyphra Optimization Engine(ZOE)提供注意力引擎支持,后者驱动 Zyphra Cloud。BF16 内核用 HIP 编写。FP8 内核是由 Python 生成器生成的手工调度汇编。
为什么 FP8 注意力在 MI355X 上表现不佳?
FP8 矩阵指令的运行速度是 BF16 峰值的 2 倍。但 FP8 只能加速 2 个矩阵乘积 QKᵀ 和 PV。它们之间的 softmax 在向量单元上以两种精度相同的速度运行。
对于针对128个键的32个查询,FP8矩阵运算需要1,024个周期。精确FP8 softmax需要896个向量周期。因此该内核需要近乎完美的重叠才能让矩阵单元保持忙碌。Zyphra的线程轨迹显示AITER的矩阵管线有31.7%的周期处于空闲状态。
MI355X在硬件中不跟踪矩阵指令的完成情况,因此必须靠指令顺序来创造这种重叠。
ZYFA是如何工作的?
ZYFA使用了3种技术:
精简P:用一次乘加替代一次指数运算
e4m3位模式若按整数读取,每当数值翻倍时增长8。所以2^x的码值约为8倍加一个常数。Lean P只需1次乘加和1次整数转换即可写出每个FP8概率。这将每块的softmax向量周期从896削减到384。
双波乒乓调度
2个波共享每个SIMD。一个发射矩阵指令,另一个计算概率,然后互换。采用精简P后,矩阵管线在稳态周期中有86%处于忙碌,高于精确P时的60%。
逐块缩放与舍入后的行和
每64键的概率块在舍入前先乘以2的幂。CDNA4的缩放矩阵指令可零成本移除该因子。ZYFA还除以舍入概率之和,该和每块仅需1条额外矩阵指令即可算出。
ZYFA有多快?
在持续时钟下,真实激活的FP8在16K处的加速比如下 1.16x(Mistral-7B)、1.17x(Qwen3-8B)和1.76x(DeepSeek-V2-Lite)。在32K的Gaussian基准上,head-dim-128 FP8内核达到2,672 TFLOP/s,而AITER为1,944。在MLA形状下达到2,464对1,237,比率为1.99x。
在捕获的Mistral-7B和Qwen3-8B层上,ZYFA分别提供1.94和1.91 TFLOP/J。AITER则为1.60和1.56。
为什么真实模型的收益小于基准测试收益?
MI355X的时钟取决于翻转的操作数位数量。AITER让更多小概率变为零或次正规数。这些切换的位更少,因此其时钟运行更高。在Mistral-7B层上,ZYFA运行于1,628 MHz,而AITER运行于2,108 MHz。尽管如此,ZYFA整体仍需更少的周期。Zyphra团队建议在真实激活和持续时钟下进行基准测试。
ZYFA的FP8注意力有多精确?
与fp64注意力相比,ZYFA的FP8误差在Mistral-7B上是AITER的0.22x,在Qwen3-8B上是0.38x。它与Composable Kernel和Triton相差仅几个百分点。在完整Mistral-7B中,与fp64的KL散度为:ZYFA为3.4×10⁻³,AITER为5.9×10⁻³。ZYFA的输出在7,200次启动中也完全逐位一致。
完整模型中大部分FP8误差来自对Q、K和V的量化。对Q和K做Hadamard旋转可将该输入误差大致减半。在Mistral-7B上,这会给注意力时间增加约0.1%。
ZYFA与其他注意力内核相比如何？
| 特性 | ZYFA FP8 | AITER FP8 | FlashAttention-3 | FlashAttention-4 |
|---|---|---|---|---|
| 开发方 | Zyphra | AMD | Dao-AILab等 | Dao-AILab等 |
| 目标GPU | AMD MI355X | AMD MI300X至MI355X | NVIDIA H100/H800 | NVIDIA B200（及Hopper） |
| 标称吞吐量 | 2,672 TFLOP/s FP8 | 1,944 TFLOP/s FP8（Zyphra实测） | 接近1.2 PFLOP/s FP8 | 1,613 TFLOP/s BF16 |
| 峰值占比 | 53% FP8 | 39% FP8（Zyphra实测） | 未披露（FP8） | 71% BF16 |
| Softmax方案 | Lean P，无指数运算 | 硬件指数 | 通过异步Tensor Core重叠 | 软件模拟指数 |
| FP8精度声明 | 误差比AITER低62-78% | 基线 | 误差比基线FP8低2.6倍 | 未披露 |
| 范围 | 仅因果预填充前向传播 | 预填充、解码、MLA、分页 | 前向与反向传播,FP8前向 | 前向与反向传播 |
| 许可证 | Apache 2.0 | MIT | BSD-3-Clause | BSD-3-Clause |
跨厂商的数据来自不同的输入和时钟频率。Zyphra并未测量NVIDIA GPU。
要点总结
- MI355X上的FP8注意力机制受softmax限制,而非矩阵运算限制。
- Lean P用1次乘加运算替代了指数运算。
- 真实模型上的收益为1.16-1.76倍,低于基准测试所显示的水平。
- 在头维度为128时,FP8误差比AITER降低62-78%。
- 仅支持预填充:不支持解码、不支持分页KV,仅限MI355X。
请查看 论文、 Zyphra博客文章 和 GitHub仓库。此外,欢迎在 Twitter 上关注我们,别忘了加入我们的 150k+ML SubReddit 并订阅 我们的新闻通讯等等！你在用 Telegram 吗？ 现在你也可以在 Telegram 上加入我们了。
需要与我们合作推广你的 GitHub Repo、Hugging Face Page、产品发布、Webinar 等？ 与我们联系
本文 Zyphra ZYFA: FP8 FlashAttention for AMD MI355X, Up to 1.76x Faster 首发于 MarkTechPost.