Optimizing Attention

Attention is the one operation that every token depends on, and it is the most expensive part of inference across models and architectures. A small gain in attention speed therefore gives a gain in the whole system, and a small error in attention can spoil the output. This lesson explains the two ways to optimize it, one that keeps results identical and one that changes them a little, and names the main techniques of each kind.

Afterwards you can explain why attention is costly even with a KV cache, what FlashAttention and PagedAttention each fix, how sliding window attention changes the cost formula, and what the other attention variants and the Mamba family do.

The lesson continues from Finding the Bottleneck, which computed the memory traffic of standard attention, and from Attention and Mixture of Experts.

Example 1: why attention is expensive and delicate

For LLMs, attention scales quadratically with input length. Each attention calculation depends on the K and V values of every earlier token. In practice, attention during decode takes linear time, because the KV cache stores the key and value results for previous tokens.

A linear algorithm can still be very expensive. Suppose a sequence of 100,000 tokens: each decode step touches 100,000 cached entries per head per layer, and this repeats for every generated token. Attention is therefore one of the most costly parts of inference across models and architectures, and optimizing it is a highly active field of research.

Attention is also delicate, because each token depends on every previous token. A small error in attention can build up quickly. Optimizations to it need care, and the check for quality is the eval work described in the model selection lesson.

flowchart TB
    A["Attention optimization"] --> B["Implementation improvements: faster kernels, lossless, still quadratic"]
    A --> C["New algorithms: better than quadratic, small quality cost"]
    B --> B1["FlashAttention"]
    B --> B2["PagedAttention"]
    C --> C1["Sliding window"]
    C --> C2["Gated, linear, compressed, multi-latent"]
    C --> C3["Other architectures: Mamba, hybrids"]

Example 2: the problem in the basic algorithm

The attention algorithm is straightforward, and the standard implementation fits in a handful of lines. It is also inefficient. In the standard algorithm of the previous lesson, the intermediate matrices S and P are written to memory at the end of one step and read back immediately in the next step.

Using the numbers from that lesson (N = 4,096, d = 128, FP16):

size of S = 2 x N x N = 33,554,432 bytes (32 MiB)
size of P = same       = 33,554,432 bytes (32 MiB)

traffic for S and P: written once and read once each
  = 4 x 33,554,432 = 134,217,728 bytes  (this is the 8 N^2 term: 134,217,728 of 138,412,032 total)
share of all memory traffic = 134,217,728 / 138,412,032 = 97 percent

Almost all the memory traffic goes to moving two intermediate results that are used once and thrown away. That is the waste that implementation work removes.

Example 3: two strategies, lossless and lossy

There are two strategies:

Example 4: FlashAttention

The most famous implementation is the series of FlashAttention papers and kernels. The basic algorithm fits in a few lines of code. FlashAttention takes tens of thousands of lines, because it implements attention in hand-fused kernels built for specific GPUs. (A fused kernel combines several steps into one program on the GPU, so intermediate results stay on the chip.) FlashAttention for the H100 uses different code from FlashAttention for the B200. The B200 is a newer NVIDIA GPU, covered in the GPU generations lesson.

FlashAttention works by removing the extra reads and writes to memory, and by laying out the attention steps to fit what the GPU does well. It helps most for compute-bound operations, such as LLM prefill and video generation.

An idealized count shows what removing the reads and writes of S and P can save. If S and P never leave the chip, only Q, K, and V are read once and O is written once:

idealized traffic = reads of Q, K, V + write of O
                  = 3 x (2 x N x d) + 2 x N x d = 8 N d
                  = 8 x 4,096 x 128 = 4,194,304 bytes (4.2 MB)

standard traffic  = 138,412,032 bytes
ratio             = 138,412,032 / 4,194,304 = 33 times less traffic

intensity with the same operations = 8,640,266,240 / 4,194,304 = 2,060 ops per byte
H100 ops:byte = 295   ->   2,060 > 295, compute-bound

This is an idealized lower bound, and not a measurement of FlashAttention. Real kernels process the matrices in blocks and cannot reach exactly this traffic. The count shows why removing the intermediate traffic moves attention from the memory-bound side of the roofline to the compute-bound side, which is the reason FlashAttention matters most when the work is already compute-bound, such as prefill.

flowchart LR
    subgraph STD["Standard attention"]
        S1["Compute S, write to memory"] --> S2["Read S, softmax, write P"] --> S3["Read P, compute O"]
    end
    subgraph FLASH["FlashAttention (fused)"]
        F1["Load blocks of Q, K, V, compute S, P, and O on chip, write O"]
    end

Example 5: PagedAttention

