PagedAttention: Virtual-Memory KV Cache Management, Continuous Batching, and PyTorch Architecture

KV Cache Memory Bottlenecks in Autoregressive Serving

During autoregressive generation in decoder-only Transformers, generating each subsequent token requires computing self-attention over all previous key and value tensor representations. Storing these intermediate states in the Key-Value (KV) Cache eliminates redundant quadratic recomputations ($\mathcal{O}(s^2)$) during sequential decoding:

$M_{\text{KV}}(s)=2\times N_{\text{layers}}\times N_{\text{heads, KV}}\times d_{\text{head}}\times s\times \text{sizeof}(\text{dtype})\text{ bytes}$

For a standard $70\text{B}$ parameter model ($N_{\text{layers}}=80,N_{\text{heads, KV}}=8,d_{\text{head}}=128$, FP16 precision), each token consumes $\approx 320\text{ KB}$. A sequence of $8192$ tokens consumes $\approx 2.62\text{ GB}$ of High Bandwidth Memory (HBM) per concurrent request.

Under traditional contiguous memory allocation, servers pre-allocate static contiguous buffers sized to the maximum possible sequence length $L_{\text{max}}$. This paradigm introduces three critical memory pathologies:

1. Reserved Over-Allocation: Memory is locked for $L_{\text{max}}$ tokens even if requests terminate early ($s\ll L_{\text{max}}$), leaving $60\text{–}80\%$ of allocated HBM idle.
2. Internal & External Fragmentation: Varied sequence lengths fragment GPU memory, preventing new requests from finding contiguous buffers even when aggregate free memory is sufficient.
3. Redundant Duplicate Allocations: Shared system prompts (e.g., in multi-turn chat or agentic workflows) duplicate identical KV cache pages for every active stream.

PagedAttention (Kwon et al., SOSP 2023) resolves these bottlenecks by translating operating system paged virtual memory principles to GPU KV cache management.

KV Cache Allocation Paradigms Comparison

Mathematical Foundations

Figure 1: PagedAttention virtual memory mapping translating contiguous logical token positions to non-contiguous physical GPU HBM blocks.

1. Paged Virtual Memory Mapping

Let $B\in \{8,16,32\}$ denote the fixed block size (number of tokens per physical memory block).

A request sequence of length $s$ is partitioned into $\lceil s/B\rceil$ logical blocks:

$\text{Logical Block Index }i=\lfloor t/B\rfloor ,\text{Block Offset }j=t\mathrm{mod}B$

Each request maintains a dynamic Block Table (page table) mapping logical block indices $i$ to physical block identifiers $\mathcal{B}(i)\in \{0,1,\dots ,N_{\text{total\_blocks}}-1\}$:

$\text{Physical Address}(\mathbf{K}[t])={\text{BasePtr}}_{\text{K}}+\mathcal{B}(\lfloor t/B\rfloor )\cdot S_{\text{block}}+(t\mathrm{mod}B)\cdot d_{\text{head}}$

where $S_{\text{block}}=B\times N_{\text{heads, KV}}\times d_{\text{head}}\times \text{sizeof}(\text{dtype})$.

2. Memory Utilization Bound and Fragmentation Analysis

Under naive contiguous allocation, memory efficiency ${\eta }_{\text{naive}}$ for a request with actual length $s$ under maximum allocation $L_{\text{max}}$ evaluates to:

${\eta }_{\text{naive}}=\frac{s}{L_{\text{max}}}\ll 1.0$

Under PagedAttention, internal fragmentation occurs strictly within the final allocated block:

$\text{Wasted Slots}=(\lceil s/B\rceil \cdot B)-s<B$

The effective memory utilization ${\eta }_{\text{paged}}$ is analytically bounded by:

${\eta }_{\text{paged}}=\frac{s}{\lceil s/B\rceil \cdot B}\ge \frac{s}{s+B}=1-\frac{B}{s+B}$

For block size $B=16$ and sequence length $s=2048$:

${\eta }_{\text{paged}}\ge 1-\frac{16}{2064}\approx 99.22\%$

Eliminating memory over-reservation increases concurrent batch capacity $N_{\text{batch}}$ by $\approx \frac{{\eta }_{\text{paged}}}{{\eta }_{\text{naive}}}\approx 3\text{–}5\times$, directly multiplying decoding throughput.

3. Non-Contiguous Online Softmax Attention Kernel

