📖 AI 总结

Meta与PyTorch团队将FlashAttention-4扩展至端到端MXFP8精度，在Blackwell架构上实现了块缩放注意力机制。该实现在LLM形状下前向达到2.85 PF/s、反向达到2 PF/s，内部形状下前向2.54 PF/s、反向1.58 PF/s，较BF16分别提升最高1.6倍和1.52倍。关键技术包括：在已饱和的TMEM中容纳缩放因子的分配策略、利用redux指令实现在线转置不变的dS量化、融合RMSNorm与GEMM的量化内核，以及零聚集jagged模块——FP8数据保持非填充位置，仅对较小的缩放因子进行散射和排列。该模块已在Meta内部用于GEM训练，是首批在生产训练负载中使用的MXFP8 FA4前反向实现之一，代码已开源。