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

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

Speakers: Rya Sanovar, Research Fellow, Microsoft; Srikant Bharadwaj; Renée St. Amant; Victor Rühle; Saravan Rajmohan

Conference: MLSys 2025

YouTube: https://slideslive.com/39042974

Overview

The proliferation of large language models (LLMs) has brought the self-attention mechanism to the forefront of AI innovation. However, efficiently executing attention, particularly during the decode phase of inference, presents significant hardware utilization challenges. This talk, presented by Rya Sanovar from Microsoft, introduces LeanAttention, a novel hardware-aware attention mechanism designed to maximize GPU occupancy and drastically improve the latency of Transformer inference for critical workloads like long-context processing and ragged batching.

LeanAttention tackles the fundamental problem of imbalanced GPU load distribution inherent in existing attention optimizations like FlashAttention and FlashDecoding. By proposing a generalized, hardware-efficient decomposition strategy inspired by Stream-K for matrix multiplications, LeanAttention ensures that all Streaming Multiprocessors (SMs) on a GPU receive an equal amount of work. This approach not only addresses a long-standing bottleneck in LLM serving but also promises substantial speedups and more consistent performance across diverse and dynamic inference scenarios, making it a crucial advancement for cloud service providers and anyone deploying LLMs at scale.

Background

▶ Watch: Introduction: The need for hardware-efficient attention (0:00)

Large Language Model (LLM) inference typically proceeds in two distinct stages: prefill and decode. The prefill stage involves computing attention for the entire input prompt against itself. This stage naturally enjoys a high scope of parallelism, as computation can be parallelized over the batch, head, and query length dimensions. In contrast, the decode phase, which follows prefill, computes attention for a single new token against the entire accumulated context generated so far. As more tokens are generated, the context length incrementally increases, but the query length remains just one token long. This fundamental difference severely limits the naive parallelization scope in the decode phase, confining it primarily to the batch and head dimensions.

This limited parallelism in the decode phase is a major source of inefficiency. The speaker highlights that attention during the decode phase can consume over half of the total processing time and end-to-end latency, even for high prompt-to-output token ratios. The root cause lies in how attention computations are currently partitioned on hardware, leading to low hardware occupancy, especially for two prevalent and challenging workloads:

1. Long Context Support: Models are increasingly supporting context lengths ranging from millions to billions of tokens. Processing these extremely long sequences with traditional methods leads to significant GPU underutilization. For instance, a model like GPT-4, potentially with 128 attention heads, running on an 8x A100 system (with over 800 cores), might only utilize 128 of those cores, leaving the vast majority idle.
2. Ragged Batch Support: This refers to batches of requests with unequal context lengths, a highly common scenario in cloud service environments. Current attention mechanisms struggle to efficiently handle the inherent heterogeneity of ragged batches, resulting in imbalanced loads across the GPU.

Let's delve into the limitations of existing attention mechanisms:

• Vanilla FlashAttention: This approach primarily leverages parallelism along the batch and head dimensions. While effective for certain scenarios, it results in highly imbalanced loads on the GPU during long-context execution. As context length increases, the latency of the decode phase grows, and more GPU cores remain idle. This is a direct consequence of not being able to parallelize along the context length dimension effectively in the decode phase.
• FlashDecoding: An improvement over vanilla FlashAttention, FlashDecoding attempts to address the underutilization by splitting the context of each head across different Streaming Multiprocessors (SMs) by a fixed factor K. While this can increase occupancy, it comes with its own set of limitations. Achieving perfect occupancy with FlashDecoding requires the number of SMs in the system to be a multiple of the total number of heads, a condition rarely met in practice. Increasing the split factor K to improve occupancy further leads to increased reduction overheads. Moreover, more splits mean less work assigned to each GPU core, which can be an inefficient use of resources.