Standard self-attention computes $\mathbf{A}=\mathrm{softmax}\left(\frac{\mathbf{Q}{\mathbf{K}}^{\top }}{\sqrt{d_k}}\right)\mathbf{V}$ over contiguous tensors. PagedAttention factors the reduction across non-contiguous physical blocks via tiled online softmax (FlashAttention-style):

$\mathrm{Attention}(\mathbf{q},\mathbf{K},\mathbf{V})={⨁}_{b=0}^{\lceil s/B\rceil -1}\mathrm{softmax}\left(\frac{\mathbf{q}{\mathbf{K}}_{\mathcal{B}(b)}^{\top }}{\sqrt{d_k}}+{\mathbf{M}}_b\right){\mathbf{V}}_{\mathcal{B}(b)}$

Let $m^{(b)}$ represent the running maximum logit, $l^{(b)}$ represent the running normalization factor, and ${\mathbf{o}}^{(b)}$ represent the unnormalized output accumulator. For each physical block $\mathcal{B}(b)$:

${\mathbf{S}}_b=\frac{\mathbf{q}{\mathbf{K}}_{\mathcal{B}(b)}^{\top }}{\sqrt{d_k}}+{\mathbf{M}}_b\in {\mathbb{R}}^{1\times B}$

$m_{\text{new}}=\max \left(m^{(b-1)},{\max }_j({\mathbf{S}}_{b,j})\right)$

${\mathbf{P}}_b=\exp \left({\mathbf{S}}_b-m_{\text{new}}\right)$

$l^{(b)}=l^{(b-1)}\cdot \exp \left(m^{(b-1)}-m_{\text{new}}\right)+{\sum }_{j=1}^B{\mathbf{P}}_{b,j}$

${\mathbf{o}}^{(b)}={\mathbf{o}}^{(b-1)}\cdot \exp \left(m^{(b-1)}-m_{\text{new}}\right)+{\mathbf{P}}_b{\mathbf{V}}_{\mathcal{B}(b)}$

Final normalized output evaluates to: $\mathbf{O}={\mathbf{o}}^{(\text{final})}/l^{(\text{final})}$.

Figure 2: Online softmax reduction iterating across non-contiguous physical blocks without materializing intermediate global attention matrices.

4. Automatic Prefix Caching (APC) and Copy-on-Write (CoW)

When multiple requests share identical prompt prefixes, PagedAttention shares physical memory blocks using reference counting:

${\text{Memory}}_{\text{shared}}(N_{\text{req}},L_{\text{prefix}},L_{\text{gen}})=M_{\text{KV}}(L_{\text{prefix}})+{\sum }_{k=1}^{N_{\text{req}}}{M}_{\text{KV}}(L_{\text{gen},k})$

When an active stream modifies or appends to a shared block ($\text{ref\_count}>1$), the allocator triggers Copy-on-Write (CoW):

1. Allocate a fresh physical block from the free list.
2. Copy the $B$ tokens from the shared block to the new block.
3. Decrement the shared block’s reference count ($\text{ref\_count}\leftarrow \text{ref\_count}-1$).
4. Update the requesting stream’s block table pointer to the newly allocated block.

Production PyTorch Implementation

Below is a complete, modular PyTorch implementation of the PagedBlockAllocator (with reference counting and Automatic Prefix Caching) and a vectorized PagedAttention decoding kernel with online softmax.

Step 1: Paged Block Allocator & Page Table Manager

from __future__ import annotations

import hashlib
import math
from collections import deque
import torch


