FlashAttention and PagedAttention: Two Fixes for Two Different Bottlenecks

Anyone who spends time around LLMs hears FlashAttention and PagedAttention mentioned as if they were rivals, and "which one should I use?" turns up constantly in ML communities. The framing is wrong. They tackle two unrelated problems at two different stages of a model's life, and a serious serving stack often runs both at once.

Why attention is costly

Self-attention in a Transformer computes this:

Attention(Q, K, V) = softmax(QKᵀ / √d) · V

The formula looks harmless, but it hides two separate bottlenecks that depend on when you run it.

1. During computation (training and prefill), the attention matrix grows as O(N²) with the sequence length N. A 10,000-token document produces a matrix of 100 million cells, which gets computed and shuffled around GPU memory over and over.
2. During generation (inference and decoding), each new token requires the model to remember the Key and Value vectors of every earlier token. That store is the KV cache, and it grows with every token and every additional user request.

A restaurant makes a handy picture. The first problem is how efficiently the chef cooks a single dish, which is the computation. The second is how well the restaurant manages storage when hundreds of orders are in flight, which is memory management. FlashAttention speeds up computation, and PagedAttention makes memory use efficient while serving. Keep that split in mind, since everything else follows from it.

FlashAttention: faster attention math

The problem

A GPU has two kinds of memory. HBM (high bandwidth memory) is big, tens of GB, but comparatively slow. SRAM, the on-chip memory, holds only a few MB but is extremely fast.

A standard implementation builds the whole N by N attention matrix and writes it to HBM, reads it back so softmax can run in SRAM, writes the result again, then reads it once more for the final multiplication by V. All that back-and-forth between slow and fast memory is the real bottleneck. The arithmetic itself isn't.

Picture a chef in a small, fast kitchen (SRAM) with every ingredient stored in a huge warehouse (HBM) a mile off. Standard attention is the chef driving over, taking a little, driving back, cooking a bit, driving back to store leftovers, and repeating for each step. The traveling wastes more time than the cooking. FlashAttention is one well-planned trip: bring back a properly sized batch, do every cooking step in the kitchen, and only return to the warehouse once the dish is finished.

How it works

FlashAttention doesn't change what attention computes. It changes how. The main techniques are:

• Tiling splits Q, K and V into small blocks that fit in SRAM.
• Online softmax avoids needing a whole row in memory. It keeps a running maximum and a running sum and updates both as each new block arrives, so the full N by N matrix is never materialized.
• Fused kernels run the matmul, the softmax and the second matmul inside one GPU kernel, with no intermediate writes to HBM.
• Recomputation in the backward pass means the huge attention matrix isn't stored for gradients. It's rebuilt on the fly during backpropagation, which spends a little extra compute to save a great deal of memory.

The output is exactly the same as standard attention. FlashAttention isn't an approximation, only a smarter way of getting there.

In code

PyTorch dispatches to a FlashAttention-style kernel by itself on supported GPUs such as the A100 and H100:

import torch
import torch.nn.functional as F

# PyTorch automatically dispatches to a FlashAttention-style
# kernel on supported GPUs (A100, H100, etc.)
q = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)  # (batch, heads, seq_len, head_dim)
k = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)


output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(output.shape)  # torch.Size([2, 8, 1024, 64])

You can also call the dedicated flash-attn library directly:

from flash_attn import flash_attn_func
import torch

q = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
k = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)

output = flash_attn_func(q, k, v, causal=True)

When it fits

• Training large Transformers from scratch, or fine-tuning them.
• Long context windows of 8k, 32k or 128k tokens, where the N² growth really hurts.
• Any case, in training or inference, where raw attention speed matters.

PagedAttention: efficient KV cache memory

The problem

This is purely a serving issue. In autoregressive decoding the model caches the Key and Value tensors of all earlier tokens so it needn't recompute them, which is the KV cache. Traditional serving systems run into trouble here:

