By the end of this lesson
  • Distinguish the compute and memory characteristics of prefill vs. decode phases
  • Derive the exact KV cache memory consumption formula for MHA, MQA, and GQA architectures
  • Explain how PagedAttention eliminates internal and external GPU memory fragmentation
  • Compare static batching with continuous (iteration-level) batching

1. The problem

Serving Large Language Models (LLMs) is fundamentally different from serving traditional web APIs or microservices.

In a conventional application:

  • A request arrives, the server queries a database or executes business logic, returns an HTTP response, and discards all request state.
  • Adding server instances linearly scales capacity because requests are stateless.

In an LLM serving system:

  • Autoregressive Generation: Generating a 500-token response requires 500 sequential forward passes through the neural network.
  • Growing Dynamic State: Every single token generated must attend to all previous tokens. Storing these intermediate states (the KV Cache) consumes gigabytes of expensive GPU High Bandwidth Memory (HBM).
  • Two Distinct Execution Regimes: The prompt ingestion phase and the token generation phase exhibit completely opposite hardware bottlenecks.

Treating an LLM server as a generic stateless worker leads to catastrophic GPU Out-Of-Memory (OOM) crashes, massive memory fragmentation, and poor throughput.


2. Mental model

Treat LLM serving like a high-speed printing press with a scratchpad:

[ Incoming Prompt ] ──► [ Stage 1: Prefill (Bulk Ingestion) ] ──► Writes Initial KV Cache
                                                                           │
                                                                           ▼
[ Emitted Stream  ] ◄── [ Stage 2: Decode (Step-by-Step) ] ◄── Reads & Appends KV Cache
  1. Prefill Phase (Ingestion): Like scanning an entire book into memory. The GPU processes thousands of prompt tokens concurrently in a single massive parallel operation.
  2. Decode Phase (Generation): Like typing one character at a time. The GPU must load the entire multi-billion parameter model from memory into its compute cores just to emit a single token, update its scratchpad (KV cache), and repeat.

3. Core concepts & metrics

Serving Latency Metrics

MetricFull NameDefinitionTarget Goal
TTFTTime To First TokenDuration from request arrival until the first token streams back to the user (dominated by Prefill). for interactive chatbots
ITL / TPOTInter-Token Latency / Time Per Output TokenTime elapsed between consecutive generated tokens (governed by Decode). (, faster than human reading speed)
ThroughputToken ThroughputTotal tokens processed and generated per second across all concurrent GPU streams.Maximize tokens/$ without violating TTFT/ITL SLOs

Hardware Bottlenecks: Compute-Bound vs. Memory-Bound

Using the Roofline Model, an accelerator’s operational ridge point is :

  • NVIDIA H100 SXM5 (80 GiB HBM3): ( only with 2:4 structured sparsity).
  • NVIDIA A100 SXM4 (80 GiB HBM2e): ( sparse).

Consequently:

  • Prefill is Compute-Bound: Matrix-matrix multiplications () across hundreds or thousands of prompt tokens yield arithmetic intensity , running Tensor Cores near peak FLOP/s.
  • Decode is Memory-Bandwidth Bound: Single-token matrix-vector/skinny-matrix multiplications () yield arithmetic intensity (batch size, typically ). Tensor Cores stall waiting for weights and KV blocks to stream across the HBM bus.

4. Internal workings

SYSTEM ARCHITECTURE BLUEPRINTDISTRIBUTED AI
LLM inference prefill and decode phases.
ARCHITECTURAL FLOW · STEP 1 OF 5

1. Prompt Ingestion & Tokenization

User prompt is tokenized and batched. All N input tokens are ingested simultaneously into GPU High Bandwidth Memory (HBM).

Architectural Invariant: Parallelism: Prompt tokens are known upfront, enabling fully parallel matrix operations across all input positions.
ARCHITECTURAL FLOW · STEP 2 OF 5

2. Prefill Phase: Compute-Bound Saturation

Model computes causal self-attention over the entire prompt sequence. Matrix-matrix multiplications (GEMM) achieve peak GPU Tensor Core utilization (MFU > 50%).

Architectural Invariant: Compute-Bound: Time to First Token (TTFT) is dominated by arithmetic intensity: FLOPs grow quadratically O(N^2) with sequence length.
ARCHITECTURAL FLOW · STEP 3 OF 5