class PagedBlockAllocator:
    """Manages physical GPU memory blocks and per-request block tables."""

    def __init__(self, num_blocks: int, block_size: int = 16) -> None:
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.free_blocks: deque[int] = deque(range(num_blocks))
        
        # Mapping: request_id -> list of physical block IDs
        self.block_tables: dict[str, list[int]] = {}
        # Reference counting for prefix sharing / CoW: block_id -> count
        self.ref_counts: dict[int, int] = {i: 0 for i in range(num_blocks)}
        # Prefix cache: hash(tokens) -> physical_block_id
        self.prefix_cache: dict[str, int] = {}

    def _hash_tokens(self, tokens: list[int]) -> str:
        return hashlib.sha256(str(tokens).encode("utf-8")).hexdigest()

    def allocate_prompt(self, request_id: str, prompt_tokens: list[int]) -> list[int]:
        num_blocks_needed = math.ceil(len(prompt_tokens) / self.block_size)
        table: list[int] = []

        for i in range(num_blocks_needed):
            chunk = prompt_tokens[i * self.block_size : (i + 1) * self.block_size]
            chunk_hash = self._hash_tokens(chunk) if len(chunk) == self.block_size else None

            # Check Prefix Cache Hit
            if chunk_hash and chunk_hash in self.prefix_cache:
                phys_id = self.prefix_cache[chunk_hash]
                self.ref_counts[phys_id] += 1
            else:
                if not self.free_blocks:
                    raise MemoryError("GPU HBM Exhausted: No free blocks available for allocation.")
                phys_id = self.free_blocks.popleft()
                self.ref_counts[phys_id] = 1
                if chunk_hash:
                    self.prefix_cache[chunk_hash] = phys_id

            table.append(phys_id)

        self.block_tables[request_id] = table
        return table

    def append_token_slot(self, request_id: str, current_seq_len: int) -> int:
        """Allocates a new block if the token crosses the block boundary."""
        table = self.block_tables[request_id]
        if current_seq_len % self.block_size == 0:
            if not self.free_blocks:
                raise MemoryError("GPU HBM Exhausted during autoregressive decode step.")
            new_block = self.free_blocks.popleft()
            self.ref_counts[new_block] = 1
            table.append(new_block)
        return table[-1]

    def free_request(self, request_id: str) -> None:
        if request_id not in self.block_tables:
            return
        table = self.block_tables.pop(request_id)
        for phys_id in table:
            self.ref_counts[phys_id] -= 1
            if self.ref_counts[phys_id] == 0:
                self.free_blocks.append(phys_id)

Step 2: Vectorized PagedAttention Kernel (Online Softmax)

from __future__ import annotations

import math
import torch


def paged_attention_decode_step(
    q: torch.Tensor,                # [B, H_q, D] (Single decode query token)
    k_cache: torch.Tensor,          # [Total_Blocks, Block_Size, H_kv, D]
    v_cache: torch.Tensor,          # [Total_Blocks, Block_Size, H_kv, D]
    block_tables: torch.Tensor,     # [B, Max_Blocks_Per_Seq] (Physical block IDs)
    seq_lens: torch.Tensor,         # [B] (Current sequence length per request)
    block_size: int = 16,
) -> torch.Tensor:
    """Executes non-contiguous paged attention across requests via online softmax."""
    batch_size, num_heads_q, head_dim = q.shape
    num_heads_kv = k_cache.shape[2]
    num_queries_per_kv = num_heads_q // num_heads_kv # Grouped-Query Attention ratio
    scale = 1.0 / math.sqrt(head_dim)

    outputs = torch.zeros_like(q)

    for b in range(batch_size):
        cur_len = seq_lens[b].item()
        num_blocks = math.ceil(cur_len / block_size)
        q_b = q[b] # [H_q, D]

        # Initialize online softmax accumulators per head
        m_i = torch.full((num_heads_q, 1), float("-inf"), device=q.device, dtype=torch.float32)
        l_i = torch.zeros((num_heads_q, 1), device=q.device, dtype=torch.float32)
        acc_o = torch.zeros((num_heads_q, head_dim), device=q.device, dtype=torch.float32)

        for block_idx in range(num_blocks):
            phys_block_id = block_tables[b, block_idx].item()
            k_block = k_cache[phys_block_id] # [Block_Size, H_kv, D]
            v_block = v_cache[phys_block_id] # [Block_Size, H_kv, D]

            # Expand KV heads for GQA if necessary
            if num_queries_per_kv > 1:
                k_block = k_block.repeat_interleave(num_queries_per_kv, dim=1)
                v_block = v_block.repeat_interleave(num_queries_per_kv, dim=1)

            # Compute valid tokens in block
            valid_tokens = min(block_size, cur_len - block_idx * block_size)
            k_block = k_block[:valid_tokens] # [Valid, H_q, D]
            v_block = v_block[:valid_tokens] # [Valid, H_q, D]

            # Block attention scores: [H_q, Valid]
            scores = torch.einsum("hd,vhd->hv", q_b, k_block) * scale

            # Online Softmax Update
            m_block = torch.max(scores, dim=-1, keepdim=True).values
            m_new = torch.maximum(m_i, m_block)
            
            p_block = torch.exp(scores - m_new)
            alpha = torch.exp(m_i - m_new)

            l_i = l_i * alpha + torch.sum(p_block, dim=-1, keepdim=True)
            acc_o = acc_o * alpha + torch.einsum("hv,vhd->hd", p_block, v_block)
            m_i = m_new

        outputs[b] = (acc_o / (l_i + 1e-8)).to(q.dtype)

    return outputs

