TurboAttention: Efficient Attention Approximation for High-Throughput LLM Serving

TurboAttention: Efficient Attention Approximation for High-Throughput LLM Serving

Speakers: Hao Kang, PhD Student, Georgia Tech; Srikant Bharadwaj, Microsoft; James Hensman, Microsoft; Tushar Krishna, Georgia Tech; Victor Rühle, Microsoft; Saravan Rajmohan, Microsoft

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In the rapidly evolving landscape of large language models (LLMs), efficient inference at scale remains a paramount challenge. Hao Kang, a second-year PhD student from Georgia Tech, presented "TurboAttention" at MLSys 2025, a novel approach developed during his research at Microsoft, addressing the critical bottlenecks of memory and computation in LLM inference. The talk outlined how TurboAttention significantly improves throughput and reduces latency by employing a combination of progressive quantization and a sparse-activated SoftMax approximation.

The core of the problem lies in the attention mechanism, which, while fundamental to LLM performance, becomes prohibitively expensive for long context windows. This expense manifests as a memory bottleneck due to the ever-growing KV cache during decoding and a computational bottleneck from the quadratic complexity of self-attention during pre-filling. TurboAttention presents a two-pronged solution: Flash Q, which fuses progressive quantization with Flash Attention to optimize KV cache management and computation, and Sparse Activated SoftMax (SAS), an innovative approximation technique that dramatically reduces the computational overhead of the SoftMax operation itself. This work is pivotal for deploying large-scale LLMs more economically and efficiently, enabling broader access to advanced AI capabilities.

Background

▶ Watch: Problem: Attention Cost in LLM Inference (0:00)

The inference phase of large language models is notoriously resource-intensive, particularly when processing long input sequences or generating extensive outputs. Profiling a 14-billion parameter model on an NVIDIA A100 GPU revealed that the attention mechanism is the primary culprit behind these inefficiencies. During the decoding phase, attention accounts for over 60% of the total inference time, predominantly due to the memory bottleneck imposed by the KV cache. The KV cache stores the key and value vectors for all previously processed tokens, growing linearly with context length, quickly exhausting GPU memory and leading to expensive offloading or re-computation. Concurrently, the pre-filling phase, where the initial prompt is processed, sees attention consume over 15% of the inference time, driven by the computational bottleneck of self-attention's quadratic complexity ($O(N^2)$) with respect to sequence length $N$.

Prior research has attempted to mitigate these issues, primarily through KV cache quantization or compression techniques. Works like Kivi Gear have focused on compressing KV caches to reduce memory footprint, thereby improving loading and calculation efficiencies. While these methods successfully address part of the memory bottleneck, they introduce new challenges. Specifically, the process of dequantizing low-bit integer KV caches back to high-precision floating points (e.g., FP16) incurs additional computational overhead. Furthermore, these prior works often overlook the inherent computational cost of the SoftMax execution, which itself contributes significantly to the attention mechanism's overall latency. Thus, two critical bottlenecks remained largely unaddressed: the dequantization overhead of compressed KV caches and the computational intensity of the SoftMax function. TurboAttention directly targets these remaining inefficiencies.

Key Findings

▶ Watch: Flash Q: Progressive Quantization with Flash Attention (2:06)

TurboAttention introduces a holistic approach to overcome the persistent memory and computational bottlenecks in LLM inference, particularly for long-context attention. The project's key findings and contributions are encapsulated in two primary methodologies: Flash Q and Sparse Activated SoftMax (SAS).

Firstly, Flash Q directly tackles the dequantization overhead and memory footprint of KV caches. By fusing a progressive quantization scheme with Flash Attention v2, TurboAttention enables efficient computation and storage. The progressive quantization utilizes symmetric 8-bit integers for computation to leverage efficient low-bit matrix multiplications, while employing a hybrid 2-bit and 4-bit asymmetric quantization for KV cache storage and transfer, significantly reducing memory bandwidth requirements. This dual-level approach ensures both computational and memory efficiency without compromising accuracy excessively.

