SageAttention

收录理由

在Transformer推理中，注意力计算往往是最耗时的一环，SageAttention正是针对这一瓶颈设计的。它把标准注意力换成量化后的CUDA和Triton内核，针对Ampere、Ada、Hopper及Blackwell等不同代际的GPU做了专门优化，在保持语言、图像和视频模型端到端效果不变的前提下，相比FlashAttention能带来约2到5倍的吞吐提升。用法上，你可以直接替换一个函数，它会根据你的GPU型号自动选用合适的内核，接入现有推理管线很省事。项目同时维护了多个版本，SageAttention2偏向高精度场景，更新版本则更追求速度，这样按需选择很方便。它还支持torch.compile和分布式推理，对做推理加速和高效注意力研究的人来说很实用，不太适合只是随便跑个demo的情况。