ZYFA：AMD CDNA4 上 FP8 FlashAttention 的速度与精度

arXiv:2611.07496v1 公告类型：交叉收录 摘要：注意力是前沿模型的核心构建模块，在长上下文场景中往往主导受计算能力限制的预填充阶段。FP8 矩阵指令的峰值吞吐量是 BF16 的两倍，但它们只能加速矩阵乘法，而注意力中两次矩阵乘法之间的 softmax 在向量单元上运行，无法受益于 FP8 加速。在 AMD Instinct MI355X 上，厂商提供的 AITER 内核在 32K 个词元时达到 FP8 矩阵峰值吞吐量的 39%，而 BF16 则达到 49%。注意力在数值上也比前馈层对量化更加敏感。此外，与 NVIDIA 的 Hopper 和 Blackwell GPU 不同，MI355X 不会在硬件中跟踪矩阵指令的完成情况，因此能否让 softmax 与矩阵乘法重叠执行取决于指令顺序。我们提出了 ZYFA，一组面向 MI355X 的预填充内核，基于三项技术构建：（1）通过一次乘加运算和一次整数转换，直接从每个分数计算对应的 FP8 概率，而不使用指数运算；（2）采用乒乓调度，让一组线程计算其 softmax，同时另一组线程执行其矩阵乘法；（3）在舍入之前，将每个概率块按 2 的幂次放大，并使用舍入后的概率之和进行归一化。在从三个模型中捕获的注意力输入上，以持续运行时的时钟频率进行计时，ZYFA 的吞吐量达到 AITER 的 1.16–1.76 倍。在内核基准测试使用的 32K 词元高斯输入上，当注意力头维度为 128 时，其吞吐量达到 AITER 的 1.26–1.37 倍；对于 DeepSeek 风格的多头潜在注意力，则达到 1.99 倍。MI355X 的持续运行时钟频率取决于操作数的值，因此随机输入上的吞吐量无法可靠预测真实模型激活值上的吞吐量。在注意力头维度为 128 时，ZYFA 的数值误差与 Composable Kernel 和 Triton 的数值误差相比，相对差异在百分之几以内，且比 AITER 低 62–78%。源代码已发布于 https://github.com/Zyphra/zyfa。