PyTorch 发布 Flash Attention 4 MXFP8 版本:Blackwell 架构训练性能再创新高
Meta 团队在 PyTorch 中开源了针对 NVIDIA Blackwell 架构优化的 Flash Attention 4 (FA4) 低精度版本,全面支持 MXFP8 前向与反向传播。该实现通过端到端的零散列模块(Zero-gather jagged module)和优化的 TMEM 分配策略,在 LLM 训练形状上实现了高达 2.85 PF/s 的前向吞吐,相比 BF16 模式提升 1.6 倍。这是首个在生产级 GEM 训练中大规模应用 MXFP8 FA4 的前沿方案,显著降低了大模型训练的计算成本与显存压力。