Source figure
来自一手来源:repository 查看原始材料 ↗
Source figure
来自一手来源:repository 查看原始材料 ↗
Source figure
来自一手来源:repository 查看原始材料 ↗

先看证据

1.16-1.76x AITER FP8 吞吐提升(真实激活)关键事实
62-78% ZYFA FP8 误差低于 AITER关键事实
2,672 TFLOP/s ZYFA FP8 峰值吞吐(32K token, 高斯输入)关键事实
53% ZYFA FP8 达到 MI355X FP8 数据手册峰值关键事实
Lean P 优化 softmax 计算量896 vs 384 周期

把机制串起来

输入先把素材压到可处理的规模

FlashAttention通过优化计算顺序和内存访问,大幅提升了注意力机制的计算效率。FP8是一种低精度浮点数格式,能显著减少模型存储和加速计算,但精度损失是主要挑战。

机制再把计算集中到关键步骤

ZYFA通过Lean P近似计算softmax,减少向量单元计算量;采用双wave乒乓调度,使矩阵计算与概率计算重叠;并对概率进行分块缩放与归一化,以在FP8精度下提高计算效率和数值稳定性,同时缓解了AMD硬件对矩阵指令完成追踪不足的问题。

结果最后落到可验证的工作结果

['适用于AMD Instinct MI355X GPU,进行模型推理预填充阶段的注意力计算。', '代码已开源,可集成至PyTorch工作流,支持`torch.compile`。', 'FP8精度提升显著,但需注意其在数值敏感任务上的误差表现。']

AMD MI355X FP8 注意力性能的意外突破

近日,Zyphra 发布了 ZYFA(Zyphra FlashAttention on AMD),一套针对 AMD Instinct MI355X GPU 的 FP8 和 BF16 预填充(prefill)注意力内核库。该开源项目旨在解决 MI355X 在使用 FP8 精度进行大语言模型(LLM)推理预填充时,由于架构特性导致的性能未能充分释放的问题。ZYFA 专注于在长上下文场景下,优化注意力机制中的矩阵乘法与 Softmax 计算,以期更高效地利用硬件资源。

这套内核库的发布引发了广泛关注,因为它直接挑战了关于 MI355X FP8 性能受限的普遍认知。Zyphra 宣称,在真实模型激活数据上,ZYFA 的 FP8 内核比 AMD 默认的 AITER 库快 1.16 倍至 1.76 倍,同时在 head dimension 为 128 时,其 FP8 误差比 AITER 低 62% 至 78%。这一显著的性能与精度提升,预示着 AMD MI355X 在 LLM 加速领域具备更深层的优化潜力,对于依赖该硬件栈进行大规模模型部署的技术负责人而言,ZYFA 提供了一个不可忽视的性能新基准。

核心症结:FP8 矩阵运算与 Softmax 向量单元的速度错配

AMD Instinct MI355X 硬件的一个关键特性是其 FP8 矩阵指令峰值吞吐量理论上是 BF16 的两倍。这意味着在执行矩阵乘法等核心计算时,FP8 能够带来显著的速度提升。然而,注意力机制不仅仅包含矩阵乘法(QKV),还包括一个关键的 Softmax 归一化步骤。问题在于,Softmax 运算主要在向量单元上执行,而向量单元的计算速度并不会因 FP8 精度的使用而同比例加速。

这种速度上的不匹配导致了一个性能瓶颈。Zyphra 的分析指出,在处理例如 32 个查询与 128 个键的注意力块时,FP8 矩阵乘法仅需约 1024 个周期,但精确的 FP8 Softmax 却需要多达 896 个向量周期。这意味着即使矩阵单元能够高速完成任务,也必须等待向量单元的 Softmax 计算,导致矩阵管线长时间空闲。Zyphra 的线程追踪数据显示,AMD 默认的 AITER 库在处理注意力时,矩阵管线有高达 31.7% 的周期处于空闲状态,使得其 FP8 注意力内核的实际吞吐量仅达到 FP8 理论峰值的 39%。这是 ZYFA 尝试解决的核心挑战。

ZYFA 的技术解法:巧用编码、并行与精细量化

为解决上述瓶颈,ZYFA 采用了三项核心技术。首先是 "Lean P" 策略,它利用 FP8 e4m3 浮点数编码的特性,通过一次乘加和一次整数转换来近似计算注意力概率。传统的指数运算需要更多的向量单元周期,而 Lean P 将每个块的 Softmax 向量周期从 896 减少到 384,大幅降低了 Softmax 的计算负担,使其与矩阵运算的速度更匹配。

其次,ZYFA 引入了"双 wave 乒乓调度"机制。MI355X 不会在硬件层面跟踪矩阵指令的完成状态,因此 Zyphra 通过 Python 生成器输出手工排程的汇编代码,让两组线程波(wave)交替执行。一组线程负责发出矩阵指令,而另一组则同步计算概率,然后它们交换角色。这种精密的调度策略将矩阵管线的忙碌比例从传统精确 Softmax 时的 60% 提升至 Lean P 策略下的 86%,显著减少了矩阵单元的空闲时间。

