LLM Inference Architecture: PagedAttention, KV-Cache & Speculative Decoding

Deploying Large Language Models (LLMs) in high-concurrency production environments presents a fundamentally different systems engineering challenge than model pre-training. Whereas training workloads execute massive, predictable matrix-matrix multiplications across complete token sequences known ahead of time, autoregressive inference serving is stateful, sequential, dynamic in sequence length, and severely bottlenecked by High-Bandwidth Memory (HBM) capacity and bandwidth.

Without specialized systems-level memory management and kernel scheduling, modern accelerator clusters capable of hundreds of TeraFLOPS routinely suffer below 15% compute utilization while wasting 60% to 80% of their HBM footprint to memory fragmentation. This foundational architectural guide dissects the complete stack of modern LLM inference engines: Roofline hardware bounds, exact Key-Value Cache (KV-Cache) sizing equations, PagedAttention virtual memory block tables, IO-aware FlashAttention tiling, iteration-level Continuous Batching, and lossless Speculative Decoding.

1. Two-Phase Execution Mechanics: Compute-Bound Prefill vs. Memory-Bound Decode

Every inference request processed by a decoder-only Transformer bifurcates into two distinct execution phases governed by opposite hardware constraints:

Under the Roofline performance model, the boundary separating memory-bound execution from compute-bound saturation is defined by the accelerator's Operational Intensity Ridge Point, the ratio of peak floating-point throughput (FLOPS) to peak memory bandwidth (Bytes/second):

Arithmetic Intensity (I) = Total Floating-Point Operations (FLOPs) / Total HBM Bytes Transferred

Hardware Ridge Point (I_ridge) = Peak Accelerator Compute (FLOPS) / HBM Bandwidth (Bytes/s)

Reference Calculation for H100 SXM5 (FP16 Tensor Core):
- Peak FP16 Tensor Throughput : 989 TFLOPS
- HBM3 Memory Bandwidth       : 3.35 TB/s
- I_ridge                     : 989 / 3.35 ≈ 295 FLOPs per Byte

Autoregressive Decode Analysis (Batch Size = 1, FP16 Weights = 2 Bytes/param):
- Every parameter performs 2 FLOPs (1 multiply, 1 accumulate) while reading 2 Bytes from HBM.
- Decode Arithmetic Intensity (B=1) ≈ 2 FLOPs / 2 Bytes = 1 FLOP/Byte (295x below compute saturation!).

This 295-fold gap between single-stream decode arithmetic intensity (1 FLOP/Byte) and hardware saturation (~295 FLOPs/Byte) explains why single-user decode latency is dictated almost entirely by HBM memory bandwidth divided by active model weight and KV-Cache byte footprints, rather than raw TFLOPS ratings.

2. Exact KV-Cache Memory Calculus: Why Context Scaling Exhausts HBM

To avoid recomputing self-attention keys and values for all preceding tokens at every autoregressive step, which would impose quadratic O(N^2) per-step complexity and cubic O(N^3) total generation complexity, inference engines cache the projected Key (K) and Value (V) tensors across all transformer layers in HBM. This stateful buffer is the Key-Value Cache (KV-Cache).

While caching reduces per-step attention compute to linear O(N) complexity, it shifts the bottleneck squarely onto memory capacity and memory bandwidth. The exact KV-Cache footprint scales linearly with layer count, KV head count, head dimension, sequence length, and concurrent batch size:

KV_Bytes_Per_Token = 2 (K and V) * n_layers * n_kv_heads * d_head * bytes_per_element
Total_KV_Cache (GB) = (KV_Bytes_Per_Token * seq_len * batch_size) / 10^9

Architectural Comparison for a 70B Parameter Model (n_layers = 80, d_model = 8192, d_head = 128):

1. Standard Multi-Head Attention (MHA, n_kv_heads = 64, FP16 = 2 Bytes):
   - KV per Token = 2 * 80 * 64 * 128 * 2 = 2,621,440 Bytes (~2.62 MB per token!)
   - 64,000-token context, Batch = 16 -> 2.62 MB * 64,000 * 16 = 2,684 GB (2.68 Terabytes of HBM!)

