FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Abstract

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3$\times$ speedup on GPT-2 (seq. length 1K), and 2.4$\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

Cited by in corpus (40)

• DINOv2: Learning Robust Visual Features without Supervision
• MentaLLaMA: Interpretable Mental Health Analysis on Social Media with Large Language Models
• Large Language Models for Next Point-of-Interest Recommendation
• Understanding the Potential of FPGA-Based Spatial Acceleration for Large Language Model Inference
• Evaluation of LLM Chatbots for OSINT-based Cyber Threat Awareness
• Exploring the Roles of Large Language Models in Reshaping Transportation Systems: A Survey, Framework, and Roadmap
• Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology
• WikiChat: Stopping the Hallucination of Large Language Model Chatbots by Few-Shot Grounding on Wikipedia
• Clinical Insights: A Comprehensive Review of Language Models in Medicine
• RL4CO: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmark
• Vision Transformers: From Semantic Segmentation to Dense Prediction
• A Lorentz-Equivariant Transformer for All of the LHC
• Anatomy of High-Performance GEMM with Online Fault Tolerance on GPUs
• Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
• Optimal Kernel Orchestration for Tensor Programs with Korch
• SmartMem: Layout Transformation Elimination and Adaptation for Efficient DNN Execution on Mobile
• VoxInstruct: Expressive Human Instruction-to-Speech Generation with Unified Multilingual Codec Language Modelling
• LLM-based speaker diarization correction: A generalizable approach
• Improved particle-flow event reconstruction with scalable neural networks for current and future particle detectors
• Teola: Towards End-to-End Optimization of LLM-based Applications
• Eloquent: A More Robust Transmission Scheme for LLM Token Streaming
• Visual Language Models as Operator Agents in the Space Domain
• beeFormer: Bridging the Gap Between Semantic and Interaction Similarity in Recommender Systems
• Apt-Serve: Adaptive Request Scheduling on Hybrid Cache for Scalable LLM Inference Serving
• SMILE: Evaluation and Domain Adaptation for Social Media Language Understanding
• Improving the Capabilities of Large Language Model Based Marketing Analytics Copilots With Semantic Search And Fine-Tuning
• MoC-System: Efficient Fault Tolerance for Sparse Mixture-of-Experts Model Training
• MCUBERT: Memory-Efficient BERT Inference on Commodity Microcontrollers
• TAGC: Optimizing Gradient Communication in Distributed Transformer Training
• A Dutch Financial Large Language Model
• Real-time respiratory motion forecasting with online learning of recurrent neural networks for accurate targeting in externally guided radiotherapy
• Multi-word Tokenization for Sequence Compression
• Scalable hybrid quantum Monte Carlo simulation of U(1) gauge field coupled to fermions on GPU
• Smurfs: Multi-Agent System using Context-Efficient DFSDT for Tool Planning
• Before It's Too Late: A State Space Model for the Early Prediction of Misinformation and Disinformation Engagement
• Dye4AI: Assuring Data Boundary on Generative AI Services
• Fusing Gathers with Integer Linear Programming
• GuARD: Effective Anomaly Detection through a Text-Rich and Graph-Informed Language Model
• Towards Robust Knowledge Representations in Multilingual LLMs for Equivalence and Inheritance based Consistent Reasoning
• Breaking the ICE: Exploring promises and challenges of benchmarks for Inference Carbon & Energy estimation for LLMs