Optimizing LLM Inference: KV Cache and PagedAttention Explained

The State of LLM Inference

Deploying Large Language Models (LLMs) in production is an entirely different beast compared to training them. While training is dominated by compute-heavy matrix multiplications, inference—the process of generating text token by token—is heavily memory bound. As developers increasingly integrate open-weight models into applications, understanding the underlying mechanics of how these models serve requests becomes critical. A naive deployment strategy quickly hits a wall: throughput plummets, latency skyrockets, and GPU memory throws out-of-memory errors long before the compute units are fully saturated.

The primary culprit behind these inference bottlenecks is not the size of the model weights alone, but rather the dynamic memory required to keep track of the conversation context. Every time an LLM generates a new word, it needs to look back at everything it has seen and generated so far. Recomputing this context from scratch for every single token would be computationally disastrous. To solve this, inference engines employ a mechanism called the KV (Key-Value) Cache. However, as we will explore, standard KV caching introduces its own set of severe memory management challenges, particularly when handling multiple concurrent users with varying context lengths.

This deep dive will uncover the architectural hurdles of serving LLMs at scale. We will dissect the role of the KV Cache, examine why naive memory allocation leads to catastrophic fragmentation, and explore how a concept borrowed from traditional operating systems—PagedAttention—has revolutionized high-throughput LLM serving. Whether you are building a custom RAG pipeline or deploying models for thousands of concurrent users, mastering these concepts is essential for maximizing hardware utilization and minimizing operational costs.

The Inference Bottleneck: Why LLMs are Memory Bound

To grasp why LLM inference struggles with memory, we must look at how the Transformer architecture generates text. Transformers are autoregressive; they generate one token at a time. For every new token, the model performs a forward pass using the entire sequence of preceding tokens. If an input prompt has 1,000 tokens, generating the 1,001st token requires processing all 1,000 previous tokens. Generating the 1,002nd token requires processing 1,001 tokens, and so on.

This recomputation is highly inefficient. The math involved in the self-attention mechanism—specifically the calculation of Keys (K) and Values (V) for the preceding tokens—remains exactly the same across these successive forward passes. To avoid this redundant arithmetic, inference engines cache the K and V tensors for all previously processed tokens. This is the KV Cache. When generating a new token, the model only computes the Query (Q), Key (K), and Value (V) for the current token, and fetches the previously computed K and V tensors from the cache to complete the attention calculation.

While the KV Cache saves massive amounts of compute, it trades it for a voracious appetite for memory. Let’s quantify this. For a 13-billion parameter model like Llama 2 13B running in 16-bit precision, the KV cache for a single token consumes roughly 800 kilobytes. That sounds small, but a 4,000-token conversation consumes over 3 gigabytes of VRAM just for the cache. If you are serving 20 concurrent requests, you suddenly need 60 gigabytes of memory strictly for caching, on top of the memory required for the model weights themselves.

Consequently, the bottleneck in serving LLMs shifts from Floating Point Operations Per Second (FLOPS) to memory bandwidth. The GPU compute cores sit idle waiting for the massive KV Cache tensors to be shuttled back and forth from High Bandwidth Memory (HBM) to the processor registers. Optimizing this memory access pattern is the single most impactful way to increase inference throughput.

Enter the KV Cache: Anatomy of Storing Past States

Understanding the structure of the KV Cache is crucial for diagnosing performance issues. In a standard Transformer, the attention mechanism is spread across multiple layers and multiple attention heads. For every token, each attention head in each layer produces a Key vector and a Value vector.

The KV Cache must store these vectors in a contiguous block of memory to allow for fast tensor operations. When a request comes in, the inference engine must allocate memory for this cache. But here lies the first major problem: at the start of a request, the engine does not know how many tokens the model will ultimately generate. A user might ask a question that requires a 10-word response, or they might ask for an essay that requires 2,000 words.