• The output length isn't known ahead of time, so they pre-allocate one large contiguous chunk per request, for instance sized for max_tokens = 2048, even if the answer ends up at 50 tokens.
• That causes heavy internal fragmentation. Up to 60 to 80% of the allocated KV cache can sit empty and wasted.
• Nothing is shared. With beam search or parallel sampling, where several outputs grow from one prompt, each branch gets a full copy of the KV cache although they start from the same text.

Think of an apartment building (GPU memory) and tenants (requests). The old method makes each new tenant reserve a whole floor upfront, just in case. Most use two or three rooms out of twenty, and nobody else may touch the empty ones. PagedAttention hands out individual rooms, fixed-size blocks or pages, wherever there's space and only when needed. A directory, the block table, records which rooms belong to whom. Related tenants, like beam-search branches with a shared prompt, can even share rooms for the parts they have in common. The idea is borrowed outright from the virtual memory paging that operating systems use for processes, hence the name.

How it works

1. The KV cache is cut into fixed-size blocks, say 16 tokens each, instead of one giant contiguous array.
2. Every sequence keeps a block table that maps logical block indexes to physical memory blocks, like a page table in an OS.
3. Blocks are allocated on demand, only when a sequence has grown enough to need one, so there's no wasteful pre-allocation.
4. Sharing works through copy-on-write. For beam search, parallel sampling or a common system prompt, sequences point at the same physical blocks for their shared prefix and copy only when their content diverges.
5. A custom CUDA kernel reads from the scattered blocks via the block table. Since the blocks are no longer contiguous, this specialized kernel is the real PagedAttention operation.

In code

This is what drives vLLM, among the most popular LLM serving engines today:

from vllm import LLM, SamplingParams

# vLLM automatically manages KV cache using PagedAttention under the hood
llm = LLM(model="meta-llama/Llama-2-7b-hf")
prompts = [
    "The future of AI is",
    "PagedAttention helps LLM serving by"
]

sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

As a user you never see PagedAttention directly. It works quietly in the background so GPU memory is used well across every concurrent request. A simplified, conceptual view of what a block table might hold looks like this, and it isn't actual vLLM internals:

# Simplified conceptual example (not actual vLLM internals)
block_table = {
    "request_1": [block_3, block_7, block_12],
    "request_2": [block_3, block_7, block_9],
    # request_1 & request_2 share blocks 3 & 7 -> same prompt prefix!
}

When it fits

• Production serving with many concurrent users.
• Chatbot APIs, RAG pipelines and any system where output lengths are unpredictable.
• Beam search or parallel sampling, where sharing memory saves a lot.
• Any setup where you want to maximize GPU throughput, since more requests per GPU means a lower cost per request.

Using them together

The two do combine, and in practice they usually do. PagedAttention is the management layer: it decides where and how KV cache blocks are stored. FlashAttention, or fused kernels in its style, is the compute layer: it decides how fast the QK transpose and softmax run. vLLM pairs PagedAttention for memory with optimized fused attention kernels inspired by FlashAttention for the math, so it gets memory efficiency and computational speed at once.

To stay with the kitchen picture, PagedAttention picks which shelf in the warehouse to take ingredients from, and FlashAttention makes sure the chef cooks efficiently when they arrive.

Which one do you need?

• Training a model, or fine-tuning with long sequences: FlashAttention.
• Serving an LLM API to hundreds or thousands of concurrent users: PagedAttention, through vLLM or something similar.
• Building a production-grade serving stack: both. Most modern frameworks, including vLLM, TensorRT-LLM and Hugging Face TGI, already combine similar ideas internally.

The short version

The shared "Attention" in the names is the only reason they get lumped together. They sit on different layers of the stack. FlashAttention makes the attention math faster and lighter on memory by controlling how data moves between HBM and SRAM. PagedAttention makes KV cache storage during inference efficient and free of fragmentation by borrowing paged virtual memory from operating systems.

One optimizes compute and the other optimizes memory management for serving. Most production serving systems today, vLLM included, use both to get models that run fast and scale cheaply. So when someone asks "PagedAttention versus FlashAttention, which is better?", the honest answer is that it isn't a versus, it's an and.