Paged Attention Kernels

• Paged Attention Kernels are memory management paradigms that partition LLM KV caches into fixed-size pages to support variable-length sequences and reduce fragmentation.
• They integrate compile-time and runtime techniques like JIT fusion and virtual memory mapping to optimize throughput on GPUs, TPUs, and other accelerator architectures.
• Empirical results demonstrate significant reductions in latency and memory overhead during long-context inference, making these kernels vital for high-throughput LLM serving.

Paged Attention Kernels are a class of memory management and compute paradigms for attention mechanisms in LLM inference that structure the key-value (KV) cache as dynamically allocated fixed-size pages rather than contiguous monolithic buffers. This approach enables efficient support for variable-length sequences, mitigates memory fragmentation, significantly improves throughput for long-context inference, and facilitates efficient implementations on GPUs, TPUs, and other accelerator architectures. Modern paged attention kernel designs integrate compile-time and runtime components for page-table indirection, often leveraging JIT fusion, hardware virtual memory, or highly-optimized block-sparse GEMM routines. These kernels constitute the foundation of high-throughput, high-capacity production LLM serving frameworks.

1. Formal Definition and Mathematical Structure

PagedAttention partitions the dynamically growing KV cache for each decoding sequence into fixed-size blocks (pages) of size $P$ tokens. Each token $t\in\{0,\dots,N-1\}$ is addressed by:

$$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$

where $b$ is the page index, $o$ is the offset within the page, and $\mathrm{PT}[\cdot]$ is a per-sequence page-table mapping logical blocks to physical pages. When decoding, appending a new token $t=N$ extends the page table only when $b_{\mathrm{new}}\geq \text{PT.size}$, resulting in the allocation of a new page—a true page-fault event. All other token accesses are $O(1)$ index computations within the fused kernel (Joshi et al., 8 Jun 2025).

The paged KV cache stores all $K_t$ and $t\in\{0,\dots,N-1\}$0 vectors within a single or multi-layer physical buffer, with page allocation and release managed by a global, lock-free free-list per accelerator device. The overall memory usage per layer for $t\in\{0,\dots,N-1\}$1 tokens is:

$t\in\{0,\dots,N-1\}$2

with $t\in\{0,\dots,N-1\}$3, $t\in\{0,\dots,N-1\}$4 the per-head dimension, and $t\in\{0,\dots,N-1\}$5 bytes for FP16.

2. Kernel Architecture and Implementation Patterns

PagedAttention kernels can be implemented via several distinct mechanisms, including:

• JIT-fused CUDA/Triton: The arithmetic for page-table indirection is inlined into a single fused attention kernel (e.g., in FlexAttention), where mask and index calculation (mask_mod, index_mod) are callable hooks, and all memory access indirection is performed on the device (Joshi et al., 8 Jun 2025, Dong et al., 2024, Ringlein et al., 7 Oct 2025).
• Virtual Memory Management (VMM)-backed kernels: Using low-level CUDA VMM APIs, vAttention and vTensor reserve a contiguous virtual address space and dynamically map physical memory to active portions of the KV cache via page tables managed on the CPU. This allows use of any contiguous, high-performance compute kernel (e.g., FlashAttention) without modification, while transparently managing physical memory as pages (Prabhu et al., 2024, Xu et al., 2024).
• Paged block-tiled attention for TPUs: Ragged Paged Attention (RPA) employs a software pipeline that fuses cache updates and attention, using fine-grained tile packing and distribution-aware JAX/Pallas kernel specialization to maximize memory and FLOP utilization under ragged serving conditions (Jiang et al., 16 Apr 2026).
• Autotuned Triton/DSL kernels: Unified Triton-based kernels leverage hierarchical tiling, shared-memory staging, prefetching, and parameter auto-tuning to attain cross-vendor efficiency for block-tiled (paged) attention, often exceeding established baselines such as FlashAttention 3 (Ringlein et al., 7 Oct 2025).

A typical attention kernel thus processes each query token as a sequence of page-sized blocks from the KV cache, fusing score calculation, softmax normalization, and value aggregation within tile-local operations to maximize arithmetic intensity and minimize global memory traffic.

3. Memory Management Schemes

PagedAttention's memory strategy centers on power-of-two-sized pages (for efficient bitwise indexing), global or local page-allocating buffers, and per-sequence page tables. Key aspects:

• Fragmentation: With $t\in\{0,\dots,N-1\}$6, worst-case internal fragmentation per sequence is $t\in\{0,\dots,N-1\}$7. For operational values ($t\in\{0,\dots,N-1\}$8), this is usually $t\in\{0,\dots,N-1\}$9 overhead (Joshi et al., 8 Jun 2025).
• Compaction and compression: Extensions such as KV-Compress enable block-wise or per-head variable-rate eviction and compaction within paged-attention allocations. By moving maximal KV blocks into contiguous memory and freeing newly-emptied pages, theoretical compression rates can be matched in hardware with minimal slack (Rehg, 2024).
• Virtualized/decoupled paging: Virtual memory mechanisms decouple physical and virtual allocation, minimizing fragmentation, supporting dynamic extension and lazy deallocation, and enabling prefix sharing by reusing page table metadata alone (Xu et al., 2024, Prabhu et al., 2024).

PagedAttention enables efficient multi-GPU scaling by partitioning the global KV cache per device and employing distributed freelists.

4. Performance Evaluation and Comparison

PagedAttention kernels exhibit substantial improvements over monolithic or statically padded KV-caching by both reducing memory use and mitigating unbounded latency growth in long-context inference. Representative empirical metrics:

Sequence Len	No Cache (ms)	PagedAttention (ms)	Latency Scaling
128	5	5.1	$$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$0 equal
256	50	6.8	$$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$1 vs linear ($$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$2)
512	500	8.2	$$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$3 vs linear
1024	5,000	12.0	$$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$4 vs linear
2048	50,000	10.2	$$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$5 vs linear

On NVIDIA L4, end-to-end latency grows only $$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$6 as context increases $$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$7, whereas recomputing attention from scratch grows exponentially (Joshi et al., 8 Jun 2025).

Memory overhead remains minimal (e.g., $$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$8160 MB per layer at $$b = \left\lfloor \frac{t}{P} \right\rfloor, \quad o = t \bmod P, \quad p = \mathrm{PT}[b] \times P + o$$9 for LLaMA-7B) compared to >10 GB for model weights, with slack appearing predominantly beyond 2048 tokens.

Extensions such as KV-Compress achieve up to $b$0 cache compression with $b$1 accuracy loss and up to $b$2 throughput gains for typical LLMs and accelerators (Rehg, 2024).

Highly-tuned Triton paged attention kernels attain $b$3 of baseline FlashAttention-3 throughput on NVIDIA and match baseline performance on AMD, confirming low overhead from paging when appropriately optimized (Ringlein et al., 7 Oct 2025).

5. Integration in Modern LLM Serving Systems

PagedAttention has become a standard primitive in production LLM serving stacks. Key integration points include:

• Drop-in replacement: PagedAttention kernels can be exposed as configuration toggles in PyTorch or model serving frameworks (IBM FMS, vLLM), requiring no architectural changes or model retraining (Joshi et al., 8 Jun 2025).
• Compiler-driven extensibility: FlexAttention supports plug-in mask_mod/score_mod functions and page-table logic, enabling composition with arbitrary attention variants (Alibi, document masking) at negligible kernel overhead (Dong et al., 2024).
• Serving-optimized memory schedulers: vTensor and vAttention frameworks leverage CPU-driven virtual memory mappers to simultaneously support multiple LLMs, elastic batch sizing, and memory sharing for activations and embeddings with peak memory reclamation of $b$4 versus baseline paged kernels (Xu et al., 2024, Prabhu et al., 2024).
• TPU and cross-accelerator support: Custom paged attention kernels have been developed for modern TPUs, making full use of hardware tiling and multi-phase pipeline specialization for decode, prefill, and mixed workloads. Empirical results confirm up to $b$5 memory bandwidth utilization (MBU) and $b$6 model FLOPs utilization (MFU) (Jiang et al., 16 Apr 2026).

6. Limitations, Trade-Offs, and Future Directions

Despite significant improvements, PagedAttention kernels exhibit certain limitations:

• Inference-only in most systems: Current implementations generally support only inference; paging for training activations and optimizer state is future work (Joshi et al., 8 Jun 2025).
• Static page-size constraints: Fragmentation overhead at sequence lengths far exceeding $b$7 tokens becomes visible due to power-of-two page allocation, motivating investigation into hierarchical or adaptive page sizing.
• Kernel overheads in tightly-coupled implementations: Handwritten paged attention kernels (e.g., early vLLM implementations) incur up to $b$8 penalty vs non-paged baselines due to runtime block-table lookups and additional CPU coordination (Prabhu et al., 2024). Compiler-fused methods (FlexAttention, vTensor) nearly eliminate this penalty.
• Hardware and software dependencies: Optimal performance requires modern GPU/TPU architectures, support for JIT-fused or highly parameterized kernels (Triton, JAX/Pallas), and system-level support for CUDA VMM or equivalent APIs.
• Scope: Many results pertain to decoder-only models. Experimental results in encoder–decoder or multimodal settings are less conclusive (Joshi et al., 8 Jun 2025).

Future research directions outlined in recent works include training-time paging, hierarchical page structures, cross-device memory tiering (disk, host, device), and context-partitioned paging for modular and multimodal architectures.

7. Extensions: Compression and Integration with Sparse/Kernel Attention

PagedAttention kernels provide the infrastructure upon which advanced KV compression, selection, and patching methods can be realized:

• KV-Compress leverages per-head, per-layer page compaction with variable-rate block eviction, realizing theoretical memory savings and throughput gains with minimal accuracy loss. Pages are always evicted/compacted as full contiguous runs to sustain physical memory efficiency (Rehg, 2024).
• MAC-Attention composes with any paged-KV manager to achieve $b$9 or $o$0 per-step KV access by reusing previously computed attention outputs for matched queries, delivering up to $o$1 KV skip rate, $o$2 kernel phase speedups, and $o$3 end-to-end latency reduction for very long context LLMs (Yao et al., 31 Mar 2026).
• Sparse/compact kernel regression views. Paged attention in the context of kernel design yields exact, context-adaptive sparsity (e.g., in α-entmax/sparsemax/kNN variants), reducing compute to $o$4 (for typical in-support size $o$5) and improving length generalization in both synthetic and downstream tasks (Santos et al., 30 Jan 2026).

Paged attention thus underlies not only memory and performance scaling, but also theoretical advances in sparse and locality-sensitive attention mechanisms.

---

References

• (Joshi et al., 8 Jun 2025) Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference
• (Rehg, 2024) KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head
• (Jiang et al., 16 Apr 2026) Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU
• (Xu et al., 2024) vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving
• (Dong et al., 2024) Flex Attention: A Programming Model for Generating Optimized Attention Kernels
• (Prabhu et al., 2024) vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention
• (Ringlein et al., 7 Oct 2025) The Anatomy of a Triton Attention Kernel
• (Yao et al., 31 Mar 2026) MAC-Attention: a Match-Amend-Complete Scheme for Fast and Accurate Attention Computation
• (Santos et al., 30 Jan 2026) Sparse Attention as Compact Kernel Regression