What is Paged FlashAttention?
Paged FlashAttention refers to using a FlashAttention-style, IO-aware attention kernel directly with a paged KV cache. The keys and values for previous tokens live in fixed-size memory pages, while the kernel uses a block table to find them and computes attention in tiles.
Terminology: "Paged FlashAttention" is a descriptive term, not a universally standardized, separate attention algorithm. The official FlashAttention implementation supports paged KV-cache access, including a block_table argument in flash_attn_with_kvcache.
Why combine paging with FlashAttention?
These techniques address different sources of memory inefficiency:
- Paged KV cache: Stores keys and values in reusable, potentially non-contiguous physical blocks, reducing wasted capacity from large per-request reservations.
- FlashAttention: Uses tiling and online softmax to compute exact attention without materializing a large attention-score matrix in GPU high-bandwidth memory.
- Together: An appropriate kernel follows the KV block mapping while performing tiled, IO-aware attention, avoiding a separate full-cache gather into a contiguous tensor.
How does it work?
- Locate the KV pages. The request has a block table mapping logical token blocks to physical cache blocks.
- Load small K/V tiles. The kernel fetches the required key and value data from those physical blocks.
- Compute attention incrementally. It forms query–key scores for each tile and maintains the normalization using an online, numerically stable softmax.
- Accumulate the output. It combines the partial results without storing a full attention-score matrix or reconstructing a full contiguous KV cache first.
Practical example
Suppose a cache page holds 256 tokens. A sequence has 600 tokens, with this illustrative mapping:
Logical page: [0, 1, 2]
Physical page: [5, 1, 9]
Token 300 (zero-based):
logical page = 300 // 256 = 1
offset = 300 % 256 = 44
physical page = block_table[1] = 1The attention kernel reads the relevant K/V data from physical page 1, offset 44, and processes it as part of a tile. The physical pages do not need to sit next to one another.
Note: The 256-token page size is illustrative and is also a supported page granularity in FlashAttention-2's documented paged-KV API; page-size constraints vary by kernel and version.
How is it different from Direct Paged Attention?
- Direct Paged Attention: Emphasizes reading paged K/V inside the attention kernel, without a separate gather.
- Paged FlashAttention: Adds the specific FlashAttention-style tiled, IO-aware computation and online-softmax strategy.
The two descriptions can refer to overlapping implementations. They are not necessarily competing products or mutually exclusive algorithms.
Does it change the attention result or complexity?
The intention is to compute the same exact attention operation up to floating-point differences. For standard full attention, FlashAttention mainly reduces memory traffic and temporary storage; it does not turn quadratic prefill attention arithmetic into linear-time computation.
In single-token decoding, a full quadratic score matrix is not the main issue. Efficient reads from the growing KV cache, kernel overhead, and GPU utilization are often more important.
Does it always improve inference speed?
No. Performance depends on GPU architecture, kernel/backend, context length, page layout, batch size, and competing overheads. Block-table indirection and irregular accesses can cost time, so benchmark time to first token, time per output token, throughput, and peak memory.
See "What is Paged KV Cache?" for the memory organization and "What is Direct Paged Attention?" for direct mapped reads.
In short: Paged FlashAttention combines paged KV-cache access with FlashAttention's tiled, IO-aware attention computation, reducing avoidable data movement while preserving the underlying attention operation.
References: Official FlashAttention implementation, FlashAttention paper, and PagedAttention paper.
Sources & Research References
Learning Sequence
Large Language Models Curriculum
What is Mixture of Experts (MoE) in Large Language Models?
Explore Related Questions
Deepen your understanding across related Large Language Models topics.