2. Grouped-Query Attention (GQA-8, n_kv_heads = 8, FP16 = 2 Bytes):
   - KV per Token = 2 * 80 * 8 * 128 * 2 = 327,680 Bytes (~327.7 KB per token, 8x reduction)
   - 64,000-token context, Batch = 16 -> 327.7 KB * 64,000 * 16 = 335.5 GB of HBM

3. GQA-8 with FP8 / INT8 Quantized KV-Cache (n_kv_heads = 8, 1 Byte):
   - KV per Token = 2 * 80 * 8 * 128 * 1 = 163,840 Bytes (~163.8 KB per token, 16x reduction)
   - 64,000-token context, Batch = 16 -> 163.8 KB * 64,000 * 16 = 167.8 GB of HBM

As these equations prove, under long-context agentic workloads (32k to 128k+ tokens) and multi-request batching, the dynamic KV-Cache rapidly eclipses the static parameter weight footprint as the dominant consumer of HBM capacity.

3. Attention Architecture Evolution: MHA, MQA, GQA, and Multi-Head Latent Attention (MLA)

To compress KV-Cache memory footprints and reduce per-step HBM read traffic without degrading reasoning quality, transformer attention mechanisms have evolved across four distinct structural paradigms:

4. PagedAttention: Operating System Virtual Memory for Zero-Waste KV-Caches

In first-generation LLM serving systems, the KV-Cache tensor for each active request had to be allocated as a single contiguous block of physical HBM memory. Because the final output length of an autoregressive request cannot be predicted a priori, engines were forced to pre-allocate a contiguous tensor sized for the maximum possible sequence length (e.g., 32,768 or 131,072 tokens) as soon as a request arrived.

In their landmark SOSP 2023 paper introducing vLLM, Woosuk Kwon and colleagues at UC Berkeley demonstrated that contiguous KV allocation wastes 60% to 80% of HBM capacity across three failure modes:

PagedAttention eliminates all three sources of waste by adapting virtual memory paging and block tables directly from operating system kernel design. Instead of requiring contiguous physical tensors, PagedAttention partitions each request's logical KV-Cache into fixed-size KV Blocks (typically storing B = 16 or 32 tokens per block) that can reside anywhere in non-contiguous physical HBM.

Block Configuration: Fixed Block Size (B) = 16 tokens per physical HBM block

Logical Token Stream for Request A (42 Tokens Generated):
- Logical Block 0 (Tokens 0..15)  --> Block Table maps to --> Physical HBM Block #104 (Full: 16/16)
- Logical Block 1 (Tokens 16..31) --> Block Table maps to --> Physical HBM Block #019 (Full: 16/16)
- Logical Block 2 (Tokens 32..41) --> Block Table maps to --> Physical HBM Block #287 (Partial: 10/16)

Systems Engineering Guarantees:
1. Near-Zero Internal Fragmentation: Waste is strictly bounded to the final partially filled block (at most B - 1 = 15 tokens per sequence, reducing memory waste below 4%).
2. Copy-on-Write (CoW) Prefix Caching: When 50 concurrent agent requests share an identical 4,000-token system prompt or repository context, their logical block tables point to the exact same physical HBM blocks via reference counting, eliminating 98% of redundant prefill compute and KV storage!

5. IO-Aware Exact Attention (FlashAttention-2 & FlashAttention-3): SRAM Tiling and Online Softmax

Even with PagedAttention managing KV storage, executing the core self-attention equation Softmax((Q * K^T) / sqrt(d_k)) * V naively creates a catastrophic HBM bandwidth bottleneck. Standard attention implementations materialize the intermediate N x N attention score matrix in HBM, read it back to compute row-wise Softmax, write the normalized N x N matrix back to HBM, and read it a third time to multiply against V.

