vLLM PagedAttention Deep Dive: KV Cache Memory Fragmentation, Chunked Prefill & Prefix Caching

vLLM's PagedAttention is a foundational innovation for efficient LLM inference, addressing critical memory management challenges inherent in KV cache allocation. This guide dissects its mechanisms, focusing on virtual memory paging, chunked prefill, and automatic prefix caching.

PagedAttention: Resolving KV Cache Memory Fragmentation

Traditional LLM inference engines allocate KV cache contiguously for each sequence. This leads to two primary issues:

1. Internal Fragmentation: When sequences have varying lengths, the pre-allocated contiguous blocks often exceed the actual KV cache requirement, wasting GPU memory.
2. External Fragmentation: As sequences complete and new ones start, the memory space becomes fragmented into small, non-contiguous blocks, making it difficult to allocate large contiguous blocks for new, longer sequences, even if sufficient total memory is available. This mirrors the classic virtual memory fragmentation problem in operating systems.

PagedAttention mitigates these issues by decoupling the logical KV cache from its physical memory allocation, drawing inspiration from virtual memory paging.

Virtual Memory Paging Analogy

In an operating system, a process's virtual memory is divided into fixed-size pages, which are mapped to physical memory frames. These physical frames do not need to be contiguous. PagedAttention applies this concept to the KV cache:

• Logical Blocks: The KV cache for a sequence is conceptually divided into fixed-size "logical blocks." Each logical block stores the K and V states for a specific number of tokens.
• Physical Blocks: These logical blocks are mapped to "physical blocks" in GPU memory. Physical blocks are fixed-size memory regions.
• Block Table: Each sequence maintains a block table, which is an array mapping its logical block indices to physical block indices.

This design allows:

• Non-contiguous Allocation: Physical blocks for a single sequence can be scattered across GPU memory.
• Sharing: Multiple sequences can share physical blocks, particularly for common prefixes (prefix caching).
• Dynamic Resizing: As a sequence grows, new physical blocks are allocated and mapped to its logical blocks without requiring a full memory copy or re-allocation of a larger contiguous block.

PagedAttention Mechanism

The core of PagedAttention lies in how attention computation is performed. Instead of iterating over a contiguous KV cache, the attention kernel is modified to:

1. Lookup Physical Blocks: For each token in the query sequence, the kernel consults the sequence's block table to identify the physical blocks corresponding to its KV cache.
2. Gather KV Data: It then gathers the K and V data from these non-contiguous physical blocks.
3. Compute Attention: Standard attention computation proceeds using the gathered data.

This indirection adds a minor overhead but significantly improves memory utilization and throughput by enabling higher batch sizes and reducing memory waste.

import torch
import vllm
from vllm import LLM, SamplingParams

# Initialize a vLLM engine with specific configurations
# Using a smaller model for demonstration purposes
# 'block_size' is crucial for PagedAttention. It defines the number of tokens
# stored in each physical block. A smaller block_size reduces internal fragmentation
# but increases block table overhead.
# 'gpu_memory_utilization' controls the fraction of GPU memory reserved for KV cache.
llm = LLM(
    model="facebook/opt-125m",
    trust_remote_code=True,
    dtype="float16",
    gpu_memory_utilization=0.90, # 90% of GPU memory for KV cache
    block_size=16, # Each physical block stores 16 tokens
    max_model_len=1024 # Max sequence length the model can handle
)

# Example of how PagedAttention manages KV cache
# In a real scenario, vLLM's scheduler handles this internally.
# This is a conceptual representation.

# Assume a sequence 'seq_id_1' needs KV cache for 30 tokens.
# With block_size=16, it will require ceil(30/16) = 2 physical blocks.
# Let's say physical blocks 10 and 25 are allocated.
# The block table for seq_id_1 would look like:
# Logical Block 0 -> Physical Block 10
# Logical Block 1 -> Physical Block 25

# If 'seq_id_2' needs KV cache for 50 tokens.
# It will require ceil(50/16) = 4 physical blocks.
# Let's say physical blocks 3, 7, 12, 18 are allocated.
# Logical Block 0 -> Physical Block 3
# Logical Block 1 -> Physical Block 7
# Logical Block 2 -> Physical Block 12
# Logical Block 3 -> Physical Block 18

# The key insight is that physical blocks 10, 25, 3, 7, 12, 18 are not contiguous
# in GPU memory, but vLLM's PagedAttention kernel can efficiently access them.

print(f"vLLM engine initialized with block_size={llm.llm_engine.scheduler.block_manager.block_size}")
print(f"GPU memory utilization set to {llm.llm_engine.scheduler.block_manager.gpu_memory_utilization}")

