FlashAttention-4 通过为非因果推理和因果训练引入新路径，解决了 Blackwell 的 4 位浮点（FP4）张量核心的性能限制。该方法称为 Direct-P，将分数直接映射到 FP4 概率，以绕过通常在矩阵乘积缩小后占主导地位的 softmax 转换开销。

  • Direct-P 在 NVIDIA GB200 上实现了高达 2.13 倍的 bfloat16 (BF16) 前向吞吐量，用于非因果推理。
  • 因果路径将前向量化直接传递到反向传播中，使用 8 位浮点（FP8）梯度操作数从保存的量化查询和键重建概率。
  • 这种因果方法加速了完整的单 GPU 80 亿参数更新，最高可达 1.14 倍。
  • 匹配的分布式训练保留 FP8 概率和值，因为每个测试的 MXFP4 概率/值训练轨迹都出现发散。

作者认为这很重要，因为它能够在 Blackwell 架构上实现高效的硬件感知 FP4 注意力机制，克服了阻止 FP4 张量核心自动加速的依赖关系。