Tri Dao's FlashAttention architecture (spanning FlashAttention-2 and Hopper/Blackwell-optimized FlashAttention-3) reformulates exact self-attention around GPU memory hierarchy physics. While HBM provides 80 GB to 192 GB of capacity at ~3.35 TB/s, on-chip SRAM inside each Streaming Multiprocessor (SM) operates at roughly 19 TB/s, nearly an order of magnitude faster, despite holding only ~192 KB to 256 KB per SM.

6. Iteration-Level Scheduling (Continuous Batching) & Chunked Prefill

In legacy request-level batching (Static Batching), if an 8-request batch contains 7 short queries that finish after 50 tokens and 1 long-form code generation query requiring 2,000 tokens, the 7 completed batch slots sit idle waiting for the longest sequence to finish before the batch can return and admit new traffic.

Modern serving engines implement Continuous Batching (Iteration-Level Scheduling, pioneered by Orca). Rather than scheduling at the granularity of complete requests, the scheduler executes decisions at every single token iteration:

7. Speculative Decoding Mathematics: Breaking the Sequential Autoregressive Barrier

Because single-token autoregressive decoding is fundamentally memory-bandwidth-bound, reading hundreds of gigabytes of weights from HBM produces exactly one token per forward pass, leaving the GPU's massive arithmetic compute units largely idle. Yet in natural language and source code generation, a high fraction of tokens (syntactic boilerplate, closing brackets, common idioms) are highly predictable.

Formalized by Yaniv Leviathan et al. at Google Research, Speculative Decoding decouples token drafting from parallel target verification:

from vllm import LLM, SamplingParams

# Initialize production inference engine with PagedAttention, FP8 KV-Cache,
# Chunked Prefill, Automatic Prefix Caching, and Lossless Speculative Decoding
llm = LLM(
    model="meta-llama/Llama-3.3-70B-Instruct",
    speculative_model="meta-llama/Llama-3.2-1B-Instruct",
    num_speculative_tokens=5,
    tensor_parallel_size=4,
    kv_cache_dtype="fp8",
    enable_chunked_prefill=True,
    max_num_batched_tokens=2048,
    enable_prefix_caching=True,  # Copy-on-Write KV block sharing
    gpu_memory_utilization=0.92,
    block_size=16
)

sampling_params = SamplingParams(temperature=0.0, max_tokens=1024)
outputs = llm.generate(
    ["Design a Raft consensus partition recovery state machine in Rust:"],
    sampling_params
)
print(outputs[0].outputs[0].text)

Given an average per-token acceptance probability alpha (e.g., 0.80) and a draft lookahead length gamma = 5, the expected number of tokens generated per target forward pass follows the closed-form geometric series (1 - alpha^(gamma + 1)) / (1 - alpha) = (1 - 0.8^6) / (1 - 0.8) ≈ 3.69 tokens per step, translating into a 2.2x to 3.0x wall-clock decoding speedup.

8. Disaggregated Prefill-Decode Serving Clusters

At hyperscale production volume, co-locating compute-bound prefill and bandwidth-bound decode workloads on the exact same GPU instances creates unavoidable interference between Time-To-First-Token (TTFT) and Inter-Token Latency (ITL) Service Level Objectives (SLOs), while forcing both phases to share identical Tensor Parallelism (TP) and Pipeline Parallelism (PP) topologies.

State-of-the-art inference clusters resolve this tension through Disaggregated Prefill-Decode Serving: physical nodes are partitioned into dedicated Prefill Pools (optimized for high-FLOPS compute saturation and low TTFT) and Decode Pools (optimized for massive HBM capacity, large batch sizes, and low ITL). Once a prefill node computes the initial PagedAttention KV blocks and emits the first token, the non-contiguous KV pages are streamed asynchronously via GPUDirect RDMA (InfiniBand or RoCEv2) directly into the HBM block table of a decode worker, achieving independent scaling across both stages of LLM execution.

“The defining breakthroughs in modern AI inference do not come from brute-force silicon scaling alone; they come from treating HBM and SRAM as a virtualized operating system through PagedAttention, IO-aware FlashAttention kernels, and parallel Speculative Decoding.”