PagedAttention & Memory Footprint Calculations
Calculate the exact VRAM required for KV cache at batch size B and sequence length L. See why naive contiguous allocation wastes up to 80% of GPU memory.
- •Compute KV Cache bytes = 2 * 2 * n_layers * n_heads * d_head * seq_len * batch_size
- •Understand physical vs virtual memory blocks in vLLM PagedAttention
- •Mitigate memory fragmentation during continuous streaming
The Arithmetic Intensity Wall & KV Cache Formula
During the autoregressive decode phase, an LLM generates exactly one token per forward step. Each new token requires streaming the entire preceding KV cache across all 32-80 layers from High-Bandwidth Memory (HBM) into on-chip SRAM. The arithmetic intensity drops to less than 1 FLOP per byte transferred! At this point, GPU compute cores sit 90% idle, waiting on memory bus bandwidth.
KV Cache Memory Wall & Hardware Fit Calculator
Notice in the interactive simulator above: Increasing batch size amortizes model weight loading overhead, but causes KV cache memory to grow linearly until hardware VRAM is completely saturated.
PagedAttention: Eliminating Fragmentation via Virtual Memory Pages
Standard PyTorch allocations require contiguous memory chunks. Because generative output lengths are unknown in advance, naive systems pre-allocate memory for the maximum possible sequence length (e.g. 8k or 32k tokens), resulting in 60% to 80% wasted VRAM due to internal and external fragmentation. Inspired by OS virtual memory paging, vLLM's PagedAttention partitions the KV cache into fixed-size physical blocks (typically 16 tokens), indexing them through a dynamic Block Table.
PagedAttention Virtual-to-Physical Block Mapping
Logical token sequences are perceived by the transformer as contiguous, while the physical blocks reside anywhere across non-contiguous GPU HBM pages. Zero memory is pre-reserved, reclaiming up to 4x throughput capacity.
Production Block Allocator Implementation
Here is an end-to-end, runnable implementation of a PagedAttention memory manager that dynamically allocates and frees 16-token physical pages on GPU during continuous streaming.
import torch
from dataclasses import dataclass
from typing import Dict, List
BLOCK_SIZE = 16 # Tokens per page
class PagedKVCacheManager:
def __init__(self, num_blocks: int, num_layers: int, num_heads: int, head_dim: int, device="cuda"):
self.num_blocks = num_blocks
self.device = device
# Pre-allocate non-contiguous physical page pool: [num_blocks, 2, num_layers, num_heads, BLOCK_SIZE, head_dim]
self.kv_storage = torch.zeros(
(num_blocks, 2, num_layers, num_heads, BLOCK_SIZE, head_dim),
dtype=torch.float16,
device=device
)
self.free_blocks: List[int] = list(range(num_blocks))
self.block_tables: Dict[str, List[int]] = {}
def allocate_page_for_request(self, request_id: str) -> int:
if not self.free_blocks:
raise MemoryError("CUDA OOM: KV-Cache Physical Page Pool Exhausted!")
block_id = self.free_blocks.pop(0)
if request_id not in self.block_tables:
self.block_tables[request_id] = []
self.block_tables[request_id].append(block_id)
return block_id
def free_request(self, request_id: str):
if request_id in self.block_tables:
for b_id in self.block_tables[request_id]:
self.free_blocks.append(b_id)
del self.block_tables[request_id]
# Initialize 10,000 physical blocks (160,000 tokens)
manager = PagedKVCacheManager(num_blocks=10000, num_layers=32, num_heads=8, head_dim=128, device="cpu")
req_block = manager.allocate_page_for_request("req_user_42")
print(f"Allocated physical block ID: {req_block} for req_user_42")