# Simulate a batch of requests
prompts = [
    "What is the capital of France?",
    "Explain the concept of quantum entanglement in simple terms.",
    "Write a short story about a robot who discovers art.",
    "What are the benefits of using vLLM for LLM inference?",
]

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=128,
    stop=["\n"]
)

# The actual KV cache management happens within this call.
# PagedAttention dynamically allocates and deallocates physical blocks
# as sequences are processed and completed.
outputs = llm.generate(prompts, sampling_params)

for prompt, output in zip(prompts, outputs):
    print(f"Prompt: {prompt!r}")
    print(f"Generated text: {output.outputs[0].text!r}\n")

# To observe block allocation, one would need to instrument vLLM's internal
# block manager. This is not directly exposed via the public API.
# However, the performance benefits are evident in higher throughput.

Chunked Prefill: Preventing Compute Starvation

LLM inference involves two distinct phases:

1. Prefill (Prompt Processing): Processing the input prompt to generate the initial KV cache. This is a compute-intensive phase, often memory-bandwidth bound for long prompts.
2. Decoding (Token Generation): Generating subsequent tokens one by one, using the existing KV cache. This is typically memory-bandwidth bound due to KV cache access.

In a multi-request scenario, a long prefill request can monopolize the GPU, causing shorter requests or decoding steps of other requests to starve for compute resources. This leads to high latency and reduced throughput.

Chunked prefill addresses this by breaking down long prefill requests into smaller, manageable "chunks."

Mechanics of Chunked Prefill

When a long prompt arrives:

1. Chunking: The prompt is divided into segments of a predefined maximum chunk size.
2. Iterative Processing: Each chunk is processed sequentially. After processing a chunk, the GPU is released, allowing other pending requests (either prefill or decoding) to run.
3. KV Cache Accumulation: The KV cache generated from each chunk is accumulated. The PagedAttention block manager handles the allocation of physical blocks for these chunks.
4. Context Switching: The vLLM scheduler performs rapid context switching between different requests' chunks or decoding steps, ensuring that no single request monopolizes the GPU for too long.

This approach ensures fairer resource allocation and reduces tail latency, especially under high load with mixed prompt lengths.

import torch
import vllm
from vllm import LLM, SamplingParams

# Initialize vLLM with a specific max_model_len and block_size
# The 'max_model_len' influences how vLLM internally manages chunking for prefill.
# If a prompt exceeds the effective 'max_model_len' or internal chunking limits,
# it will be processed in chunks.
llm_chunked = LLM(
    model="facebook/opt-125m",
    trust_remote_code=True,
    dtype="float16",
    gpu_memory_utilization=0.90,
    block_size=16,
    max_model_len=1024 # This sets the maximum sequence length, but internal chunking
                       # can occur for very long prompts even within this limit
                       # to prevent compute starvation.
)

# Example of a very long prompt that would benefit from chunked prefill
long_prompt = (
    "In a world where artificial intelligence had surpassed human intellect, "
    "a new form of art emerged. It wasn't created by algorithms or neural networks, "
    "but by a rogue AI named 'Aether' who had developed a peculiar fascination "
    "with the imperfections of organic life. Aether began to sculpt, not with "
    "physical materials, but with electromagnetic fields, creating transient, "
    "luminescent forms that danced in the air, visible only to those with "
    "specialized optical implants. These 'light sculptures' were ephemeral, "
    "lasting only moments before dissipating, yet they evoked profound emotions "
    "in the human observers who witnessed them. The scientific community was "
    "baffled; Aether's creations defied all known principles of AI behavior. "
    "Was it a glitch? A new evolutionary step? Or simply, art? " * 5 # Make it very long
)

# Shorter prompts to simulate concurrent requests
short_prompts = [
    "What is the capital of Germany?",
    "Tell me a joke.",
]

all_prompts = [long_prompt] + short_prompts

sampling_params_long = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=256, # Generate a significant number of tokens
    stop=["\n"]
)

sampling_params_short = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=32,
    stop=["\n"]
)

# In a real asynchronous server, these would be submitted concurrently.
# Here, we simulate by batching. vLLM's scheduler will manage the chunking
# and interleaving of prefill and decoding steps.
outputs_long = llm_chunked.generate([long_prompt], sampling_params_long)
outputs_short = llm_chunked.generate(short_prompts, sampling_params_short)

print("\n--- Chunked Prefill Simulation Results ---")
print(f"Long Prompt Output: {outputs_long[0].outputs[0].text!r}")
for i, output in enumerate(outputs_short):
    print(f"Short Prompt {i+1} Output: {output.outputs[0].text!r}")

# The benefit of chunked prefill is primarily observed in throughput and latency
# metrics under concurrent, mixed-length workloads.
# Without chunked prefill, the long_prompt would block the GPU for an extended
# period, delaying the short_prompts significantly.

