Inside LLM Inference Engines: KV Caching, PagedAttention, and Continuous Batching

The Shift from Training to Inference

During the training phase of Large Language Models (LLMs), workloads are overwhelmingly compute-bound ($O(N)$ FLOPs per parameter with massive matrix multiplications utilizing tensor cores at peak utilization). In contrast, serving generative LLMs in production is plagued by an entirely different bottleneck: memory bandwidth and VRAM capacity.

Understanding why modern runtimes—such as vLLM, TensorRT-LLM, and SGLang—outperform naive PyTorch inference scripts by $5\times$ to $20\times$ requires diving into the mechanical sympathy of GPU memory hierarchies and KV cache dynamics.

---

Anatomy of Autoregressive Generation

In a standard transformer decoder model (e.g., Llama-3, Mistral, DeepSeek), inference consists of two distinct operational phases:

1. The Prefill Phase (Prompt Ingestion):

   • The user submits an input prompt of length $S$.
   • All $S$ tokens are processed in parallel via self-attention: $Q K^T / \sqrt{d_k}$.
   • High arithmetic intensity ($\text{FLOPs} / \text{Byte}$), making it compute-bound and highly efficient on modern GPU Tensor Cores.

2. The Decode Phase (Autoregressive Token Generation):

   • The model generates tokens one by one: $t_1, t_2, \dots, t_n$.
   • Each new token requires computing attention against all previous tokens in the context.
   • For every single generated token, the entire model weights ($\sim 140\text{ GB}$ for Llama-3-70B in FP16) and historical key/value vectors must be transferred from GPU High-Bandwidth Memory (HBM) into SRAM caches.
   • Extremely low arithmetic intensity, making the decode phase strictly memory-bandwidth bound.

Memory Hierarchy Transfer per Decode Step:
[ GPU HBM (80 GB / 2-3 TB/s) ] ===(PCIe/NVLink)===> [ On-Chip SRAM (L2 / 100+ TB/s) ]
For batch_size = 1, we read 140 GB weights just to generate a single 2-byte token!

---

The KV Cache: Mathematical Formulation & Memory Footprint

To prevent recomputing attention representations of past tokens at every step ($O(L^2)$ redundant compute), inference engines store past Key and Value activations in GPU memory—the KV Cache.

The memory footprint $M_{\text{KV}}$ for a single active sequence is:

$$M_{\text{KV}} = 2 \times n_{\text{layers}} \times n_{\text{kv_heads}} \times d_{\text{head}} \times \text{bytes_per_element} \times \text{seq_len}$$

Let us calculate the memory requirement for a production model:

• Model: Llama-3-70B
• Layers ($n_{\text{layers}}$): 80
• Grouped-Query Attention KV Heads ($n_{\text{kv_heads}}$): 8
• Head Dimension ($d_{\text{head}}$): 128
• Precision: 16-bit Float (2 bytes)
• Sequence Length: 8,192 tokens

$$M_{\text{KV}} = 2 \times 80 \times 8 \times 128 \times 2 \times 8,192 \approx 2.68 \text{ GB per sequence!}$$

If a server handles a concurrent batch of 32 requests, the KV cache alone demands 85.7 GB of VRAM—exceeding the entire 80 GB capacity of an NVIDIA H100 GPU before even accounting for model weights.

---

The KV Cache Memory Crisis: Internal & External Fragmentation

In naive serving frameworks, tensors for each sequence must be stored contiguously in virtual GPU memory. Because the final generated length of a response is unknown in advance, the system must pre-allocate memory for the maximum possible length (max_seq_len, e.g., 8,192 tokens).

This introduces catastrophic memory waste:

1. Internal Fragmentation: A sequence that terminates early at 150 tokens still reserves 8,192 tokens worth of physical VRAM.
2. External Fragmentation: Memory allocations of varying sizes create fragmented holes across the address space, preventing new requests from finding contiguous blocks.
3. Reserved Reservation Buffer: Traditional systems waste 60% to 80% of their available GPU memory on unutilized pre-allocations!

---

PagedAttention: Virtual Memory for LLM Generation

In 2023, Woosuk Kwon et al. from UC Berkeley published PagedAttention (the core engine behind vLLM). PagedAttention borrows the classic Virtual Memory Paging design from Operating Systems and applies it to GPU memory management.

Logical KV Cache:   [ Block 0 ] -> [ Block 1 ] -> [ Block 2 ]
                        |              |              |
                      Page 12        Page 4         Page 87
                        v              v              v
Physical GPU HBM:   [ Page 4 ] ... [ Page 12 ] ... [ Page 87 ]

How PagedAttention Operates:

1. Physical Blocks: GPU VRAM is partitioned into fixed-size physical blocks (e.g., 16 or 32 tokens per block).
2. Block Table: A lightweight software page table maps a sequence's logical token indices to arbitrary physical blocks scattered across GPU HBM.
3. On-Demand Allocation: As a sequence generates new tokens, new blocks are allocated dynamically only when the current 16-token block fills. Internal fragmentation is capped at the size of a single block ($< 4%$).
4. Zero-Copy Page Sharing:
   • Parallel Sampling / Beam Search: Multiple candidate branches point to the same physical prompt blocks.
   • Prefix Caching: System prompts shared across thousands of users (e.g., code assist or customer support templates) reside in physical memory once and are shared read-only across all requests via Copy-on-Write (CoW).

---

Continuous Batching (Cellular Iteration-Level Scheduling)

Traditional deep learning batching batches requests at the request level: the batch runs until the longest sequence completes. If Request A finishes in 20 tokens and Request B runs for 2,000 tokens, Request A's slot sits idle (filled with padding tokens) for 1,980 iterations!

Continuous Batching (pioneered by Orca and implemented in vLLM/TRT-LLM) re-evaluates the batch at every single iteration:

• The moment a request outputs the End-of-Sequence (<EOS>) token, its physical KV blocks are instantly freed.
• A newly queued request enters the running batch at the very next iteration.
• GPU tensor cores remain 100% saturated without artificial padding bubbles.

---

Chunked Prefill & Speculative Decoding

Two cutting-edge optimizations dominate modern inference architectures:

1. Chunked Prefill: Long prompt prefills take tens of milliseconds, starving running decode steps and introducing huge Time Between Tokens (TBT) jitter. By slicing prefill tokens into 512-token chunks and interleaving them with decode steps, systems maintain low latency SLAs.
2. Speculative Decoding: A small, fast "draft model" (e.g., 1B parameter) generates 4-5 speculative tokens quickly. The large target model (70B) verifies all 5 tokens in a single parallel forward pass, achieving $2\times$ to $3\times$ speedups without any loss in output accuracy.

Mastering these systems engineering principles allows software engineers to deploy massive models at fractional infrastructure cost.