3. KV Cache Allocation & Population

Keys and Values for all prompt tokens across all transformer layers are saved into the KV Cache in GPU VRAM to avoid redundant recomputation.

Architectural Invariant: Memory Footprint: KV Cache requires 2 * 2 * n_layers * n_heads * d_head bytes per token. For 70B models, this reaches ~1.3MB/token.
ARCHITECTURAL FLOW · STEP 4 OF 5

4. Decode Phase: Memory-Bandwidth-Bound Step

Only 1 token is processed at a time. The query vector attends to all cached historical KV vectors via matrix-vector multiplication (GEMV).

Architectural Invariant: Memory-Bound Bottleneck: Arithmetic intensity drops to ~1 FLOP/byte. Token generation speed is strictly capped by GPU HBM memory bandwidth.
ARCHITECTURAL FLOW · STEP 5 OF 5

5. Autoregressive Cache Append & Loopback

The newly sampled token appends its own K/V vector to the cache. The token loops back as input for step t+2 until an EOS (end-of-sequence) token is emitted.

Architectural Invariant: Iterative Generation: Inter-Token Latency (ITL) dictates user perceived generation speed (streaming response).

Why We Need the Key-Value (KV) Cache

In a Transformer decoder layer, self-attention computes:

When generating token :

  • The query vector represents the current token.
  • It must be multiplied against the key vectors and value vectors of all past tokens.
  • Without a cache, generating token would require recomputing and for all preceding tokens from scratch, resulting in an computational explosion.

The KV Cache solution: Store the and vectors for past tokens in GPU memory. At step , compute and for the new token only, append them to the cache, and compute attention against the cached vectors in operations.


Attention Architectures: MHA vs. MQA vs. GQA

The memory footprint of the KV cache depends heavily on the attention architecture:

Multi-Head Attention (MHA)      Multi-Query Attention (MQA)      Grouped-Query Attention (GQA)
      Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8             Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8             Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
      │  │  │  │  │  │  │  │             └──┬──┴──┬──┴──┬──┴──┘             └──┬──┘  └──┬──┘  └──┬──┘
      K1 K2 K3 K4 K5 K6 K7 K8                        K1                                K1       K2       K3
      V1 V2 V3 V4 V5 V6 V7 V8                        V1                                V1       V2       V3
(1 KV head per Query head)         (1 single KV head shared)           (Groups share KV heads, e.g. 8:1)
  • MHA (e.g. GPT-3): Number of KV heads equals number of Query heads (). Generates massive KV cache overhead.
  • MQA (e.g. Falcon): A single KV head is shared across all Query heads (). Reduces KV cache by , but can slightly degrade reasoning quality.
  • GQA (e.g. Llama 2/3, Mistral): Query heads are grouped (e.g. 8 query heads share 1 KV head). Delivers the memory efficiency of MQA with the representational capacity of MHA.

PagedAttention & Memory Management (vLLM)

Traditional serving engines suffered from severe memory fragmentation:

  1. Over-allocation: Servers pre-allocated contiguous memory blocks based on the model’s maximum context window (e.g. 8,192 tokens), even if the prompt only generated 50 tokens.
  2. External Fragmentation: Varying request lifetimes fragmented contiguous GPU memory space.

PagedAttention (Kwon et al., SOSP 2023) applies OS virtual memory paging to GPU tensors:

Logical KV Cache:   [ Block 0 ] ──► [ Block 1 ] ──► [ Block 2 ] ──► [ Block 3 ]
                           │               │               │               │
Block Table:          Physical #7     Physical #2    Physical #19     Physical #4
                           │               │               │               │
Physical GPU HBM:   [ P_Block 2 ] ... [ P_Block 4 ] ... [ P_Block 7 ] ... [ P_Block 19 ]
  • KV cache is divided into fixed-size physical blocks (typically 16 or 32 tokens).
  • A Block Table maps contiguous logical sequence positions to non-contiguous physical GPU blocks.
  • Memory is allocated dynamically on demand. Waste is reduced from to , allowing to higher concurrency on the same GPU hardware.

Static Batching vs. Continuous (Iteration-Level) Batching

