# 2.5 Optimizing Attention

_Inference Engineering_ by Philip Kiely. © 2026 Baseten Labs, Inc. All rights reserved. From Chapter 2: Models. [Full book index](https://www.baseten.co/inference-engineering/llms.txt)

For LLMs, attention scales quadratically with the length of the input sequence. Each calculation of attention depends on the K and V values of each previous token. In practice, attention is a linear-time operation during decode as the KV cache stores the results of key and value computations for previous tokens.

Even a linearly scaling algorithm gets very expensive. Attention is one of the most expensive parts of inference across models and architectures. Naturally, optimizing attention is an important and highly active research area.

Attention is a sensitive process because each token depends on every previous token. Small errors in attention can accumulate quickly, making attention optimization a delicate process.

Figure 2.14 showed that the attention algorithm itself is straightforward. However, that basic implementation is inefficient. The intermediate matrices `S` and `P` are stored at the end of one step, then immediately loaded in the next step.

There are two strategies for optimizing attention:

- **Implementation improvements:** Write higher-performance kernels that use memory and compute more efficiently.
- **New algorithms:** Create algorithms for attention that scale in better-than-quadratic time with minimal quality loss.

Implementation improvements are still limited by attention’s quadratic time complexity, but are lossless (do not affect quality) and make inference feasible for long sequences on today’s hardware. Other algorithmic approaches trade off quality for time and space complexity, though training techniques can minimize the impact.

The most famous implementation of attention is the FlashAttention series of papers and kernels. Where the basic algorithm can be implemented in a handful of lines of code, FlashAttention uses tens of thousands of lines to implement attention in hand-fused kernels built for specific GPUs – FlashAttention for H100 uses different code than FlashAttention for B200.

FlashAttention works by eliminating excess reads and writes from memory and laying out the attention algorithm to precisely fit the GPU’s capabilities. FlashAttention is especially useful for compute-bound operations like LLM prefill and video generation.

Another important implementation is PagedAttention. KV caches quickly grow large, filling GPU memory and taking time to read. PagedAttention partitions the KV cache into blocks (pages) that can be accessed via a lookup table. This means the KV cache can be stored across the GPU with fragmented memory rather than requiring a single contiguous block of memory.

While FlashAttention and PagedAttention are valuable optimizations, they don’t change the fact that attention is a quadratic algorithm. New variants of attention improve the underlying time and space complexity:

- **Sliding window attention:** Computes attention for a sliding window of the previous `w` tokens, turning attention from `O(N^2)` to `O(Nw)` where `w` is often in the range of 8K to 32K.
- **Gated attention:** Various types of layers introduced in training allow for approximating attention for certain chunks of context in linear time with respect to chunk length.
- **Linear attention:** Replaces the quadratic softmax equation with a linear-time algorithm that approximates attention.
- **Compressed attention:** Periodically compresses context from earlier in the sequence, attention considers both compressed context and uncompressed recent tokens.
- **Multi-latent attention:** Approximates attention in low-dimensional latent space.

Intuitively, it makes sense that tokens near each other in a sequence affect each other more than tokens from much earlier. The sentence I am writing now follows closely from the previous sentence, but less so from the sentence at the start of this chapter.

This intuition can be extended through training. Algorithms like sliding window attention, when applied during training, create models that keep quality high when the same technique is used in inference.

Another avenue of research is avoiding attention altogether by using a different architecture than transformers. Mamba is a selective state-space model that replaces self-attention with a recurrent state update, achieving linear scaling on sequence length. Hybrid models sometimes mix Mamba-style state-space model blocks with transformer blocks. Applications of state-space models are still limited, though hybrid models are becoming more popular with open models like NVIDIA Nemotron 3 Nano adopting hybrid architectures.