注意力运算优化

搜索文档
清华SageAttention3,FP4量化5倍加速!且首次支持8比特训练
机器之心· 2025-06-18 17:34
核心观点 - 清华大学陈键飞团队提出的SageAttention3实现了5倍于FlashAttention的推理加速,并在多种视频和图像生成大模型中保持端到端精度表现 [2][5] - SageAttention3是首个针对BlackWell架构的全FP4量化即插即用注意力算子,在RTX 5090上达到1040 TOPS速度,比H100上的FlashAttention3快1.65倍 [2] - 该团队还首次提出可训练的8比特注意力(SageBwd)用于大模型训练加速,在微调任务中保持与全精度注意力相同结果 [2] 技术突破 推理加速 - SageAttention系列迭代加速效果:V1(2.1倍)、V2(3倍)、V2++(3.9倍)、V3(5倍) [2] - 在RTX5090上实现HunyuanVideo端到端3倍加速(164s vs 489s),视频质量几乎无损 [5][33] - 算子速度相比FlashAttention2和xformers分别提升5倍和10倍 [22] FP4量化创新 - 采用BlackWell硬件支持的Microscaling FP4量化,选择NVFP4格式(准确率99.52% vs MXFP4的98.37%) [15] - 提出两阶段量化解决P矩阵缩放因子范围狭窄问题,显著提升精度 [15][16] - 算法流程优化实现高效Attention算子,两行代码即可加速任意模型推理 [5] 8比特训练 - 对Q/K/V采用Per-block INT8量化,对P采用无量化Overhead的Per-token量化 [17] - 反向传播中保留dOVT为FP16精度,量化其他四个矩阵乘法 [18] - 在Base Model微调任务中与BF16注意力表现完全一致 [33] 应用效果 - 已应用于多个知名大模型:Vidu、CogvideoX、Mochi、Wan、HunyuanVideo、Flux、Llama3、Qwen等 [1] - 在CogvideoX(2B)上端到端速度从64s提升至27s [33] - 训练加速方面,8比特注意力在微调任务中完全无损,但预训练任务与全精度仍有差距 [34] 硬件适配 - 充分利用BlackWell架构的FP4 Tensor Core,RTX5090速度达FP16 Tensor Core的8倍 [9] - 实现GPU底层CUDA Kernel优化,在算子速度和模型精度上均有显著提升 [21]