arXiv:2410.16663cs.LG2024-10被引 8

让FlashAttention适配国产芯片和低端显卡,大幅提速大模型推理。

FastAttention: Extend FlashAttention2 to NPUs and Low-resource GPUs

  • 针对NPU和低配GPU设计分层分块策略,降低内存与通信开销。
  • 在昇腾芯片上实现10.7倍加速,Llama-7B吞吐提升5.16倍。
  • 支持8张V100组合推理,最长输入达256K,适合资源受限场景。

FlashAttention系列虽广泛用于大语言模型推理,但仅支持高端GPU架构(如Ampere和Hopper),难以迁移至NPUs和低资源GPU。此外,在多设备推理中效率低下。本文提出FastAttention,首次将FlashAttention系列适配于昇腾NPU和基于Volta架构的GPU,以提升大模型推理效率。针对昇腾NPU,提出两级分块策略以加速运行时计算,分块掩码策略节省内存,分块AllReduce策略减少通信开销;针对Volta GPU,重新设计共享内存操作数布局,并引入简单有效的CPU-GPU协同策略以优化内存利用。在昇腾NPU上,相比标准注意力实现,FastAttention获得10.7×加速,Llama-7B的吞吐量提升5.16×;在Volta架构上,较xformers快1.43×,使用FasterTransformer的Pangu-38B实现1.46×端到端加速。结合所提协同策略,8张V100 GPU可支持最大256K输入长度。代码即将开源。

原文摘要 · Abstract (English)

FlashAttention series has been widely applied in the inference of large language models (LLMs). However, FlashAttention series only supports the high-level GPU architectures, e.g., Ampere and Hopper. At present, FlashAttention series is not easily transferrable to NPUs and low-resource GPUs. Moreover, FlashAttention series is inefficient for multi- NPUs or GPUs inference scenarios. In this work, we propose FastAttention which pioneers the adaptation of FlashAttention series for NPUs and low-resource GPUs to boost LLM inference efficiency. Specifically, we take Ascend NPUs and Volta-based GPUs as representatives for designing our FastAttention. We migrate FlashAttention series to Ascend NPUs by proposing a novel two-level tiling strategy for runtime speedup, tiling-mask strategy for memory saving and the tiling-AllReduce strategy for reducing communication overhead, respectively. Besides, we adapt FlashAttention for Volta-based GPUs by redesigning the operands layout in shared memory and introducing a simple yet effective CPU-GPU cooperative strategy for efficient memory utilization. On Ascend NPUs, our FastAttention can achieve a 10.7$\times$ speedup compared to the standard attention implementation. Llama-7B within FastAttention reaches up to 5.16$\times$ higher throughput than within the standard attention. On Volta architecture GPUs, FastAttention yields 1.43$\times$ speedup compared to its equivalents in \texttt{xformers}. Pangu-38B within FastAttention brings 1.46$\times$ end-to-end speedup using FasterTransformer. Coupled with the propose CPU-GPU cooperative strategy, FastAttention supports a maximal input length of 256K on 8 V100 GPUs. All the codes will be made available soon.

大模型推理NPU优化低资源加速FlashAttention

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。