The KV cache grows quickly. It fills GPU memory, and reading it takes time. PagedAttention splits the KV cache into blocks, called pages, which are found through a lookup table. The cache can then be stored in fragmented memory across the GPU, and it does not need one contiguous block.

A worked case with illustrative numbers: a request has 1,000 tokens so far, and a page holds 16 tokens.

pages needed = ceil(1,000 / 16) = ceil(62.5) = 63 pages
last page    = 1,000 - 62 x 16 = 8 tokens used, 8 slots free
wasted slots = 8 of 1,008 (0.8 percent)

Contiguous alternative: reserve space for a maximum length of, say, 8,192 tokens up front
  reserved = 8,192 slots, used 1,000, wasted 7,192 (88 percent)

The page size and the maximum length are examples, and engines choose their own. With pages, memory is handed out as the sequence grows, and pages need not sit next to each other because the lookup table records where each page is. This is the same idea as virtual memory in an operating system, which stores a program's memory in non-contiguous pages.

flowchart LR
    T["Lookup table for one request"] --> P1["Page 1 at GPU address A"]
    T --> P2["Page 2 at GPU address Q"]
    T --> P3["Page 3 at GPU address C"]

FlashAttention and PagedAttention are valuable, but neither changes the fact that attention is a quadratic algorithm. The KV caching lesson returns to paged caches.

Example 6: new attention algorithms

The following variants change the underlying time and space complexity:

Here O(...) is "big O" notation, which describes how cost grows with size. O(N^2) means cost grows with the square of N, and O(N w) means cost grows with N times w.

A worked case for sliding window attention on a sequence of N = 100,000 tokens:

full attention:           N x N = 100,000 x 100,000 = 10,000,000,000 comparisons
window w = 8,192 (8K):    N x w = 100,000 x 8,192   =    819,200,000 comparisons   (12.2 times fewer)
window w = 32,768 (32K):  N x w = 100,000 x 32,768  =  3,276,800,000 comparisons   (3.1 times fewer)

(8K and 32K here mean 8,192 and 32,768.) The window makes the cost grow linearly with N for a fixed w. For short sequences, shorter than the window, nothing is saved.

The idea is intuitive: tokens near each other affect each other more than tokens far apart. The sentence being read now follows closely from the sentence before it, and much less from a sentence at the start of a long text. The intuition can be extended through training. When a method such as sliding window attention is applied during training, the resulting models keep their quality when the same method is used in inference. This is the reason the source says that training techniques can reduce the quality effect.

Example 7: avoiding attention with other architectures

Another line of research avoids attention altogether by using an architecture other than the transformer. Mamba is a selective state-space model. It replaces self-attention with a recurrent state update, which means that it keeps a fixed-size summary (the state) and updates it as each token arrives. This gives linear scaling with sequence length.

attention at token t: look at all t earlier tokens (cache grows with t)
Mamba at token t:     update one fixed-size state   (state does not grow with t)

Hybrid models sometimes mix Mamba-style state-space blocks with transformer blocks. Uses of state-space models are still limited, but hybrid models are becoming more popular as open models adopt hybrid architectures, for example NVIDIA Nemotron 3 Nano.

Interview points

• Attention is quadratic in sequence length and linear per step during decode thanks to the KV cache. It is one of the most expensive parts of inference, and it is sensitive because errors accumulate across tokens.
• Two strategies: implementation improvements (lossless, still quadratic, make long sequences feasible) and new algorithms (better than quadratic, with some quality cost that training can reduce).
• The standard algorithm writes the intermediates S and P and reads them back straight away. Most of its memory traffic is this.
• FlashAttention uses tens of thousands of lines of hand-fused, GPU-specific kernels (H100 and B200 differ). It removes extra memory reads and writes and fits the GPU's capabilities, and it helps most for compute-bound work such as LLM prefill and video generation.
• PagedAttention divides the KV cache into pages found through a lookup table, so the cache can live in fragmented memory.
• Sliding window attention uses the previous w tokens, turning O(N^2) into O(N w), with w often 8K to 32K.
• Other variants: gated attention, linear attention, compressed attention, and multi-latent attention.
• Mamba is a selective state-space model with a recurrent state update and linear scaling. Hybrid models mix Mamba-style and transformer blocks, for example NVIDIA Nemotron 3 Nano.

Summary

• Attention costs the most and tolerates errors least, so its optimization is careful work.
• Lossless implementation improvements (FlashAttention, PagedAttention) make today's hardware handle long sequences but do not change the quadratic cost.
• New algorithms (sliding window, gated, linear, compressed, multi-latent) lower the complexity at some quality cost, which training can reduce.
• State-space models such as Mamba, and hybrids with transformer blocks, avoid full attention.
• The next section starts with the hardware these workloads run on, beginning with Inside a GPU.