# Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers

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

## Key facts

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Textbooks Are All You Need II: phi-1.5 technical report](https://mlsanity.com/paper/2309.05463/) (2023)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Long-context LLMs Struggle with Long In-context Learning](https://mlsanity.com/paper/2404.02060/) (2024)
- [Data Engineering for Scaling Language Models to 128K Context](https://mlsanity.com/paper/2402.10171/) (2024)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [Long Context Compression with Activation Beacon](https://mlsanity.com/paper/2401.03462/) (2024)
- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [FlashDecoding++: Faster Large Language Model Inference on GPUs](https://mlsanity.com/paper/2311.01282/) (2023)
- [Striped Attention: Faster Ring Attention for Causal Transformers](https://mlsanity.com/paper/2311.09431/) (2023)
- [Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU](https://mlsanity.com/paper/2301.03598/) (2023)

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

- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference](https://mlsanity.com/paper/2410.18038/) (2024)
- [SwizzlePerf: Hardware-Aware LLMs for GPU Kernel Performance Optimization](https://mlsanity.com/paper/2508.20258/) (2025)
- [Switch Attention: Towards Dynamic and Fine-grained Hybrid Transformers](https://mlsanity.com/paper/2603.26380/) (2026)
- [SALE : Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling](https://mlsanity.com/paper/2505.24179/) (2025)
- [InnerQ: Hardware-Aware Tuning-Free Quantization of KV Cache for Large Language Models](https://mlsanity.com/paper/2602.23200/) (2026)
- [ShuntServe: Cost-Efficient LLM Serving on Heterogeneous Spot GPU Clusters](https://mlsanity.com/paper/2606.18600/) (2026)
- [ATTENTION2D: Communication Efficient Distributed Self-Attention Mechanism](https://mlsanity.com/paper/2503.15758/) (2025)
- [Exact Causal Attention with 10% Fewer Operations](https://mlsanity.com/paper/2510.05175/) (2025)
- [Optimizing Attention on GPUs by Exploiting GPU Architectural NUMA Effects](https://mlsanity.com/paper/2511.02132/) (2025)
- [LLM-CoOpt: A Co-Design and Optimization Framework for Efficient LLM Inference on Heterogeneous Platforms](https://mlsanity.com/paper/2602.09323/) (2026)
- [VTC: DNN Compilation with Virtual Tensors for Data Movement Elimination](https://mlsanity.com/paper/2604.09558/) (2026)
- [ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers](https://mlsanity.com/paper/2604.23798/) (2026)
- [FoldAttention: Declared-Reference Softmax for Fast Decode and Deterministic Backward](https://mlsanity.com/paper/2609.33410/) (2026)
- [MOIRA: Mass-Oriented Indexing with Ragged Attention for Long-Context Decoding](https://mlsanity.com/paper/2610.04313/) (2026)

## Abstract (arXiv)

Transformer-based models have emerged as one of the most widely used architectures for natural language processing, natural language generation, and image generation. The size of the state-of-the-art models has increased steadily reaching billions of parameters. These huge models are memory hungry and incur significant inference latency even on cutting edge AI-accelerators, such as GPUs. Specifically, the time and memory complexity of the attention operation is quadratic in terms of the total context length, i.e., prompt and output tokens. Thus, several optimizations such as key-value tensor caching and FlashAttention computation have been proposed to deliver the low latency demands of applications relying on such large models. However, these techniques do not cater to the computationally distinct nature of different phases during inference. To that end, we propose LeanAttention, a scalable technique of computing self-attention for the token-generation phase (decode-phase) of decoder-only transformer models. LeanAttention enables scaling the attention mechanism implementation for the challenging case of long context lengths by re-designing the execution flow for the decode-phase. We identify that the associative property of online softmax can be treated as a reduction operation thus allowing us to parallelize the attention computation over these large context lengths. We extend the "stream-K" style reduction of tiled calculation to self-attention to enable parallel computation resulting in an average of 2.6x attention execution speedup over FlashAttention-2 and up to 8.33x speedup for 512k context lengths.