--- title: Paged Attention Kernels url: https://www.emergentmind.com/topics/paged-attention-kernels type: topic --- # Paged Attention Kernels Paged Attention Kernels are a class of memory management and compute paradigms for attention mechanisms in large language model (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 [2506.07311]. The paged KV cache stores all $K_t$ and $V_t$ 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 $N$ tokens is: \[ \text{KV Memory} = M \cdot P \cdot d \cdot 2 \approx N \cdot d \cdot 2 + (P-(N \bmod P))d2 \] with $M = \lceil N/P \rceil$, $d$ the per-head dimension, and $2$ 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 [2506.07311, 2412.05496, 2511.11581]. - **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 [2405.04437, 2407.15309]. - **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 [2604.15464]. - **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 [2511.11581]. 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 $P=2^\ell$, worst-case internal fragmentation per sequence is $\leq P-1$. For operational values ($P \ll N$), this is usually $<5\%$ overhead [2506.07311]. - **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 [2410.00161]. - **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 [2407.15309, 2405.04437]. 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 | $\approx$ equal | | 256 | 50 | 6.8 | $10\times$ vs linear ($1.3\times$) | | 512 | 500 | 8.2 | $10\times$ vs linear | | 1024 | 5,000 | 12.0 | $10\times$ vs linear | | 2048 | 50,000 | 10.2 | $10\times$ vs linear | On NVIDIA L4, end-to-end latency grows only $\sim2\times$ as context increases $16\times$, whereas recomputing attention from scratch grows exponentially [2506.07311]. Memory overhead remains minimal (e.g., $\sim$160 MB per layer at $N=2048$ 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 $8\times$ cache compression with $<1\%$ accuracy loss and up to $5\times$ throughput gains for typical LLMs and accelerators [2410.00161]. Highly-tuned Triton paged attention kernels attain $>100\%$ of baseline FlashAttention-3 throughput on NVIDIA and match baseline performance on AMD, confirming low overhead from paging when appropriately optimized [2511.11581]. ## 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 [2506.07311]. - **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 [2412.05496]. - **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 $\sim70\%$ versus baseline paged kernels [2407.15309, 2405.04437]. - **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 $86\%$ memory bandwidth utilization (MBU) and $73\%$ model FLOPs utilization (MFU) [2604.15464]. ## 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 [2506.07311]. - **Static page-size constraints**: Fragmentation overhead at sequence lengths far exceeding $2048$ 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 $20\%{-}28\%$ penalty vs non-paged baselines due to runtime block-table lookups and additional CPU coordination [2405.04437]. 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 [2506.07311]. 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 [2410.00161]. - **MAC-Attention** composes with any paged-KV manager to achieve $O(1)$ or $O(r)$ per-step KV access by reusing previously computed attention outputs for matched queries, delivering up to $99\%$ KV skip rate, $14\times{-}18\times$ kernel phase speedups, and $2\times{-}3\times$ end-to-end latency reduction for very long context LLMs [2604.00235]. - **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(n\cdot s)$ (for typical in-support size $s$) and improving length generalization in both synthetic and downstream tasks [2601.22766]. Paged attention thus underlies not only memory and performance scaling, but also theoretical advances in sparse and locality-sensitive attention mechanisms. --- **References** - [2506.07311] Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference - [2410.00161] KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head - [2604.15464] Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU - [2407.15309] vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving - [2412.05496] Flex Attention: A Programming Model for Generating Optimized Attention Kernels - [2405.04437] vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention - [2511.11581] The Anatomy of a Triton Attention Kernel - [2604.00235] MAC-Attention: a Match-Amend-Complete Scheme for Fast and Accurate Attention Computation - [2601.22766] Sparse Attention as Compact Kernel Regression Source: https://www.emergentmind.com/topics/paged-attention-kernels