Production Serving Architectures & PagedAttention (vLLM)

Analyze memory-bandwidth bottlenecks in LLM inference, derive the KV cache memory formula, explain PagedAttention virtual memory block allocation, and implement continuous batching in Python and Rust.

06. Multimodal AI, Graph Neural Networks & Production MLOps12 min read

Deploying Large Language Models in production environments presents hardware challenges that differ fundamentally from training. While training is dominated by massive parallel matrix multiplications across thousands of accelerators, real-time LLM inference is severely constrained by memory bandwidth and GPU Video RAM (VRAM) fragmentation.

The seminal innovation that transformed production LLM serving throughput is PagedAttention, developed by Woosuk Kwon et al. at UC Berkeley (vLLM).

PagedAttention EfficiencyHigh-Throughput LLMOps
Memory Waste: 60-80% -> < 4%

Paged virtual memory paging eliminates internal and external KV cache fragmentation, dramatically increasing serving throughput.

1. The LLM Serving Bottleneck: Prefill vs Decode#

Generative Transformer inference splits into two computationally distinct phases:

Architecture Diagram
Rendering diagram...

The Roofline Model and Arithmetic Intensity#

Arithmetic Intensity measures the ratio of floating-point operations performed per byte of memory transferred:

I=Floating Point Operations (FLOPs)Memory Access (Bytes)I = \frac{\text{Floating Point Operations (FLOPs)}}{\text{Memory Access (Bytes)}}
  • During Prefill: Matrix dimensions are large (B×S×dB \times S \times d). The arithmetic intensity is high (I≫100 FLOPs/byteI \gg 100 \text{ FLOPs/byte}), fully saturating the GPU's Tensor Cores.
  • During Decode: For batch size B=1B=1, generating a single token requires loading the entire model parameter tensor WW (e.g. 140 GB for a 70B parameter FP16 model) just to compute a single vector-matrix product (2×70×1092 \times 70\times 10^9 FLOPs). The arithmetic intensity is approximately 1 FLOP/byte1 \text{ FLOP/byte}. The GPU compute cores sit idle >90%>90\% of the time waiting for memory transfer from GPU High Bandwidth Memory (HBM).

2. Mathematical Derivation of the KV Cache Footprint#

To prevent recomputing Key and Value representations for historical tokens at every auto-regressive step, inference engines maintain a Key-Value (KV) Cache.

Derivation of Per-Token KV Memory#

For each token in the context:

  • One Key vector per attention head: dheadd_{\text{head}} elements.
  • One Value vector per attention head: dheadd_{\text{head}} elements.
  • Across nlayersn_{\text{layers}} layers and nkv_headsn_{\text{kv\_heads}} key-value heads.

The total memory consumed by a single token in bytes is:

Memtoken=2×nlayers×nkv_heads×dhead×Pbytes\text{Mem}_{\text{token}} = 2 \times n_{\text{layers}} \times n_{\text{kv\_heads}} \times d_{\text{head}} \times P_{\text{bytes}}

where PbytesP_{\text{bytes}} is the precision in bytes (e.g., 2 bytes for FP16/BF16, 1 byte for INT8).

Real-World Example: LLaMA-3-70B#

LLaMA-3-70B uses Grouped-Query Attention (GQA):

  • nlayers=80n_{\text{layers}} = 80
  • nkv_heads=8n_{\text{kv\_heads}} = 8
  • dhead=128d_{\text{head}} = 128
  • Precision: 16-bit float (Pbytes=2P_{\text{bytes}} = 2)
Memtoken=2×80×8×128×2=327,680 bytes≈320 KB / token\text{Mem}_{\text{token}} = 2 \times 80 \times 8 \times 128 \times 2 = 327,680 \text{ bytes} \approx 320 \text{ KB / token}

For a concurrent serving batch of B=64B = 64 requests, each with a context length of L=4,096L = 4,096 tokens:

Total KV Cache=64×4096×320 KB=83,886,080 KB≈80.0 GB\text{Total KV Cache} = 64 \times 4096 \times 320 \text{ KB} = 83,886,080 \text{ KB} \approx 80.0 \text{ GB}

The KV cache alone exceeds the total VRAM capacity of an entire NVIDIA A100 (80GB) GPU.

3. The Memory Fragmentation Crisis in Naive Serving#

In naive inference frameworks, memory for the KV cache must be allocated as a single, contiguous physical tensor in GPU memory. Because the final sequence length of a user prompt is unknown in advance, engines must pre-allocate memory corresponding to the maximum possible sequence length (e.g. Lmax⁡=4,096L_{\max} = 4,096):

Architecture Diagram
Rendering diagram...

This causes two severe forms of fragmentation:

  1. Internal Fragmentation: If a request terminates after generating 50 tokens, the reserved memory for the remaining 3,946 tokens is completely wasted.
  2. External Fragmentation: Sequences finish at different times, leaving scattered holes in physical memory that cannot accommodate larger contiguous requests.

