Faster large language model decoding with asymmetric sparse attention
More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding
Computation and LanguageArtificial Intelligence
Summary
Large language models slow down when they pay attention to too many details at once. The authors noticed that focusing less on some parts and more on others can speed things up. They created a new attention method called SAGA that treats keys and values differently to be faster without losing quality. Their tests show this method can double decoding speed while keeping answers just as good. They also offer a way to convert existing models to use this method without retraining from scratch.
What this means in practice
- •For machine learning engineers: Speed up LLM inference in production by adopting asymmetric sparse attention to reduce decoding time with minimal quality loss.
- •For ai platform operators: Convert existing pretrained language models to faster architectures with fine-tuning to save cloud costs without retraining.
Authors
Noam Elata, Itay Lamprecht, Mikey Shechter, Daniel Ohayon, Itay Hubara, Daniel Soudry
Abstract
utoregressive generation in Large Language Models (LLMs) is constrained by the memory and computational demands of attention mechanisms. Sparse attention methods mitigate this cost by selecting only high-probability entries of the attention matrix. We observe that in many such methods, this renders the probability-value multiplication negligible, shifting the bottleneck to the query-key step. Key heads can therefore be reduced to accelerate inference, while retaining more value heads preserves capacity with limited additional decoding cost. We introduce Sparse Asymmetric Group-Query Attention (SAGA), which decouples key and value head counts to exploit this principle, and pair it with approximate top-N (Atop-N) attention, a simple sparse attention method designed to study the interaction between sparsity and head-count asymmetry. We formalize the benefits of this asymmetry theoretically and validate them empirically through latency measurements and quality evaluations on models up to 1.5B parameters. Together, SAGA and Atop-N achieve end-to-end decoding speedups exceeding $2\times$ over our full-attention GQA baseline at long contexts. Models trained from scratch with SAGA nearly match the quality of comparable GQA variants on the evaluated benchmarks. To facilitate adoption, we introduce an efficient fine-tuning method that converts pretrained models to the SAGA architecture, enabling practitioners to benefit from our approach without costly retraining.