最后,ZYFA 实施了"分块缩放与舍入行和"技术。它在将每个 64-key 块的概率舍入为 FP8 之前,先将其按 2 的幂进行缩放,并利用 CDNA4 架构的 scaled matrix instruction 在不增加成本的情况下抵消此缩放因子。此外,ZYFA 在归一化时使用舍入后概率的行和,而非未舍入前的总和。这解决了当小概率值被舍入为零时,分母仍计入这些值,从而导致结果偏差的问题,进一步提升了 FP8 注意力计算的数值稳定性与精度。

性能实测:真实模型激活与基准测试的数字证据

ZYFA 的创新技术带来了可衡量的性能提升。在真实模型激活(16K 序列长度)的单层内核测试中,ZYFA 的 FP8 注意力内核相对于 AITER,在 Mistral-7B 上实现了 1.16 倍的吞吐量提升,Qwen3-8B 上为 1.17 倍,而在 DeepSeek-V2-Lite 上更是达到了 1.76 倍。相比之下,BF16 内核的性能提升相对温和,约为 1.01 至 1.02 倍,进一步凸显了 ZYFA 在 FP8 优化上的专注与成效。

除了真实激活数据,在高斯随机输入、32K 序列长度的内核级基准测试中,ZYFA 在 head dimension 为 128 的 FP8 吞吐量达到了 2,672 TFLOP/s,是其 FP8 峰值吞吐的 53%;而 AITER 仅为 1,944 TFLOP/s,峰值利用率为 39%。在 DeepSeek 风格的 MLA 形状(QK 维度 192、V 维度 128)测试中,ZYFA 更是以 2,464 TFLOP/s 对比 AITER 的 1,237 TFLOP/s,取得了 1.99 倍的显著优势。同时,ZYFA 在 Mistral-7B 和 Qwen3-8B 层面的能效比(TFLOP/J)也优于 AITER,表明其不仅速度更快,而且更节能。更重要的是,在 head dimension 128 的单次注意力调用中,ZYFA 的 FP8 相对 fp64 的 RMS 误差比 AITER 低 62% 至 78%,证明其在提升性能的同时,显著改善了数值精度。

性能背后的考量:时钟频率与指标解读

虽然基准测试显示了高达 1.99 倍的性能提升,但真实模型激活上的加速比(1.16-1.76 倍)相对较低,这背后涉及到 MI355X 硬件的一个重要特性:其持续时钟频率会受到处理数据中操作数位翻转数量的影响。AITER 倾向于生成更多零值或次正规数,这些数值的位翻转次数较少,使得 AITER 在某些情况下能够以更高的时钟频率运行。

例如,在 Mistral-7B 层测试中,ZYFA 运行时钟频率为 1,628 MHz,而 AITER 却能达到 2,108 MHz。尽管 ZYFA 的时钟频率较低,但它通过更高效的周期利用率,在每个周期内完成了更多有意义的计算工作,从而仍能提供整体的性能优势。这强调了在评估内核性能时,不仅要看绝对吞吐量,更要深入分析其背后的时钟频率、指令调度效率以及计算密度。因此,高斯随机输入测试虽然能展示理论峰值能力,但真实模型激活测试更能反映实际应用中的性能表现。

应用边界与工程决策:何时采用 ZYFA?

ZYFA 的发布为 AMD Instinct MI355X 用户提供了明确的性能优化路径,但其适用范围和工程决策仍需审慎考量。目前,ZYFA 内核主要覆盖等长查询(query)和键(key)的因果前向预填充注意力,尚不支持推理过程中的解码(decode)阶段、分页 KV 缓存(paged KV cache)或变长批处理(variable batch size)等高级功能。这意味着对于端到端的 LLM 服务场景,ZYFA 的性能优势仍需与其他优化(例如高效的 KV 缓存管理和批处理策略)相结合才能全面体现。

此外,当前的结果均由 Zyphra 团队发布,尚未提及独立的第三方复现验证。虽然其技术机制清晰且数据详尽,但在关键任务部署前,建议技术团队进行自己的测试与验证。值得注意的是,ZYFA 报告的误差降低主要针对单次注意力内核操作,不等同于完整模型输出的精度提升或任务效果的直接改善。在实际应用中,FP8 精度带来的大部分误差可能源于 Q、K、V 输入的量化本身,而非注意力内核。因此,在使用 ZYFA 时,仍需结合模型量化策略和整体精度需求进行综合评估。ZYFA 推荐在 ROCm 7.2.3 和 PyTorch 2.12 环境下使用,并明确支持 CDNA4 (gfx950) 架构,这为技术负责人提供了清晰的部署指引。