Step 3: Production Serving Pipeline Simulation

from __future__ import annotations

import torch
from paged_allocator import PagedBlockAllocator
from paged_attention_kernel import paged_attention_decode_step


def run_continuous_batching_simulation() -> None:
    # 1. Hardware & Memory Geometry
    num_physical_blocks = 128
    block_size = 16
    num_layers = 32
    num_heads_q = 32
    num_heads_kv = 8 # GQA
    head_dim = 128

    allocator = PagedBlockAllocator(num_blocks=num_physical_blocks, block_size=block_size)

    # Physical HBM KV Cache Pool
    k_cache = torch.randn(num_physical_blocks, block_size, num_heads_kv, head_dim, device="cuda", dtype=torch.bfloat16)
    v_cache = torch.randn(num_physical_blocks, block_size, num_heads_kv, head_dim, device="cuda", dtype=torch.bfloat16)

    # 2. Admit Prompt with Shared System Prefix (Automatic Prefix Caching)
    system_prefix = [101, 2054, 2003, 1037] * 4 # 16 tokens (1 Block)
    req1_prompt = system_prefix + [1001, 1002, 1003]
    req2_prompt = system_prefix + [2001, 2002, 2003]

    table1 = allocator.allocate_prompt("req_1", req1_prompt)
    table2 = allocator.allocate_prompt("req_2", req2_prompt)

    print(f"Req 1 Physical Block Allocation: {table1}")
    print(f"Req 2 Physical Block Allocation: {table2}")
    print(f"Shared Prefix Block Ref Count: {allocator.ref_counts[table1[0]]} (Expected: 2)")

    # 3. Simulate Iteration-Level Continuous Batching Step
    max_blocks = max(len(table1), len(table2))
    block_table_tensor = torch.zeros(2, max_blocks, dtype=torch.long, device="cuda")
    block_table_tensor[0, :len(table1)] = torch.tensor(table1)
    block_table_tensor[1, :len(table2)] = torch.tensor(table2)

    seq_lens = torch.tensor([len(req1_prompt), len(req2_prompt)], device="cuda")
    q = torch.randn(2, num_heads_q, head_dim, device="cuda", dtype=torch.bfloat16)

    out = paged_attention_decode_step(q, k_cache, v_cache, block_table_tensor, seq_lens, block_size)
    print(f"Decoded Output Tensor Shape: {out.shape} (Success)")


if __name__ == "__main__":
    if torch.cuda.is_available():
        run_continuous_batching_simulation()

Empirical Benchmark Evaluation

Serving throughput, memory efficiency, and time-to-first-token (TTFT) across high-concurrency benchmarks:

Troubleshooting Common PagedAttention Faults

1. CUDA Out-of-Memory on Engine Warmup

• Symptom: Server crashes during initialization before receiving traffic: torch.cuda.OutOfMemoryError.
• Root Cause: gpu_memory_utilization set too high ($\ge 0.98$), leaving insufficient headroom for PyTorch CUDA context, model weights, activations, and NCCL communication buffers.
• Remedy: Set gpu_memory_utilization = 0.90 (or 0.85 for multi-GPU Tensor Parallelism).

2. Zero Prefix Cache Hit Rate in Production

• Symptom: vllm:prefix_cache_hit_rate metric reports $0.0\%$, and system prompt prefill latencies spike.
• Root Cause: Non-deterministic prompt string formatting (e.g., dynamic timestamps, random whitespace, or varying BOS/EOS token additions) altering hash signatures.
• Remedy: Ensure system prompt strings and token IDs are strictly identical across requests; place dynamic variables at the end of the prompt sequence.

3. CPU Swap Thrashing & High Inter-Token Latency Spikes

• Symptom: Inter-token latency suddenly degrades from $20\text{ ms}$ to $400\text{ ms}$ under heavy concurrency.
• Root Cause: Active working set exceeds physical GPU HBM, triggering synchronous PCIe block swaps between GPU and host CPU RAM.
• Remedy: Increase max_num_seqs throttling limits, configure --swap-space 8 (GB), or deploy horizontal model replicas.

References

1. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., & Stoica, I. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. Proceedings of the 29th ACM Symposium on Operating Systems Principles (SOSP 2023).
2. Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. International Conference on Learning Representations (ICLR 2024).
3. Zheng, L., et al. (2024). SGLang: Efficient Execution of Structured Language Model Programs. arXiv:2312.07104.
4. DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.