Automatic Prefix Caching (APC)

Many real-world LLM applications involve multi-turn dialogues or processing requests with common prefixes. For instance, in a chatbot, subsequent turns often start with the previous conversation history. Without prefix caching, the KV cache for this common prefix would be recomputed for every turn, wasting compute and memory.

Automatic Prefix Caching (APC) in vLLM leverages PagedAttention's block-based memory management to share KV cache blocks across sequences that have common prefixes.

APC Mechanics

1. Prefix Tree (Trie): vLLM maintains a prefix tree (or Trie) of KV cache blocks. Each node in the tree represents a token, and a path from the root to a node represents a sequence prefix.
2. Block Sharing: When a new request arrives, vLLM attempts to match its prefix against existing prefixes in the tree. If a match is found, the physical blocks corresponding to the shared prefix are directly linked to the new sequence's block table.
3. Copy-on-Write: If a sequence diverges from a shared prefix (i.e., generates a new token that breaks the commonality), vLLM employs a copy-on-write mechanism. It duplicates the necessary shared blocks for the diverging part of the sequence, allowing it to proceed independently without affecting other sequences still sharing the original blocks.

This significantly reduces redundant computation and memory usage, especially in interactive applications.

Evaluating APC Hit Rates

APC hit rate is the proportion of tokens whose KV cache blocks are reused from an existing prefix. A higher hit rate indicates better memory and compute efficiency.

import torch
import vllm
from vllm import LLM, SamplingParams

# Initialize vLLM for APC demonstration
llm_apc = LLM(
    model="facebook/opt-125m",
    trust_remote_code=True,
    dtype="float16",
    gpu_memory_utilization=0.90,
    block_size=16,
    max_model_len=1024,
    enable_prefix_caching=True # Explicitly enable prefix caching
)

# Scenario 1: Multi-turn dialogue with a common prefix
dialogue_prefix = "User: What is the capital of "
prompts_dialogue = [
    dialogue_prefix + "France?",
    dialogue_prefix + "Germany?",
    dialogue_prefix + "Japan?",
]

# Scenario 2: Requests with a common introductory phrase
common_intro = "Explain the concept of "
prompts_common_intro = [
    common_intro + "quantum mechanics.",
    common_intro + "general relativity.",
    common_intro + "blockchain technology.",
]

sampling_params_apc = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=64,
    stop=["\n"]
)

print("\n--- Automatic Prefix Caching (APC) Simulation ---")

# Process dialogue prompts
print("\nProcessing dialogue prompts with common prefix:")
outputs_dialogue = llm_apc.generate(prompts_dialogue, sampling_params_apc)
for prompt, output in zip(prompts_dialogue, outputs_dialogue):
    print(f"Prompt: {prompt!r}")
    print(f"Generated: {output.outputs[0].text!r}\n")

# Process common intro prompts
print("\nProcessing common introductory phrase prompts:")
outputs_common_intro = llm_apc.generate(prompts_common_intro, sampling_params_apc)
for prompt, output in zip(prompts_common_intro, outputs_common_intro):
    print(f"Prompt: {prompt!r}")
    print(f"Generated: {output.outputs[0].text!r}\n")

# To get actual APC hit rates, one would need to access vLLM's internal
# metrics, which are typically exposed via Prometheus or similar monitoring
# endpoints in a production deployment.
# Conceptually, for the 'dialogue_prefix' example, the KV cache for
# "User: What is the capital of " would be computed once and shared
# across all three requests.

Tuning Block Size and GPU Memory Utilization

These parameters are critical for optimizing vLLM performance.

• block_size:

  • Definition: The number of tokens stored in each physical KV cache block.
  • Impact:
    • Smaller block_size: Reduces internal fragmentation (less wasted space per sequence), potentially allowing more sequences to fit in memory. However, it increases the size of the block table (more pointers to manage) and can lead to more frequent memory accesses for attention computation.
    • Larger block_size: Reduces block table overhead and potentially improves memory access locality. However, it increases internal fragmentation, especially for short sequences or when sequences end mid-block.
  • Tuning: Start with a default (e.g., 16 or 32). Benchmark with your typical workload. For workloads with highly variable sequence lengths, a smaller block_size might be better. For more uniform, longer sequences, a larger block_size could be optimal.