Secondly, the Sparse Activated SoftMax (SAS) approximation provides a novel solution to the computational intensity of the SoftMax function. Recognizing that current GPU architectures (like A100) often require converting FP16 inputs to FP32 for exponential computations, incurring substantial overhead, SAS re-engineers the SoftMax calculation. It decomposes the exponential function into an integer part, handled by a small, pre-computed lookup table, and a decimal part, approximated by a low-degree polynomial. This method drastically reduces the need for expensive FP32 operations and parallelizes the approximation, leading to significant speedups.

The empirical results demonstrate the effectiveness of TurboAttention. When evaluated with various LLM sizes (including 7B and 14B parameter models) on challenging reasoning datasets, the 4-bit TurboAttention achieved near-lossless inference accuracy. While the 2-bit version was acknowledged as "lossy," it offered further speed improvements for scenarios where some accuracy trade-off is acceptable. Crucially, TurboAttention showcased substantial performance gains, achieving up to 1.8x faster attention latency compared to Flash Attention v2 on a single NVIDIA A100 GPU for a 7-billion parameter model, across diverse context lengths and batch sizes. These findings highlight TurboAttention as a robust and practical solution for enhancing the throughput and reducing the latency of LLM serving in real-world applications.

Technical Deep Dive

▶ Watch: Sparse Activated Softmax Approximation (SAS) Introduction (3:50)

TurboAttention's technical innovation stems from its synergistic combination of advanced quantization techniques and a novel SoftMax approximation, meticulously designed to integrate with modern GPU architectures and existing efficient attention mechanisms.

The first core component, Flash Q, is built upon the foundation of progressive quantization, a concept introduced in papers like Q-Serve. This involves a two-level quantization strategy:

1. Computation: For matrix multiplications (matmul), symmetric int8 quantization is applied. Symmetric quantization has a fixed zero point, enabling efficient low-bit integer arithmetic directly on the GPU's tensor cores. This is crucial for reducing computational overhead.
2. Storage and Transfer: For storing and transferring KV caches, asymmetric 4-bit quantization is used. Asymmetric quantization allows for a floating zero point, which often provides higher accuracy for a given bit-width, making it suitable for memory-efficient storage.

TurboAttention enhances this by fusing it with Flash Attention v2. Flash Attention is known for its memory-efficient block-wise computation of attention, avoiding materialization of the large $N \times N$ attention matrix in GPU global memory. Flash Q introduces block-wise int8 quantization, allowing the int8 matrix multiplications to be performed within Flash Attention's block-wise structure. This ensures that the computational benefits of int8 matmul are realized directly where attention is computed. Furthermore, to push memory efficiency even further, TurboAttention employs hybrid 2-bit and 4-bit mixed precision for KV cache compression. This adaptive scheme allows for finer-grained control over memory footprint, potentially storing less critical parts of the KV cache at 2-bit precision and more important parts at 4-bit, dynamically optimizing the memory-accuracy trade-off.

The second, equally critical component is the Sparse Activated SoftMax (SAS) approximation. The traditional SoftMax computation involves several steps: converting inputs from FP16 to FP32, performing exponential computations, divisions, and additions. The speaker highlighted a key bottleneck: current GPU architectures (specifically A100) primarily support exp operations efficiently in FP32 or int32, making the FP16 to FP32 conversion a significant overhead. SAS cleverly circumvents this by mathematically decomposing the exponential function, specifically exp(-X).

The decomposition separates exp(-X) into an integer part and a decimal part:

$$ \exp(-X) = \exp(\text{integer part}) \times \exp(\text{decimal part}) $$

For the integer part, since exp(integer part) is a constant for any given integer, a lookup table can be pre-built. The talk emphasizes that this lookup table can be remarkably small because the input values to SoftMax often exhibit sparsity, meaning many values correspond to negligible SoftMax outputs. Only six items in the lookup table were found to be sufficient for "lossless inference" in their evaluation.

For the decimal part, which typically falls within a very small range, a polynomial approximation can be used. The research found that a degree-two polynomial yields "near lossless" results, while a degree-three polynomial achieves "almost lossless" accuracy. The crucial aspect here is that the search in the lookup table and the polynomial approximation computation can be parallelized, further accelerating the SoftMax calculation. This innovative approach effectively bypasses the expensive FP16-to-FP32 conversion and leverages the inherent properties of SoftMax inputs to achieve significant speedups without substantial accuracy degradation.

