FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Abstract
Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FlashAttention exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4$\times$ compared to optimized baselines), with no approximation. However, FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40\% of the theoretical maximum FLOPs/s. We observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2$\times$ speedup compared to FlashAttention, reaching 50-73\% of the theoretical maximum FLOPs/s on A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72\% model FLOPs utilization).
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
FlashAttention-2 introduces a hardware-optimized attention kernel that resolves GPU occupancy and memory bandwidth bottlenecks through refined work partitioning, delivering near-GEMM efficiency and establishing a new industry standard for Transformer acceleration. The paper represents a masterclass in ML systems engineering, translating deep architectural insights into a widely adopted, production-grade implementation that has tangibly lowered the compute barrier for long-sequence modeling and accelerated the broader LLM ecosystem.
Methodology
The paper presents a highly refined GPU kernel optimization for the self-attention mechanism, directly addressing the hardware utilization bottlenecks of its predecessor. Rather than introducing mathematical approximations or algorithmic complexity reductions, the authors focus on low-level architectural efficiency: (1) algorithmic restructuring to minimize non-matrix-multiply FLOPs, (2) cross-thread-block parallelization to boost SM occupancy even for single-head computations, and (3) warp-level work distribution to minimize shared memory synchronization overhead. This represents a sophisticated application of hardware-aware algorithm design, demonstrating deep understanding of CUDA execution models, memory hierarchies, and occupancy constraints. The methodology is rigorous, well-motivated by profiling data, and directly targets the gap between theoretical peak FLOPs and realized throughput in attention kernels.
Results & evaluation
The empirical validation is robust and practically oriented. The authors benchmark against highly optimized baselines (including Triton and xFormers) on A100 hardware, demonstrating a consistent ~2x speedup over FlashAttention-1 and achieving 50-73% of theoretical peak FLOPs/s. Crucially, they validate end-to-end training performance on GPT-style models, reporting up to 225 TFLOPs/s per GPU (72% model FLOPs utilization). The experiments cover varying sequence lengths and head dimensions, confirming that the optimizations generalize across typical LLM configurations. While the evaluation is heavily NVIDIA A100-focused, the results are sufficiently comprehensive to establish clear performance boundaries and practical utility.
Reproducibility
Excellent. The implementation is open-sourced, well-documented, and has been rapidly integrated into major frameworks (PyTorch, Hugging Face Transformers, vLLM, etc.). The paper provides clear algorithmic pseudocode, explicit work-partitioning strategies, and references to the CUTLASS 3.x abstractions used, enabling other systems researchers to replicate or adapt the kernel. The reliance on standard CUDA/Triton tooling further lowers the barrier to reproduction and extension.
Limitations
The optimizations are tightly coupled to NVIDIA GPU microarchitectures (specifically Ampere/Hopper), limiting direct portability to AMD, Intel, or TPU ecosystems without significant re-engineering. The paper does not alter the asymptotic O(N^2) complexity of attention, meaning it mitigates but does not solve the fundamental scaling bottleneck for extremely long sequences. Additionally, as GPU architectures evolve (e.g., H100's dedicated attention hardware, Blackwell's tensor cores), the relative gains of manual kernel tuning may diminish, requiring continuous maintenance. The evaluation also lacks extensive ablation on non-standard attention variants (e.g., MQA, GQA, sliding window), which are now standard in modern LLMs.
Broader impact
FlashAttention-2 has fundamentally shifted the compute economics of Transformer training and inference, enabling longer context windows, faster iteration cycles, and reduced energy consumption across academia and industry. By closing the gap between attention and GEMM efficiency, it has become a foundational infrastructure component, indirectly accelerating progress in long-context reasoning, multimodal modeling, and code/audio/video generation. The work also establishes a blueprint for hardware-aware ML systems research, demonstrating that careful low-level optimization can yield outsized returns compared to purely algorithmic innovations. FlashAttention-2 introduces a hardware-optimized attention kernel that resolves GPU occupancy and memory bandwidth bottlenecks through refined work partitioning, delivering near-GEMM efficiency and establishing a new industry standard for Transformer acceleration. The paper represents a masterclass in ML systems engineering, translating deep architectural insights into a widely adopted, production-grade implementation that has tangibly lowered the compute barrier for long-sequence modeling and accelerated the broader LLM ecosystem.
Curator note
Dao; further 2x improvement over FlashAttention
Author of this paper? Send a correction.