In practice, naive systems waste 60% to 80% of available GPU memory, drastically limiting concurrent batch sizes.

4. The PagedAttention Virtual Memory Architecture#

Inspired by classical operating system virtual memory and OS paging, PagedAttention partitions the KV cache into fixed-size physical blocks (e.g. 16 tokens per block).

The continuous sequence is broken down into Logical Blocks, which are mapped to non-contiguous Physical Blocks distributed arbitrarily across GPU memory:

Architecture Diagram
Rendering diagram...

Key Architectural Benefits#

  1. Near-Zero Memory Waste: Blocks are allocated strictly on-demand as tokens are generated. Only the final active block of a sequence has unused token slots, reducing memory waste to <4%< 4\%.
  2. Zero-Copy Branching (Copy-on-Write): In parallel sampling or beam search, multiple candidate generations share identical prompt prefixes. Instead of duplicating the prompt KV cache, both candidates point to the same physical blocks. Only when candidate sequences branch does the manager allocate new independent blocks.
  3. Prefix Caching: System instructions shared across millions of user queries (such as system prompt guardrails) are cached once and referenced across all incoming requests without duplicating VRAM.

5. Continuous (In-Flight) Batching#

Traditional static batching processes a batch until every request completes. If one request generates 1,000 tokens while others generate 10 tokens, the entire batch waits idly:

Architecture Diagram
Rendering diagram...

Continuous (Iteration-Level) Batching operates at the step level:

  • As soon as a sequence generates an <EOS> (End-of-Sequence) token, its memory blocks are immediately recycled back to the free pool.
  • A new waiting request from the queue enters the batch in the next decode step without waiting for existing requests to finish.

6. From-Scratch Implementation: PagedAttention Block Manager & Engine#

We build a complete zero-dependency Paged Block Allocator, Block Table mapper, and multi-tenant continuous request scheduler from first principles in Python and Rust.

Python Implementation (Pure Standard Library)#

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
""" PagedAttention Block Manager & Continuous Batching Scheduler from scratch using only the Python standard library. """ from typing import List, Dict, Optional BLOCK_SIZE = 4 # Each block stores 4 tokens class PhysicalBlock: def __init__(self, block_id: int): self.block_id = block_id self.ref_count = 0 self.tokens: List[int] = [] def is_full(self) -> bool: return len(self.tokens) >= BLOCK_SIZE def append_token(self, token: int): assert not self.is_full(), "Cannot append to full block" self.tokens.append(token) class BlockManager: """Manages the pool of physical GPU memory blocks.""" def __init__(self, num_blocks: int): self.blocks = [PhysicalBlock(i) for i in range(num_blocks)] self.free_pool: List[int] = list(range(num_blocks)) def allocate(self) -> int: if not self.free_pool: raise MemoryError("GPU Out of Memory: No free blocks remaining.") block_id = self.free_pool.pop(0) self.blocks[block_id].ref_count = 1 self.blocks[block_id].tokens.clear() return block_id def free(self, block_id: int): block = self.blocks[block_id] block.ref_count -= 1 if block.ref_count == 0: block.tokens.clear() self.free_pool.append(block_id) class Sequence: def __init__(self, seq_id: int, prompt_tokens: List[int]): self.seq_id = seq_id self.tokens = prompt_tokens[:] self.block_table: List[int] = [] # Logical block index -> physical block ID self.is_finished = False class PagedServingEngine: def __init__(self, total_blocks: int = 16): self.manager = BlockManager(total_blocks) self.active_sequences: Dict[int, Sequence] = {} def add_request(self, seq_id: int, prompt: List[int]): seq = Sequence(seq_id, prompt) # Allocate blocks for initial prompt for token in prompt: self._append_token_to_seq(seq, token) self.active_sequences[seq_id] = seq print(f"Request {seq_id} admitted. Allocated blocks: {seq.block_table}") def _append_token_to_seq(self, seq: Sequence, token: int): # Check if we need a new physical block if not seq.block_table or self.manager.blocks[seq.block_table[-1]].is_full(): new_block_id = self.manager.allocate() seq.block_table.append(new_block_id) last_block = self.manager.blocks[seq.block_table[-1]] last_block.append_token(token) def step(self): """Simulates one iteration of continuous decode generation.""" finished_ids = [] for seq_id, seq in self.active_sequences.items(): if seq.is_finished: continue # Generate next simulated token next_token = 100 + len(seq.tokens) seq.tokens.append(next_token) self._append_token_to_seq(seq, next_token) # Simulated termination condition if len(seq.tokens) >= 10: seq.is_finished = True finished_ids.append(seq_id) # Free finished requests immediately (Continuous Batching) for seq_id in finished_ids: seq = self.active_sequences.pop(seq_id) for b_id in seq.block_table: self.manager.free(b_id) print(f"Request {seq_id} completed. Recycled {len(seq.block_table)} blocks. Free pool: {len(self.manager.free_pool)}") if __name__ == "__main__": engine = PagedServingEngine(total_blocks=8) # Admit requests with varying prompt lengths engine.add_request(seq_id=1, prompt=[1, 2, 3]) engine.add_request(seq_id=2, prompt=[10, 20]) # Run execution iterations for it in range(8): if not engine.active_sequences: break print(f"\n--- Iteration {it + 1} ---") engine.step()