Static Batching:
Seq 1: [Prompt][==== Output ====] [       IDLE PADDING       ]
Seq 2: [Prompt][================ Output ================]
Seq 3: [Prompt][= Output =]       [       IDLE PADDING       ]
                                  ▲
                         Batch wait bottleneck

Continuous Batching (Orca / vLLM):
Step 1: [Seq 1 Token][Seq 2 Token][Seq 3 Token]
Step 2: [Seq 1 Token][Seq 2 Token][Seq 3 Token]
Step 3: [Seq 1 Token][Seq 2 Token][Seq 4 (NEW) Token]  <-- Seq 3 finished; Seq 4 injected immediately!

Continuous batching schedules requests at the iteration level rather than the request level. As soon as a request emits an end-of-sequence (<EOS>) token, its memory blocks are reclaimed and a waiting request enters the execution step immediately.


5. Mathematics: Sizing the KV Cache

The Exact Formula

Where:

  • Factor : One matrix for Keys (), one matrix for Values ().
  • : Total transformer decoder layers.
  • : Number of KV attention heads.
  • : Dimension per head ().
  • : Bytes per parameter ( for FP16/BF16, for FP8).

For an entire active batch:


Complete Worked Example: Llama-3-70B

Let us calculate the exact memory requirements for hosting Llama-3-70B on an 8-GPU node of NVIDIA H100 SXM5 ( addressable HBM3, or decimal):

Model Parameters:

  • Total Parameters:
  • Layers ():
  • Query Heads (): ()
  • KV Heads (): (Grouped-Query Attention, query-to-KV ratio; contrast with MHA or MQA )
  • Head Dimension ():
  • Precision: 16-bit BF16 (; FP8 KV cache uses )

1. Model Weights Memory:

Across 8 GPUs using Tensor Parallelism ():

2. KV Cache per Token:

(Note: Under FP8 KV cache quantization at , this halves to .)

3. KV Cache Sizing for 64 Concurrent Streams at 4,096 Context Length:

4. Total GPU Node Budget (in consistent binary , matching nvidia-smi):

On a () H100 node, is available for the KV cache pool—supporting 385 concurrent streams even at full -token capacity, and 1,052 concurrent streams at an average active sequence length of tokens via PagedAttention!


6. Implementation: KV Cache Memory & Serving Simulator in TypeScript

export interface ModelConfig {
  name: string;
  parametersBillion: number;
  nLayers: number;
  nKvHeads: number;
  dHead: number;
  bytesPerElement: number; // 2 for FP16, 1 for FP8
}

export interface ServingClusterConfig {
  gpuCount: number;
  vramPerGpuGb: number;
  activationOverheadGbPerGpu: number;
}

export class KVCacheMemorySimulator {
  private readonly model: ModelConfig;

  constructor(model: ModelConfig) {
    this.model = model;
  }

  // KV cache memory per token in bytes: 2 * n_layers * n_kv_heads * d_head * bytes_per_element
  public getBytesPerToken(): number {
    return (
      2 *
      this.model.nLayers *
      this.model.nKvHeads *
      this.model.dHead *
      this.model.bytesPerElement
    );
  }

  public getWeightsMemoryGb(): number {
    const rawBytes = this.model.parametersBillion * 1e9 * this.model.bytesPerElement;
    return Math.round((rawBytes / (1024 ** 3)) * 100) / 100;
  }