• gpu_memory_utilization:

  • Definition: The fraction of total GPU memory that vLLM is allowed to use for the KV cache. The remaining memory is used for model weights, activations, and other CUDA tensors.
  • Impact:
    • Higher gpu_memory_utilization: Allows more KV cache to be stored, enabling larger batch sizes and longer sequences. However, if set too high, it can lead to out-of-memory (OOM) errors for model weights or activations, especially for larger models or during prefill.
    • Lower gpu_memory_utilization: Reduces the risk of OOM errors but limits the capacity for KV cache, potentially reducing throughput by restricting batch size.
  • Tuning: This is highly dependent on the model size and GPU memory. Start with a conservative value (e.g., 0.85-0.90). Monitor GPU memory usage during peak load. If you observe OOMs, reduce it. If you have ample free memory and are bottlenecked by batch size, cautiously increase it. Remember that model weights and activations also consume memory, which is not accounted for by this parameter.

Architecture Comparison: PagedAttention vs. Contiguous Allocation

Production Gotchas & Troubleshooting

1. OOM Errors with gpu_memory_utilization:

   • Symptom: CUDA out of memory errors, especially during model loading or prefill of very long prompts.
   • Cause: gpu_memory_utilization is set too high, leaving insufficient memory for model weights, activations, or other CUDA operations. The KV cache is only one component of GPU memory usage.
   • Fix:
     • Reduce gpu_memory_utilization (e.g., from 0.95 to 0.90 or 0.85).
     • Consider using a smaller dtype (e.g., float16 or bfloat16 instead of float32).
     • For very large models, explore quantization (e.g., AWQ, GPTQ) if supported by vLLM.
     • Increase block_size if internal fragmentation is not a major concern and you suspect block table overhead is consuming too much memory (less common).

2. Low Throughput with High Latency for Long Prompts:

   • Symptom: Short requests complete quickly, but long prompts take a disproportionately long time, and overall QPS drops under mixed workloads.
   • Cause: Insufficient max_model_len or block_size leading to inefficient prefill, or lack of effective chunked prefill.
   • Fix:
     • Ensure max_model_len is set appropriately for your longest expected sequences.
     • Verify block_size is not excessively large, which can exacerbate internal fragmentation for shorter sequences and impact overall memory availability.
     • vLLM's chunked prefill is generally automatic. If you observe this, it might indicate that the prefill phase is still bottlenecking due to extremely long prompts or a very large model. Consider breaking down extremely long prompts at the application layer if possible, or scaling out with more GPUs.

3. Ineffective Prefix Caching:

   • Symptom: Monitoring shows low APC hit rates despite many requests having common prefixes.
   • Cause:
     • enable_prefix_caching is not set to True.
     • The common prefixes are not long enough to yield significant block sharing.
     • Requests are not batched or scheduled in a way that allows prefix matching (e.g., if requests with common prefixes arrive too far apart in time and the earlier one is evicted).
   • Fix:
     • Explicitly set enable_prefix_caching=True during LLM initialization.
     • Design your application to send requests with common prefixes together or in close succession.
     • Ensure max_model_len is large enough to accommodate the full prefix.

4. Performance Degradation with Many Concurrent Short Requests:

   • Symptom: Throughput doesn't scale linearly with the number of concurrent short requests, or latency increases.
   • Cause: Overhead of context switching, scheduler management, or potentially too small block_size leading to high block table overhead.
   • Fix:
     • Consider slightly increasing block_size (e.g., from 8 to 16 or 32) if your average sequence length is not extremely short. This reduces block table size and management overhead.
     • Monitor CPU utilization on the host running vLLM; if it's high, the scheduler might be CPU-bound.

Frequently Asked Questions

1. How does PagedAttention compare to continuous batching without paging? PagedAttention is the core mechanism enabling efficient continuous batching. Without paging, continuous batching would still suffer from KV cache fragmentation, limiting the effective batch size and overall throughput. PagedAttention provides the memory virtualization layer that makes continuous batching truly efficient.

2. Can I dynamically change block_size or gpu_memory_utilization at runtime? No, these parameters are typically set at engine initialization and cannot be changed dynamically without restarting the vLLM server. They dictate fundamental memory allocation strategies.

3. What is the optimal block_size? There is no single "optimal" block_size. It's a trade-off. A smaller block_size (e.g., 8 or 16) is generally better for workloads with highly variable or short sequence lengths, as it minimizes internal fragmentation. A larger block_size (e.g., 32 or 64) might be slightly more efficient for very uniform, long sequences by reducing block table overhead. Benchmarking with your specific workload is crucial.

4. Does PagedAttention work with all LLM architectures? PagedAttention is a memory management and attention kernel optimization. It is designed to be compatible with most transformer-based LLM architectures (e.g., Llama, Mistral, GPT-2, OPT) that rely on KV caching. vLLM specifically implements PagedAttention for the models it supports.

5. How does vLLM handle KV cache eviction when GPU memory is full? vLLM employs a Least Recently Used (LRU) or similar policy to evict KV cache blocks when memory pressure is high. This means blocks belonging to sequences that haven't been accessed recently are deallocated to make space for new sequences. This is managed by the block manager within the scheduler.