The challenges are exacerbated by ragged batching. Even with FlashDecoding, partitioning ragged batches using fixed-split strategies still leads to imbalanced loads. Since requests rarely arrive with identical context lengths in a real-world cloud environment, these inefficiencies are pervasive. Furthermore, relying solely on batching to increase GPU occupancy faces a fundamental memory-bound limitation. As model sizes and supported context lengths grow, the need to independently cache the KV (Key-Value) context for each instance in a batch can quickly exceed the memory capacity of the system. This makes batching itself an unsustainable solution for very large contexts.

These issues collectively define the core research question LeanAttention seeks to answer: how can we decompose any attention workload, regardless of its configuration (single decode, long context, ragged batch), to ensure maximum possible GPU occupancy consistently?

Key Findings

▶ Watch: FlashAttention limitations: Imbalanced GPU loads (3:00)

LeanAttention introduces a paradigm shift in how attention computations are decomposed and executed on GPUs, leading to several critical findings and performance breakthroughs:

• Consistent Maximum GPU Occupancy: LeanAttention's core design ensures near 100% GPU occupancy, consistently across a wide range of diverse workload configurations. Unlike FlashDecoding, where occupancy varies significantly with workload characteristics, LeanAttention maintains a maximum possible occupancy, effectively eliminating idle GPU cores. For an example decode workload on an A100 GPU, LeanAttention was shown to occupy nearly all SMs, whereas FlashDecoding only occupied about half.
• Significant Speedups for Single Decode Requests: For medium to high context lengths in single decode requests, LeanAttention achieves more than twice the speedup over state-of-the-art methods like FlashDecoding. This directly translates to reduced inference latency for interactive LLM applications.
• Multifold Speedups for Ragged Batching: In ragged batching scenarios, the benefits of LeanAttention increase multifold as the heterogeneity and skewness of loads within the batch grow. This is a crucial improvement for cloud service providers dealing with dynamic and unpredictable request patterns.
• Broad Model Compatibility: The performance improvements of LeanAttention are not limited to specific model architectures. It delivers significant speedups across various popular LLMs, including Llama, Mistral, and Phi-3 models.
• End-to-End Latency Improvements: The enhanced GPU occupancy and raw computational speedups directly translate into significant improvements in end-to-end latency, making LLM serving more responsive and efficient.
• Accuracy Preservation: Crucially, LeanAttention achieves these performance gains without any loss of accuracy. It generates exact attention outputs by recognizing and utilizing the fully associative nature of its custom reductive operator, softmax rescaling, which is essential for correctly combining partial results from unequally split computations.
• Generalization of Existing Methods: LeanAttention is designed as a generalized attention mechanism. FlashAttention and FlashDecoding are shown to be special cases of LeanAttention's Stream-K'd mapping, implying that LeanAttention will either perform identically to them or, more often, perform better.
• Public Availability: LeanAttention has been made publicly available on ONNX Runtime, with plans for an open-source release, facilitating its adoption and integration into existing ML ecosystems.

These findings collectively demonstrate LeanAttention as a robust and highly efficient solution for optimizing the decode phase of Transformer inference, addressing critical performance bottlenecks that limit the scalability and responsiveness of modern LLMs.

Technical Deep Dive

▶ Watch: FlashDecoding limitations: Inefficient fixed splitting (4:00)

The core innovation of LeanAttention lies in its "lean decomposition" strategy, which aims to achieve 100% GPU occupancy by distributing work equally among all Streaming Multiprocessors (SMs) on the GPU. This fundamentally differs from prior approaches that might split work equally per head or by fixed factors, often leading to underutilization. To achieve equal work distribution per SM, LeanAttention often necessitates splitting the work of each attention head into unequal portions.

The mechanism behind this is a Stream-K'd linear mapping, inspired by the original Stream-K algorithm designed for dense matrix multiplications. Here's how it works for attention:

1. Linearization of Context: Instead of processing heads or context segments in isolation, LeanAttention "rolls out the inner mode" of the context length dimension. This means conceptually flattening the entire attention computation space (considering all heads and their respective context lengths) into a single, contiguous linear sequence of work.
2. Equal Work Partitioning: This linearized work is then partitioned into equal-sized chunks. The number of chunks corresponds directly to the number of GPU cores (SMs) available to distribute the workload over.
3. Mapping to GPU Cores: Each equal-sized chunk is then mapped to a specific GPU core. This ensures that every core receives an identical amount of computation. The consequence is that all cores execute concurrently and finish their computation simultaneously, eliminating idle time and maximizing hardware occupancy.

This Stream-K'd linear mapping is particularly beneficial for ragged batching scenarios. By treating the entire batch as a single, linear stream of context segments (even if they originate from different requests with varying lengths), LeanAttention can easily decompose this heterogeneous workload into equal chunks for each SM, ensuring balanced load distribution.

A critical challenge with splitting attention heads into unequal portions is correctly combining their partial outputs. The standard attention formula involves a softmax operation followed by a weighted sum. When computations are split, the intermediate results (e.g., partial logits or partial sums) need to be reduced (combined) correctly. This reduction operation must be fully associative to ensure that the order of combination does not affect the final result, thereby guaranteeing exact attention outputs without any accuracy loss.

LeanAttention introduces a novel reductive operator termed softmax rescaling. The talk highlights that LeanAttention both "recognizes and utilizes its fully associative nature." This property of softmax rescaling is key: it allows for the correct aggregation of partial attention scores from the unequally split head computations. Without a fully associative operator, combining partial results would either lead to accuracy degradation or require complex, non-parallelizable synchronization steps, negating the benefits of the parallel decomposition. By leveraging softmax rescaling, LeanAttention maintains the mathematical exactness of the attention output while enabling its highly parallel and hardware-efficient decomposition.

The implementation details also underscore the technical sophistication. LeanAttention's kernels were implemented using Cutlass and CuTe templates. Cutlass is a high-performance linear algebra library for CUDA, providing highly optimized building blocks for common operations like matrix multiplication. CuTe (CUDA Tensor Core Element-wise) is a C++ template library for writing high-performance CUDA kernels, offering fine-grained control over hardware resources. The choice of these low-level, high-performance libraries indicates a deep optimization effort to extract maximum performance from the underlying GPU hardware.

In essence, LeanAttention engineers a generalized, hardware-aware decomposition of the attention mechanism that dynamically adapts to the workload, ensuring maximal and consistent GPU occupancy. This is achieved through a Stream-K inspired linear mapping of the attention computation space and a mathematically sound, fully associative softmax rescaling operator for reducing partial results.

Experimental Setup & Results

▶ Watch: LeanAttention design: Equalizing work for 100% occupancy (7:00)

The evaluation of LeanAttention focused on demonstrating its efficacy across critical LLM inference workloads, primarily targeting the decode phase where traditional methods show significant weaknesses.

Hardware:

The primary hardware platform used for evaluation was an NVIDIA A100 GPU. For comparative discussions regarding the underutilization of FlashAttention, the speaker also referenced an 8x A100 system, noting its availability of over 800 cores (specifically, 860+ cores) and contrasting it with the 128 heads of a hypothetical GPT-4 model, illustrating the scale of underutilization.

Baselines:

LeanAttention was compared against FlashDecoding, which is considered the state-of-the-art for optimized attention execution in the decode phase. Vanilla FlashAttention was also referenced to highlight the initial limitations that FlashDecoding aimed to address.

Metrics:

The key performance indicators used to evaluate LeanAttention included:

• GPU Occupancy: A direct measure of how effectively the GPU's Streaming Multiprocessors (SMs) are being utilized.
• Speedup: Measured as the ratio of execution time of the baseline to LeanAttention, quantifying performance gains.
• End-to-End Latency: Reflecting the overall improvement in the time taken for inference.

Workloads:

Two primary workload configurations were used to stress-test LeanAttention:

1. Single Decode Request: This scenario involved processing a single token against an accumulating context of varying lengths, particularly focusing on medium to high context lengths. This mimics the core operation of sequential token generation.
2. Ragged Batching: This crucial workload involved processing batches of requests where each request had a different, unequal context length. The experiments explored various fixed batch sizes while increasing the heterogeneity and skewness of the context lengths within the batch to simulate real-world cloud serving conditions.

Models:

LeanAttention's performance was validated across a range of popular and representative LLM architectures, including:

• Llama
• Mistral
• Phi-3

Headline Results:

The experimental results vividly demonstrated LeanAttention's superior performance:

• Occupancy Advantage: For a specific example decode workload on an A100 GPU, LeanAttention achieved near 100% occupancy of the available SMs. In stark contrast, FlashDecoding for the same workload occupied only approximately half of the SMs. This consistent maximum occupancy was shown to hold true across different workload configurations, whereas FlashDecoding's occupancy varied significantly.
• Decode Speedup: For single decode requests, LeanAttention delivered more than twice the speedup over FlashDecoding for medium to high context lengths. This indicates a substantial improvement in the fundamental token generation speed.
• Ragged Batching Efficiency: In ragged batching scenarios, the benefits of LeanAttention scaled with the complexity of the workload. As the heterogeneity and skewness of context lengths within a batch increased, the speedups provided by LeanAttention grew "multifold." This highlights its robustness in dynamic, real-world serving environments.
• General Model Performance: The speedups were consistently observed across Llama, Mistral, and Phi-3 models, confirming the general applicability of LeanAttention's optimization.
• End-to-End Impact: The improved computational efficiency translated directly into "significant speedups in end-to-end latency," underscoring its practical value for deployers.

Ablations and Analysis:

The paper (referenced by the speaker) includes a reduction overhead analysis. This is important because splitting work, especially into unequal portions, could introduce overheads for combining the partial results. The analysis found that for most open-source models (typically having 12 or more attention heads), the number of splits per head on the evaluated processors was usually less than eight. This low number of splits resulted in "negligible reduction costs" relative to the overall computation. Furthermore, LeanAttention benefits from the same "temporal skew" observed in the original Stream-K algorithm. This means the "host CTA" (Cooperative Thread Array), responsible for reducing all partial outputs, typically completes its own attention computation before it has received all partial outputs, effectively hiding the reduction latency.

In summary, LeanAttention's experimental validation confirms its ability to unlock significantly higher GPU utilization and provide substantial speedups for the most challenging aspects of LLM inference, without compromising accuracy.

Practical Implications

▶ Watch: LeanAttention for ragged batching with linear mapping (8:00)

LeanAttention represents a significant advancement with profound practical implications for various stakeholders in the AI/ML ecosystem, particularly those involved in deploying and serving large language models.

For Practitioners and Infrastructure Teams:

• Enhanced GPU Utilization: The most immediate benefit is the ability to achieve near 100% GPU occupancy consistently. This means infrastructure teams can extract maximum computational value from their expensive GPU hardware, reducing the need for over-provisioning and lowering operational costs. Idle GPU cycles, a common problem with traditional attention mechanisms, are largely eliminated.
• Improved Throughput: By maximizing parallel execution and minimizing idle time, LeanAttention directly contributes to higher throughput for LLM inference. This is crucial for cloud service providers and organizations running large-scale LLM inference services, allowing them to serve more requests per unit of time and hardware.
• Efficient Long Context Support: As LLMs evolve to support ever-longer context windows (millions to billions of tokens), LeanAttention provides a hardware-efficient way to handle these demands. This allows for the deployment of more capable models without incurring prohibitive latency penalties.
• Robust Ragged Batching: The ability to efficiently handle ragged batches is a game-changer for real-world inference scenarios where requests arrive with varying context lengths. Infrastructure teams can now implement dynamic batching strategies more effectively, improving overall system responsiveness and resource allocation without suffering from load imbalance.

For Model Builders and Deployers:

• Faster Inference: Model deployers can expect significant reductions in end-to-end inference latency, particularly for the decode phase. This directly translates to a snappier user experience for interactive LLM applications, chatbots, and agents.
• Wider Model Applicability: LeanAttention's benefits span across various popular LLMs (Llama, Mistral, Phi-3), suggesting a broad applicability. This means model builders can leverage these optimizations regardless of their chosen Transformer architecture.
• Simplified Deployment for Complex Workloads: Deploying LLMs that handle long contexts or dynamic, heterogeneous batches becomes less of a performance headache. LeanAttention abstracts away much of the underlying hardware-aware optimization, allowing deployers to focus on model quality and service reliability.
• Future-Proofing: As LLM capabilities continue to expand, particularly in context length and multi-modal integration, LeanAttention provides a robust foundation for maintaining performance efficiency on current and future GPU architectures.

Tradeoffs and Limitations:

While LeanAttention offers substantial benefits, it's important to consider potential tradeoffs and nuances:

• Workload Dependency: The speaker explicitly notes that "the speedups you get with LeanAttention is highly contingent on the workload that you're dealing with." This implies that while it offers a general solution, its maximal impact is seen in scenarios where traditional methods struggle with occupancy – i.e., long-context decode and ragged batching. For workloads where FlashAttention or FlashDecoding already achieve high occupancy (e.g., very large batch sizes with uniform short contexts), the relative speedup might be less dramatic, though LeanAttention is designed to perform at least as well.
• Implementation Complexity: Implementing such a low-level, hardware-aware optimization requires specialized knowledge and tools (Cutlass, CuTe templates). While its availability on ONNX Runtime and planned open-source release will democratize access, integrating and maintaining these highly optimized kernels might still be more involved than using higher-level frameworks.
• Memory Footprint / Energy Efficiency: The talk does not explicitly detail the memory footprint or energy consumption. However, increased GPU occupancy and reduced idle time typically correlate with better energy efficiency as the GPU is doing useful work rather than waiting. Further analysis in these areas would provide a more complete picture of its resource efficiency.

In essence, LeanAttention delivers a critical piece of the puzzle for scalable and efficient LLM serving, enabling the full potential of modern GPUs to be harnessed for the most demanding inference workloads.

Key Takeaways

• GPU Underutilization is a Major Bottleneck: Current attention mechanisms like FlashAttention and FlashDecoding suffer from low GPU occupancy during the Transformer decode phase, especially for long context lengths and ragged batching, leading to significant latency.
• LeanAttention Maximizes GPU Occupancy: It introduces a hardware-aware, Stream-K'd decomposition strategy that ensures near 100% GPU occupancy by distributing equal amounts of work to each Streaming Multiprocessor (SM), even if it means unequal splits per attention head.
• Achieves Significant Speedups: LeanAttention delivers more than 2x speedup over state-of-the-art methods for single decode requests with medium to high context lengths, and multifold speedups for heterogeneous ragged batching scenarios.
• Guarantees Exact Outputs with Softmax Rescaling: A novel "softmax rescaling" technique is employed as a fully associative reductive operator, allowing for the correct aggregation of partial results from unequal splits without any loss of attention accuracy.
• Generalizes Existing Optimizations: FlashAttention and FlashDecoding are special cases of LeanAttention's generalized approach, ensuring that LeanAttention always performs at least as well as, or better than, these prior methods.
• Practical for LLM Deployment: Its consistent performance gains across models (Llama, Mistral, Phi-3) and availability on ONNX Runtime make it a crucial tool for practitioners and infrastructure teams aiming to improve the efficiency and responsiveness of LLM serving.

About the Speaker(s)

Rya Sanovar is a research fellow at Microsoft. The work on LeanAttention was conducted during her internship at Microsoft and formed part of her undergraduate thesis. She presented the talk at MLSys 2025. Her co-authors on this research include Srikant Bharadwaj, Renée St. Amant, Victor Rühle, and Saravan Rajmohan.