  public evaluateClusterCapacity(
    cluster: ServingClusterConfig,
    avgContextTokens: number,
    maxContextTokens: number
  ): {
    kvCacheBytesPerToken: number;
    weightsGbTotal: number;
    usableKvMemoryGb: number;
    maxStreamsContiguous: number;
    maxStreamsPagedAttention: number;
    concurrencyGainPercent: number;
  } {
    const totalVramGb = cluster.gpuCount * cluster.vramPerGpuGb;
    const weightsGb = this.getWeightsMemoryGb();
    const activationOverheadGb = cluster.gpuCount * cluster.activationOverheadGbPerGpu;
    const usableKvMemoryGb = Math.max(0, totalVramGb - weightsGb - activationOverheadGb);

    const bytesPerToken = this.getBytesPerToken();
    const usableKvMemoryBytes = usableKvMemoryGb * (1024 ** 3);

    // Contiguous allocation pre-allocates for maxContextTokens
    const contiguousBytesPerStream = maxContextTokens * bytesPerToken;
    const maxStreamsContiguous = Math.floor(usableKvMemoryBytes / contiguousBytesPerStream);

    // PagedAttention allocates blocks on demand based on actual average context
    const pagedBytesPerStream = avgContextTokens * bytesPerToken;
    const maxStreamsPagedAttention = Math.floor(usableKvMemoryBytes / pagedBytesPerStream);

    const concurrencyGainPercent =
      maxStreamsContiguous > 0
        ? Math.round(((maxStreamsPagedAttention - maxStreamsContiguous) / maxStreamsContiguous) * 100)
        : 0;

    return {
      kvCacheBytesPerToken: bytesPerToken,
      weightsGbTotal: weightsGb,
      usableKvMemoryGb: Math.round(usableKvMemoryGb * 10) / 10,
      maxStreamsContiguous,
      maxStreamsPagedAttention,
      concurrencyGainPercent,
    };
  }
}

export function runDemo(): void {
  console.log('=== LLM Serving: KV Cache & PagedAttention Capacity Simulation ===');

  // Model: Llama-3-70B (GQA: 8 KV heads, 80 layers, d_head 128)
  const llama70bFp16: ModelConfig = {
    name: 'Llama-3-70B (FP16)',
    parametersBillion: 70,
    nLayers: 80,
    nKvHeads: 8,
    dHead: 128,
    bytesPerElement: 2,
  };

  const cluster8xH100: ServingClusterConfig = {
    gpuCount: 8,
    vramPerGpuGb: 80, // 640 GB total HBM3
    activationOverheadGbPerGpu: 3.5, // 28 GB total CUDA buffer
  };

  const simFp16 = new KVCacheMemorySimulator(llama70bFp16);
  const resFp16 = simFp16.evaluateClusterCapacity(cluster8xH100, 1500, 4096);

  console.log(`\n[Model: ${llama70bFp16.name} on 8x H100 (640 GB)]`);
  console.log(`  KV Cache per Token: ${(resFp16.kvCacheBytesPerToken / 1024).toFixed(1)} KB`);
  console.log(`  Weights Footprint: ${resFp16.weightsGbTotal} GB (across 8 GPUs: ${(resFp16.weightsGbTotal / 8).toFixed(1)} GB/GPU)`);
  console.log(`  Usable KV Cache Pool: ${resFp16.usableKvMemoryGb} GB`);
  console.log(`  Max Concurrency (Naive Contiguous): ${resFp16.maxStreamsContiguous} streams`);
  console.log(`  Max Concurrency (PagedAttention):   ${resFp16.maxStreamsPagedAttention} streams`);
  console.log(`  Throughput Increase from PagedAttention: +${resFp16.concurrencyGainPercent}%`);

  // Quantization comparison: FP8 KV Cache
  const llama70bFp8: ModelConfig = {
    ...llama70bFp16,
    name: 'Llama-3-70B (FP8 KV Cache)',
    bytesPerElement: 1,
  };
  const simFp8 = new KVCacheMemorySimulator(llama70bFp8);
  const resFp8 = simFp8.evaluateClusterCapacity(cluster8xH100, 1500, 4096);

  console.log(`\n[Quantization: FP8 KV Cache Impact]`);
  console.log(`  KV Cache per Token: ${(resFp8.kvCacheBytesPerToken / 1024).toFixed(1)} KB (50% reduction)`);
  console.log(`  Max Concurrency (PagedAttention + FP8): ${resFp8.maxStreamsPagedAttention} streams (2x serving capacity!)`);
}

runDemo();

Simulation Execution Output & State Transitions

=== LLM Serving: KV Cache & PagedAttention Capacity Simulation ===

[Model: Llama-3-70B (FP16) on 8x H100 (640 GB)]
  KV Cache per Token: 320.0 KB
  Weights Footprint: 130.39 GB (across 8 GPUs: 16.3 GB/GPU)
  Usable KV Cache Pool: 481.6 GB
  Max Concurrency (Naive Contiguous): 385 streams
  Max Concurrency (PagedAttention):   1052 streams
  Throughput Increase from PagedAttention: +173%

