FLAMIA：让 Blackwell 低精度硬件真正加速 Attention 训练

长上下文 Transformer 训练中，Attention 往往仍是主要计算瓶颈。NVIDIA Blackwell GPU 为 FP8、FP4 等低精度矩阵乘法提供了更高吞吐，但把理论算力转化为实际训练速度，并不是把数据直接转换成低精度格式那么简单。量化本身会产生额外开销，共享缩放因子还可能破坏因果约束，而前向与反向采用不同数值路径也会影响训练稳定性。

Yandex Research 在论文中提出 FLAMIA（FLAsh MIcroscaling Attention），将 FlashAttention 的思路与 Blackwell 的低精度硬件结合，重点优化 Attention 的完整前向和反向流程。方案提供两条主要路径：MXFP8，以及在 Query-Key 分数计算和反向重计算中使用 NVFP4、其余乘法使用 FP8 的混合路径。研究者没有将所有步骤都压到 FP4，因为这样会带来更多量化成本，并可能使损失表现略有恶化。

核心设计

• 复用量化结果。 FLAMIA 按 32 个 token 的块共享二次幂缩放因子，并针对分数梯度采用 32×32 的块结构。这样，同一份低精度操作数可以服务于不同矩阵方向，减少重复量化、显存访问和中间数据搬运。
• 防止量化造成因果泄漏。 在标准因果 Attention 中，当前位置不应看到未来 token。但如果一个共享尺度由包含未来位置的数据共同决定，即使最终应用了因果掩码，尺度本身也可能携带未来信息。FLAMIA 采用文档局部尺度、避免饱和以及按 key 折叠尺度等办法降低这种风险，并在支持的批处理调用中提供更严格的因果选项。
• 统一前向与反向数值。 反向阶段重算概率时，匹配前向使用的量化、分块、缩放和指数计算约定。论文强调，单步梯度相似度较高，并不自动意味着长时间训练稳定。
• 面向硬件的融合与调度。 内核将量化融合进计算流程，并让 softmax 与矩阵乘法重叠执行；反向低精度计算还采用协作线程阵列，以提高 Blackwell 的利用率。对于长度不同、以 packed layout 存放的多个文档，FLAMIA 使用文档局部量化和工作负载感知调度，无需先重新打包。

结果与意义

在多种因果 packed-document 工作负载上，包含量化时间在内，MXFP8 路径相较每种工作负载中最快的评测基线，在 B200 上实现约 1.10–1.40 倍加速，在 B300 上实现约 1.06–1.42 倍加速。研究还在 Llama 风格模型预训练中验证了训练稳定性与质量保持。

FLAMIA 的价值在于，它展示了低精度加速不能只靠更小的数据格式。真正的收益来自量化策略、因果语义、前后向数值和 GPU 调度的协同设计。对于 Blackwell 平台上的长上下文训练，这类内核优化有望缩小硬件峰值吞吐与实际 Attention 性能之间的差距；但其收益仍取决于文档长度分布、批处理方式和具体工作负载，不能简单外推为所有模型或场景的固定加速比。

来源：Hugging Face Daily Papers