The end-to-end workflow of TurboAttention integrates these components seamlessly across the pre-filling and decoding stages:

• Pre-fill Stage:

1. The FP16 query (Q), key (K), and value (V) tensors are converted into 8-bit integers using symmetric quantization.
2. Attention computation is performed using Flash Attention with these int8 tensors, leveraging the SAS support for the SoftMax operation.
3. The resulting KV caches are then further compressed into hybrid 2-bit and 4-bit mixed precision for efficient storage.

• Decoding Stage:

1. The existing KV caches are dequantized from their hybrid 2/4-bit compressed format back into int8.
2. The newly generated query (Q) and value (V) tensors (from the current token) are quantized from FP16 into int8.
3. The attention computation proceeds with these int8 tensors using Flash Attention and SAS.

Crucially, the speaker noted that "all operators are fused with matmul," indicating a highly optimized implementation that minimizes data movement and kernel launch overheads. This detailed design ensures that TurboAttention effectively addresses both the memory bandwidth and computational intensity of LLM attention, leading to substantial performance gains.

Experimental Setup & Results

▶ Watch: TurboAttention's Full Pre-fill and Decoding Workflow (6:00)

The evaluation of TurboAttention was conducted to rigorously assess its performance and accuracy benefits under realistic LLM serving conditions. The primary hardware platform used for profiling and benchmarking was a single NVIDIA A100 GPU.

For accuracy profiling, the researchers evaluated TurboAttention across three different LLM models of varying sizes, including a 7-billion parameter model and a 14-billion parameter model. These models were tested on three reasoning datasets, which are known to be particularly "fragile" or sensitive to quantization-induced accuracy degradation. The results indicated that 4-bit TurboAttention achieved near-lossless accuracy, demonstrating its practical viability for many applications. As expected, 2-bit TurboAttention was found to be "lossy" but offered further improvements in speed, providing a trade-off option for use cases prioritizing extreme throughput over absolute precision.

Performance profiling focused on attention latency on the A100 GPU, specifically using a 7-billion parameter model. The evaluation varied both context lengths and batch sizes to cover a broad range of operational scenarios. The key finding was that TurboAttention consistently outperformed Flash Attention v2 across these conditions. This performance advantage was attributed to the combined benefits of progressive quantization (Flash Q) and the Sparse Activated SoftMax (SAS) approximation. The most significant improvement observed was an 1.8x speedup in attention latency when comparing TurboAttention to Flash Attention v2 under the same batch size and context length on a single A100.

Further details on the SAS component's efficiency were provided:

• The lookup table for the integer part of the exponential approximation required only six items to achieve "lossless inferences," underscoring the sparsity of relevant values.
• For the polynomial approximation of the decimal part, a degree-two polynomial was found to be "near lossless," while a degree-three polynomial achieved "almost lossless" accuracy, confirming that low-degree polynomials are sufficient given the small input range.

During the Q&A session, a comparison with Flash Attention v3 (FA3) on H100 was raised. The speaker clarified that while FA3 incorporates advanced optimizations like warp specialization and ping-pong scheduling, which are beneficial for H100, these optimizations could potentially be inherited or applied alongside TurboAttention. The speaker noted that "FA3 generally in small batches, their latency is better than TurboAttention. But at larger batches, it is not because we compress KV cache further to like hybrid 3 bits. So, this saves more memory and it is faster in decoding compared with Flash Attention version three." This suggests TurboAttention's strength lies in its superior memory efficiency for KV caches, making it particularly advantageous for larger batch sizes and longer contexts, even against the latest Flash Attention iterations.

Practical Implications

▶ Watch: Performance & Accuracy Results: Up to 1.8x Faster (6:51)

TurboAttention presents compelling practical implications for anyone involved in deploying and managing large language models, from infrastructure teams to model builders and application developers. The primary benefit is a significant leap in the efficiency and scalability of LLM serving. By reducing both the memory footprint and computational intensity of the attention mechanism, TurboAttention enables:

1. Higher Throughput and Lower Latency: The observed 1.8x speedup over Flash Attention v2 directly translates to more requests processed per second and faster response times for individual queries. This is crucial for applications requiring real-time interaction or handling a high volume of user traffic.
2. Reduced Infrastructure Costs: Less GPU memory usage (due to hybrid 2/4-bit KV cache compression) means more tokens can be stored on-chip, reducing the need for expensive memory transfers or even allowing larger models or batch sizes to fit within existing hardware. The computational speedups imply that fewer GPUs or less powerful GPUs could achieve the same performance, leading to substantial cost savings in cloud or on-premise deployments.
3. Extended Context Windows: The memory-efficient KV cache management directly addresses one of the biggest hurdles for long-context LLMs. Practitioners can now serve models with much longer input sequences without hitting memory walls as quickly, unlocking new capabilities for summarization, code generation, and complex reasoning tasks that demand extensive context.
4. Flexible Accuracy-Performance Trade-offs: The availability of both 4-bit (near-lossless) and 2-bit (lossy but faster) quantization options allows practitioners to tailor their deployments to specific application requirements. For highly sensitive tasks, 4-bit can be chosen; for less critical applications where speed is paramount, 2-bit offers an aggressive optimization.
5. Compatibility and Future-Proofing: The design of TurboAttention, particularly its fusion with Flash Attention, suggests good compatibility with existing and future hardware optimizations. The speaker indicated that H100-specific optimizations like TMA (Tensor Memory Accelerator) from Flash Attention v3 could be applied in conjunction with TurboAttention, ensuring that the benefits are not limited to current-generation GPUs. Furthermore, its strength in larger batches compared to Flash Attention v3 implies its continued relevance for high-throughput scenarios.

However, there are also practical considerations and limitations. The "lossy" nature of 2-bit quantization means careful evaluation is necessary to ensure it meets application-specific accuracy requirements. The current state of the kernel is undergoing internal checking for open-sourcing, which, once completed, will be a significant step towards broader adoption and integration into popular LLM serving frameworks like vLLM. Until then, its immediate accessibility for practitioners might be limited. Nevertheless, TurboAttention represents a significant step forward in making high-performance LLM serving more accessible and economically viable.

Key Takeaways

• Attention Bottlenecks: The attention mechanism is the dominant bottleneck in LLM inference, accounting for over 60% of decoding time (KV cache memory) and 15% of pre-filling time (N^2 computation) for long contexts.
• Two-Pronged Solution: TurboAttention introduces Flash Q (progressive quantization fused with Flash Attention) and Sparse Activated SoftMax (SAS) approximation to address these bottlenecks.
• Flash Q for Memory & Computation: It uses symmetric int8 quantization for efficient computation and hybrid 2/4-bit asymmetric quantization for memory-efficient KV cache storage, seamlessly integrating with Flash Attention's block-wise operations.
• SAS for SoftMax Efficiency: SAS approximates the exponential function in SoftMax by decomposing it into a small lookup table for the integer part and a low-degree polynomial for the decimal part, circumventing expensive FP16-to-FP32 conversions and achieving parallel computation.
• Significant Performance Gains: TurboAttention achieved up to 1.8x faster attention latency than Flash Attention v2 on a single NVIDIA A100 GPU for a 7B model, with 4-bit quantization being "near-lossless."
• Practical Impact: This work offers substantial improvements in LLM inference throughput, latency, and cost efficiency, particularly for long-context applications, and remains competitive with Flash Attention v3 for larger batch sizes due to superior KV cache compression.

About the Speaker(s)

The primary presenter for TurboAttention was Hao Kang, a second-year PhD student at Georgia Tech. He is advised by Professor Tushar Krishna, who is also listed as a co-author on the work. This research was conducted by Hao Kang during his time researching at Microsoft, where Srikant Bharadwaj served as his mentor and is also a co-author, alongside James Hensman, Victor Rühle, and Saravan Rajmohan from Microsoft.