[Quantization: FP8 KV Cache Impact]
  KV Cache per Token: 160.0 KB (50% reduction)
  Max Concurrency (PagedAttention + FP8): 2104 streams (2x serving capacity!)

7. Practical applications

  1. Long-Context RAG (Retrieval-Augmented Generation): A (-token) context window prompt on Llama-3-70B consumes ( decimal) of KV cache for a single user session. Serving long-context models requires prompt caching (RadixAttention) to avoid re-computing and duplicating identical documentation prefixes across requests.
  2. Multi-Turn Agentic Loops: Agent workflows repeatedly call tools and feed results back into the prompt. Without KV cache prefix sharing, each tool step incurs full prefill latency.
  3. Speculative Decoding: A small “draft model” (e.g. Llama-3-8B or 1B) proposes candidate tokens sequentially; the target 70B model verifies all candidates in a single parallel forward pass, increasing decode speed by while preserving exact target distribution equivalence.

8. Design Alternatives & Inference Engine Trade-Off Comparison

Serving Architecture / EngineKV Cache ManagementPrefill vs. Decode Co-SchedulingOptimal Workload & Trade-Off
vLLM (v0.6+)PagedAttention (16/32-token blocks + APC hash prefix cache)Chunked Prefill (default enabled in v0.6+)General-purpose open-source multi-tenant serving; low CPU overhead in v0.6 async scheduler.
SGLang (v0.3+)RadixAttention (LRU radix tree over paged KV blocks)Chunked Prefill + Zero-Overhead SchedulerSuperior for multi-turn agents, tree-of-thought, and structured JSON decoding with high prefix reuse.
TensorRT-LLM (v0.12+)Paged KV Cache + FP8/INT4 KV quantizationIn-Flight Batching + Custom NVLink kernelsHighest raw tokens/sec on homogeneous NVIDIA H100/B200 fleets; requires ahead-of-time engine compilation.
Disaggregated Serving (Splitwise / DistServe)Separate Prefill & Decode GPU pools connected via RDMA KV transferZero prefill interference on decode GPUsGuarantees strict TTFT and TPOT SLAs simultaneously at large fleet scale, at the cost of RDMA KV transfer bandwidth.

9. Failure modes and limitations

  1. GPU Out-Of-Memory (OOM) during Decode: Unlike standard software where allocating memory throws a catchable exception, an uncaught CUDA OOM crash terminates the entire serving worker process, dropping all concurrent user sessions.
  2. Preemption and Thrashing: When GPU memory is exhausted during generation, the scheduler must preempt an active request:
    • Swap to CPU: Evicts KV blocks over PCIe to host RAM (adds significant latency when swapped back).
    • Recompute: Discards the generated KV cache and re-runs prefill when memory becomes available.
  3. Prefill Head-of-Line Blocking: A single massive 64K-token prompt prefill can saturate GPU tensor cores for several hundred milliseconds, starving active decode streams and creating noticeable spikes in inter-token latency (TPOT). Mitigation: Use Chunked Prefill (splitting the prompt into 512-token chunks executed alongside decode steps) or Disaggregated Prefill-Decode (running prefill and decode on separate dedicated GPU clusters).

10. Operational perspective

Production Monitoring Metrics

  • vllm:num_requests_waiting: Number of requests queued in the scheduler waiting for free memory blocks. A sustained backlog indicates the server is memory-saturated.
  • vllm:gpu_cache_usage_factor: Fraction of allocated GPU KV cache blocks in use. Should hover between .
  • vllm:num_preemptions_total: Cumulative count of preempted requests. Any value indicates severe memory pressure and degraded user experience.

11. Knowledge check

Complete the 4-question interactive assessment quiz at the end of this lesson to verify your understanding of prefill vs. decode Roofline bottlenecks, GQA KV-cache sizing, continuous batching, and PagedAttention.


12. Exercises

Exercise 1: KV Cache Sizing for Mistral-7B

  • Model Architecture:
    • Data Type: FP16 ()
  • Tasks:
    1. Calculate the KV cache memory size required for a single token in bytes and KiB ().
    2. A server handles 128 concurrent requests, each averaging 2,048 tokens. What is the total KV cache memory requirement in binary GiB and decimal GB?
    3. If we quantize the KV cache to FP8 ( per element), how much memory is saved?