Rust Implementation (Pure Standard Library)#

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
// PagedAttention Block Manager and Iteration Scheduler from scratch in pure Rust std. const BLOCK_SIZE: usize = 4; struct PhysicalBlock { block_id: usize, tokens: Vec<usize>, ref_count: usize, } impl PhysicalBlock { fn new(block_id: usize) -> Self { Self { block_id, tokens: Vec::with_capacity(BLOCK_SIZE), ref_count: 0, } } fn is_full(&self) -> bool { self.tokens.len() >= BLOCK_SIZE } } struct BlockManager { blocks: Vec<PhysicalBlock>, free_pool: Vec<usize>, } impl BlockManager { fn new(num_blocks: usize) -> Self { let blocks = (0..num_blocks).map(PhysicalBlock::new).collect(); let free_pool = (0..num_blocks).collect(); Self { blocks, free_pool } } fn allocate(&mut self) -> Result<usize, &'static str> { if self.free_pool.is_empty() { return Err("GPU Out of Memory: No free blocks remaining"); } let b_id = self.free_pool.remove(0); self.blocks[b_id].ref_count = 1; self.blocks[b_id].tokens.clear(); Ok(b_id) } fn free(&mut self, block_id: usize) { if self.blocks[block_id].ref_count > 0 { self.blocks[block_id].ref_count -= 1; if self.blocks[block_id].ref_count == 0 { self.blocks[block_id].tokens.clear(); self.free_pool.push(block_id); } } } } struct Sequence { seq_id: usize, tokens: Vec<usize>, block_table: Vec<usize>, is_finished: bool, } fn main() { let mut manager = BlockManager::new(8); let mut sequences = Vec::new(); // Initialize sequence 1 with 3 prompt tokens let mut seq1 = Sequence { seq_id: 1, tokens: vec![101, 102, 103], block_table: Vec::new(), is_finished: false, }; // Allocate blocks for sequence 1 let b0 = manager.allocate().unwrap(); seq1.block_table.push(b0); for &t in &seq1.tokens { manager.blocks[b0].tokens.push(t); } sequences.push(seq1); println!("Sequence 1 admitted. Physical blocks: {:?}", sequences[0].block_table); println!("Remaining free blocks: {}", manager.free_pool.len()); // Append 2 tokens, triggering dynamic block allocation across the boundary for new_token in [104, 105] { let current_block_id = *sequences[0].block_table.last().unwrap(); if manager.blocks[current_block_id].is_full() { let next_b = manager.allocate().unwrap(); sequences[0].block_table.push(next_b); manager.blocks[next_b].tokens.push(new_token); } else { manager.blocks[current_block_id].tokens.push(new_token); } sequences[0].tokens.push(new_token); } println!("Appended tokens. Updated physical blocks: {:?}", sequences[0].block_table); println!("Free blocks available in GPU pool: {}", manager.free_pool.len()); }

7. Knowledge Check#

Concept check

Question

question: "Why is the decode phase of LLM inference characterized as memory-bandwidth-bound rather than compute-bound?"

options:

  • "(Recommended) Because generating each single token requires streaming all model weights and KV cache tensors from GPU HBM to SRAM with an arithmetic intensity of approximately 1 FLOP/byte"
  • "Because GPUs do not possess hardware matrix multiplication tensor cores"
  • "Because transformer attention operates sequentially without any matrix operations"
  • "Because tokenization takes more time than forward neural network passes"

explanation: "In the decode phase, tokens are generated one at a time. The arithmetic intensity is near 1 FLOP/byte because hundreds of gigabytes of weights and KV cache must be loaded from HBM to SRAM for tiny vector operations, leaving Tensor Cores idle most of the time."

:::

Concept check

Question

question: "How does PagedAttention eliminate memory fragmentation in GPU VRAM?"

options:

  • "(Recommended) By partitioning the continuous KV cache into fixed-size physical blocks and mapping them through a block table, allocating memory strictly on-demand"
  • "By compiling all weights into 1-bit binary representations"
  • "By forcing all sequences to have exactly 128 tokens"
  • "By executing attention on the CPU instead of the GPU"

explanation: "PagedAttention adapts virtual memory paging to LLM inference. Logical blocks of tokens are dynamically mapped to non-contiguous physical blocks via a block table, eliminating contiguous reservation waste and dropping fragmentation to under 4%."

:::