Historically, inference engines handled this uncertainty by over-provisioning. If a model supported a maximum context length of 4,096 tokens, the engine would pre-allocate a contiguous block of VRAM large enough to hold 4,096 tokens’ worth of KV Cache for that request. This contiguous allocation was necessary because PyTorch and similar deep learning frameworks are optimized for operations on contiguous tensors.

Let’s look at how a basic KV Cache might be initialized in PyTorch:

import torch

def initialize_kv_cache(batch_size, max_seq_len, num_layers, num_heads, head_dim):
    # Pre-allocate contiguous memory for the maximum possible sequence length
    # This leads to massive memory waste if the actual generation is short.
    k_cache = torch.zeros(
        (batch_size, num_layers, num_heads, max_seq_len, head_dim),
        dtype=torch.float16,
        device='cuda'
    )
    v_cache = torch.zeros(
        (batch_size, num_layers, num_heads, max_seq_len, head_dim),
        dtype=torch.float16,
        device='cuda'
    )
    return k_cache, v_cache

# Example for a hypothetical 13B model
k_cache, v_cache = initialize_kv_cache(
    batch_size=8, max_seq_len=4096, num_layers=40, num_heads=40, head_dim=128
)
print(f"Allocated {(k_cache.element_size() * k_cache.nelement()) / (1024**3):.2f} GB for K cache")

This approach works fine if you are running a single batch locally. However, in a production environment serving unpredictable user requests, this contiguous pre-allocation strategy creates severe inefficiencies. A request generating 100 tokens consumes the exact same VRAM footprint as a request generating 4,000 tokens, simply because the engine had to prepare for the worst-case scenario. This artificial inflation of memory requirements severely limits the number of concurrent requests the GPU can handle.

The Memory Fragmentation Problem

The pre-allocation strategy introduces two distinct types of memory waste: internal fragmentation and external fragmentation. If you have ever studied traditional operating systems, these terms should sound familiar. In the context of LLM inference, they represent the biggest barrier to high throughput.

Internal Fragmentation: This occurs within an allocated block of memory. When the inference engine pre-allocates space for 4,096 tokens, but the model only generates 500 tokens before hitting a stop sequence, the remaining space for 3,596 tokens sits completely empty. This memory is reserved and cannot be used by any other request. In practical serving scenarios, internal fragmentation can result in 60% to 80% of the allocated KV Cache memory being completely wasted.

External Fragmentation: This happens when free memory exists, but it is scattered in small, non-contiguous chunks. Because standard deep learning operations require tensors to be stored contiguously in memory, the engine might have enough total free VRAM to serve another request, but no single contiguous block large enough to fit the maximum sequence length. The GPU is forced to reject or queue the new request despite having sufficient absolute memory available.

Furthermore, standard caching makes it impossible to share memory between requests. Consider a scenario where multiple users prompt the model with the same system prompt or long context document. Under a contiguous memory allocation scheme, the exact same Key and Value tensors for that shared document are redundantly computed and stored separately for every single user, multiplying the memory waste by the number of concurrent users.

To fix this, we need a fundamental shift in how memory is managed during inference. We need a system that allocates memory dynamically, as needed, and does not require contiguous blocks. We need virtual memory for LLMs.

PagedAttention: Operating Systems to the Rescue

To solve the KV Cache fragmentation crisis, researchers developed PagedAttention, a revolutionary algorithm that borrows the concept of virtual memory and paging from traditional operating systems. Instead of allocating a single, massive, contiguous chunk of memory for a request, PagedAttention divides the KV Cache into small, fixed-size blocks (pages).

Each block typically holds the Keys and Values for a small number of tokens (e.g., 16 or 32 tokens). When a new request arrives, the engine doesn’t pre-allocate for the maximum sequence length. Instead, it allocates just enough blocks to hold the initial prompt. As the model generates new tokens, the engine dynamically allocates new blocks on the fly.

Because the blocks are fixed-size and dynamically allocated, they do not need to be contiguous in physical GPU memory. PagedAttention maintains a block table—essentially a page table—that maps the logical, contiguous token sequence of a request to the scattered physical blocks in VRAM.

Let’s illustrate how this dramatically improves memory utilization:

class BlockTable:
    def __init__(self, block_size):
        self.block_size = block_size
        # Maps logical block indices to physical block indices in the pre-allocated pool
        self.logical_to_physical = {}
        self.num_tokens = 0

    def add_token(self, physical_allocator):
        if self.num_tokens % self.block_size == 0:
            # Only allocate a new physical block when the current one is full
            new_physical_block = physical_allocator.allocate_block()
            logical_idx = self.num_tokens // self.block_size
            self.logical_to_physical[logical_idx] = new_physical_block
        self.num_tokens += 1

# This dynamically maps logical token sequences to fragmented physical memory,
# eliminating external fragmentation entirely.

By decoupling logical memory from physical memory, PagedAttention completely eliminates external fragmentation. Any free block anywhere in VRAM can be allocated to any request. Furthermore, internal fragmentation is localized only to the very last block of a request. If the block size is 16 tokens, the maximum wasted memory per request is exactly 15 tokens’ worth of cache, regardless of the sequence length. This drops memory waste from ~70% down to under 4%.

Beyond solving fragmentation, PagedAttention enables complex memory sharing. Because the KV cache is managed at the block level, multiple requests can safely share the exact same physical blocks. If five users send prompts that share the same system instructions, the PagedAttention block table simply points all five requests to the same physical blocks containing the KV cache for those instructions. This allows for near-zero cost prefix sharing, making techniques like prompt caching highly efficient.

Implementing vLLM: Practical Considerations

PagedAttention is the core innovation behind vLLM, an open-source inference engine that has become the industry standard for high-throughput LLM serving. Deploying models with vLLM requires a shift in how we configure server resources.

When you start a vLLM server, it profiles the model and the GPU to determine how much VRAM is required for the model weights and execution workspace. It then takes almost all of the remaining free VRAM and pre-allocates it as a massive pool of physical KV cache blocks. The engine manages this pool directly, bypassing the standard PyTorch memory allocator.

Here is an example of spinning up an OpenAI-compatible server using vLLM:

# Install vLLM
pip install vllm

# Start an OpenAI-compatible API server leveraging PagedAttention
python -m vllm.entrypoints.openai.api_server     --model mistralai/Mistral-7B-Instruct-v0.2     --gpu-memory-utilization 0.90     --max-model-len 8192     --block-size 16

A critical parameter here is `–gpu-memory-utilization`. By default, vLLM reserves 90% of available GPU memory for its block pool. If you are running other processes on the same GPU, you must lower this value, otherwise vLLM will throw an Out of Memory error during initialization as it attempts to greedily claim the VRAM. The `–block-size` parameter dictates the granularity of the paging system. A smaller block size reduces internal fragmentation but slightly increases the overhead of managing the block tables.

For developers building applications on top of these APIs, the shift to PagedAttention engines means you no longer need to strictly limit concurrent requests to a tiny number. An A100 GPU that could previously handle 5 concurrent requests using standard PyTorch inference can often handle 40 to 60 concurrent requests using vLLM, drastically reducing the cost per token served.

Furthermore, to learn about how AI fits into broader architectures, consider reading about Demystifying RAG: Building a Retrieval-Augmented Generation Pipeline, or explore LLM Function Calling to connect these highly-optimized models to external tools.

For a deeper technical reference, please consult the official vLLM GitHub repository.

Conclusion

Serving Large Language Models is a memory management problem disguised as a compute problem. The KV Cache, while essential for avoiding redundant calculations, traditionally caused massive inefficiencies due to contiguous memory allocation requirements and devastating fragmentation.

PagedAttention fundamentally changed the landscape of AI infrastructure by applying time-tested operating system concepts to deep learning. By breaking the cache into dynamically allocated, non-contiguous blocks, engines like vLLM eliminate external fragmentation, minimize internal waste, and enable zero-copy memory sharing between requests. As models grow larger and context windows expand from thousands to millions of tokens, mastering these memory optimization techniques is no longer optional—it is the prerequisite for viable AI deployment.