Exercise 2: Roofline Arithmetic Intensity Crossover (Prefill vs. Decode on NVIDIA H100)

  • Given: An NVIDIA H100 SXM5 GPU delivers dense BF16 compute and HBM3 bandwidth ().
  • Tasks:
    1. Calculate the arithmetic intensity of a single linear weight projection ( per token over of BF16 weights) at batch size , , and .
    2. Determine the minimum effective token batch size required to transition the GPU from memory-bandwidth bound to compute bound.
View Reference Architecture & Worked Solutions (AI-C03-001)

Worked Solution — Exercise 1: KV Cache Sizing for Mistral-7B

  1. Single-Token KV Cache Memory (): (Note: Because Mistral-7B uses Grouped-Query Attention with instead of , its KV footprint is smaller than an equivalent Multi-Head Attention model at .)

  2. Total KV Cache for Concurrent Requests at Average Tokens:

    • Total active tokens in the batch:
    • Total KV cache memory in FP16:
    • Adding Mistral-7B’s FP16 model weights () and of CUDA activation workspace yields , fitting comfortably on a single A100 or H100 GPU with PagedAttention.
  3. Memory Saved via FP8 () KV Cache Quantization:

    • Halving from to reduces per-token KV storage to .
    • Total KV cache for tokens drops to:
    • Net Memory Saved: (, a reduction), enabling either higher concurrent request capacity ( streams) or longer average context windows ( tokens) on the exact same GPU.

Worked Solution — Exercise 2: Roofline Arithmetic Intensity Crossover on H100

  1. Arithmetic Intensity Across Batch Sizes:
    • Loading of BF16 weights from HBM3 once and performing across active tokens yields:
    • At (single-stream decode): ( of peak H100 Tensor Core compute).
    • At (moderate decode batch): ( of peak compute; still memory-bandwidth bound).
    • At (chunked prefill or high-concurrency batch): (; fully compute bound).
  2. Minimum Batch Size for Compute Saturation:

13. Summary and decision guide

Concept / TechnologyRule of Thumb
Prefill PhaseCompute-bound ( on H100); optimizes for Time to First Token (TTFT).
Decode PhaseMemory-bandwidth bound (); optimizes for Inter-Token Latency (ITL / TPOT).
GQA vs. MHAGQA ( ratio) reduces KV cache memory by compared to MHA with negligible quality loss.
PagedAttentionEliminates external memory fragmentation by allocating non-contiguous physical blocks (< 4% internal waste).
Continuous BatchingReplaces static batching; immediately schedules incoming requests upon sequence completion.
KV Cache QuantizationFP8 KV caching () halves KV memory footprint, doubling maximum serving concurrency.

14. Authoritative References & Version Context

INTERACTIVE ISLAND

GPU VRAM & PagedAttention KV Cache Simulator

Inspect how modern LLM inference engines (vLLM) eliminate memory fragmentation using OS-style virtual memory paging, block tables, and copy-on-write prefix caching.

PHYSICAL GPU VRAM BLOCK POOL (16 Blocks × 4 Tokens = 64 Max Tokens)
ACTIVE BATCH REQUESTS & LOGICAL-TO-PHYSICAL BLOCK TABLES
MASTERY GATE · STEP 83 OF 94

Knowledge check & lesson completion

Score 4 / 4 (100%) on this check to mark this lesson Completed ✓, advance your position on the Master Spine, and add these concepts to your Spaced-Repetition Review Deck.

1. What is the primary operational difference between the prefill and decode phases during LLM inference?
2. In Grouped-Query Attention (GQA), such as Llama-3-70B (80 layers, 64 query heads, 8 KV heads, head dimension 128), how much KV cache memory does a single token occupy in 16-bit (BF16/FP16) precision?
3. How does PagedAttention (introduced in vLLM) improve GPU memory utilization compared to traditional contiguous KV cache allocations?
4. Why does static batching perform poorly for interactive autoregressive LLM serving?
PERSONAL NOTEBOOK

Chapter Notes

Personal engineering insights and observations saved to this browser.

0 notes
Tip: You can also highlight text in the article above to quote directly.

Loading personal notes...