FastAttention：将 FlashAttention2 扩展到 NPU 和低资源 GPU
摘要
FlashAttention 系列已广泛应用于大语言模型（LLM）的推理中。然而，FlashAttention 系列仅支持高级 GPU 架构，例如 Ampere 和 Hopper。目前，FlashAttention 系列不易移植到 NPU 和低资源 GPU。此外，FlashAttention 系列对于多 NPU 或 GPU 推理场景效率低下。在这项工作中，我们提出了 FastAttention，它开创性地将 FlashAttention 系列适配到 NPU 和低资源 GPU，以提升 LLM 推理效率。具体来说，我们以 Ascend NPU 和基于 Volta 的 GPU 为代表来设计我们的 FastAttention。我们通过提出一种新颖的两级分块策略以实现运行时加速、分块-掩码策略以节省内存以及分块-AllReduce 策略以减少通信开销，将 FlashAttention 系列迁移到 Ascend NPU。此外，我们通过重新设计共享内存中的操作数布局并引入一种简单而有效的 CPU-GPU 协作策略以实现高效内存利用，将 FlashAttention 适配到基于 Volta 的 GPU。在 Ascend NPU 上，与标准注意力实现相比，我们的 FastAttention 可以实现 10.7 倍的加速。采用 FastAttention 的 Llama-7B 的吞吐量比采用标准注意力的高出 5.16 倍。在 Volta 架构 GPU 上，与 \texttt{xformers} 中的等价实现相比，FastAttention 实现了 1.43 倍的加速。采用 FastAttention 的 Pangu-38B 使用 FasterTransformer 带来了 1.46 倍的端到端加速。结合所提出的 CPU-GPU 协作策略，FastAttention 在 8 块 V100 GPU 上支持最大 256K 的输入长度。所有代码将很快公开。
引用
@article{arxiv.2410.16663,
title = {FastAttention: Extend FlashAttention2 to NPUs and Low-resource GPUs},
author = {Haoran Lin and Xianzhi Yu and Kang Zhao and Lu Hou and Zongyuan Zhan and Stanislav Kamenev and Han Bao and Ting Hu and Mingkai Wang and Qixin Chang and Siyue Sui and Weihao Sun and Jiaxin Hu and Jun Yao and Zekun Yin and Cheng Qian and Ying Zhang and Yinfei Pan and Yu Yang and Weiguo Liu},
journal= {arXiv preprint arXiv:2410.16663},
year = {2024}
}