# FastAttention: Extend FlashAttention2 to NPUs and Low-resource GPUs

arXiv: https://arxiv.org/abs/2410.16663 · HTML: https://mlsanity.com/paper/2410.16663/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 41 references that are in the map.
- First posted on arXiv 22 Oct 2024.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [Instruction Tuning with GPT-4](https://mlsanity.com/paper/2304.03277/) (2023)
- [FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU](https://mlsanity.com/paper/2303.06865/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale](https://mlsanity.com/paper/2207.00032/) (2022)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [Fast Distributed Inference Serving for Large Language Models](https://mlsanity.com/paper/2305.05920/) (2023)
- [PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation](https://mlsanity.com/paper/2104.12369/) (2021)
- [NVIDIA Tensor Core Programmability, Performance & Precision](https://mlsanity.com/paper/1803.04014/) (2018)
- [PanGu-Sigma: Towards Trillion Parameter Language Model with Sparse Heterogeneous Computing](https://mlsanity.com/paper/2303.10845/) (2023)

## Built on by (most cited first, 3 of 3)

- [Pangu Light: Weight Re-Initialization for Pruning and Accelerating LLMs](https://mlsanity.com/paper/2505.20155/) (2025)
- [Tackling the Dynamicity in a Production LLM Serving System with SOTA Optimizations via Hybrid Prefill/Decode/Verify Scheduling on Efficient Meta-kernels](https://mlsanity.com/paper/2412.18106/) (2024)
- [STEEL: Sparsity-Aware Fused Attention for Energy-Efficient Long-Sequence Inference on AMD's XDNA NPU](https://mlsanity.com/paper/2607.09385/) (2026)

## Abstract (arXiv)

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× speedup compared to the standard attention implementation. Llama-7B within FastAttention reaches up to 5.16× higher throughput than within the standard attention. On Volta architecture GPUs, FastAttention yields 1.43× speedup compared to its equivalents in textttxformers. Pangu-38B within FastAttention brings 1.46× 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.