DataArena
Inference Acceleration: KV Caching, PagedAttention & Quantization
ML Training, Inference & Safety • Inference Infrastructure10 min read

Inference Acceleration: KV Caching, PagedAttention & Quantization

High-throughput LLM serving mechanics: vLLM, memory virtualization, speculative decoding, and FP8/INT4 precision.

Practice in Prompt LabTest in AI Interview

Core Conceptual Brief

Serving LLMs in production is memory-bandwidth bound. Discover Key-Value (KV) caching, PagedAttention memory virtualization, speculative decoding, and quantization trade-offs.

Theoretical & Empirical Architecture

During LLM generation, computing attention for each newly generated token requires the Keys and Values of all preceding tokens. Recomputing them from scratch on every step would result in O(N2)O(N^2) redundant compute per generation step.

Key-Value (KV) Caching stores previous attention representations in GPU VRAM, turning autoregressive decoding into an O(1)O(1) compute operation per step—but introducing massive memory bandwidth bottlenecks.

Under the Hood: Key Engineering Mechanics

1PagedAttention (vLLM)

In traditional serving frameworks, memory for a request's KV cache had to be allocated contiguously based on maximum possible sequence length (e.g. reserving 8k tokens even if the request only needed 200). This caused 60-80% VRAM fragmentation. PagedAttention borrows virtual memory paging from operating systems, allocating KV cache in small fixed-size non-contiguous memory blocks (pages), boosting serving throughput by 2-4x.

2Speculative Decoding

A small, ultra-fast draft model (e.g. 1B parameters) speculatively generates KK tokens. The large target model (e.g. 70B parameters) evaluates all KK tokens simultaneously in a single forward pass using modified attention masking. Any accepted tokens are retained for free, accelerating generation speed by 2x to 3x with zero loss in mathematical output quality.

3Quantization: FP16 to FP8 and INT4 (AWQ/GPTQ)

Foundation models are trained in FP16 or BF16 (16 bits per weight). Quantizing weights to FP8 or INT4 reduces memory size by 2x to 4x, allowing a 70B parameter model that previously required two 80GB A100 GPUs to fit on a single consumer or cloud GPU with negligible degradation on standard benchmarks.

Calculating KV Cache Memory Footprint

def calculate_kv_cache_bytes(
    num_layers: int, 
    num_heads: int, 
    head_dim: int, 
    seq_len: int, 
    batch_size: int, 
    precision_bytes: int = 2  # FP16 = 2 bytes
) -> int:
    # Key and Value vectors for each layer and head
    bytes_per_token = 2 * num_layers * num_heads * head_dim * precision_bytes
    total_bytes = bytes_per_token * seq_len * batch_size
    return total_bytes

# Example: Llama-3-70B (80 layers, 8 KV heads in GQA, head_dim 128)
# At 8192 context length and batch size 16:
kv_bytes = calculate_kv_cache_bytes(80, 8, 128, 8192, 16, precision_bytes=2)
print(f"KV Cache VRAM Required: {kv_bytes / (1024**3):.2f} GB")
# Output: ~40.96 GB VRAM dedicated purely to KV cache!

Production Guardrails & Takeaways

  • LLM generation is memory-bandwidth bound; KV caching saves compute at the expense of heavy VRAM consumption.
  • PagedAttention eliminates memory fragmentation by allocating non-contiguous KV pages.
  • Speculative decoding accelerates generation by verifying draft tokens in parallel forward passes.

Real Staff AI Systems Interview Question

Q: Why does an LLM's generation speed (tokens per second) remain almost constant whether generating 1 token or 50 tokens in batch, but slows down dramatically when batch size increases?

Staff Answer:

Autoregressive generation is memory-bandwidth bound, not compute-bound: the GPU must transfer billions of model weights from High Bandwidth Memory (HBM) to compute SRAM on every single token step. With small batch sizes, compute cores sit idle waiting for memory transfer. As batch size increases, arithmetic intensity rises until memory bandwidth is saturated, causing throughput per stream to drop.

Inference Acceleration: KV Caching, PagedAttention & Quantization | AI Systems Engineering Guide · DataArena