Skip to content

LLM Inference Engines & Memory Management: An Architectural Deep Dive

Inference for Large Language Models (LLMs) is fundamentally a systems engineering problem constrained by memory bandwidth, compute density, and memory capacity. To achieve high throughput and low latency, inference engines must aggressively optimize memory allocation and scheduling at the iteration level.

This document dissects the core mechanisms powering state-of-the-art inference engines, exposing the hardware constraints and scheduling algorithms that dictate their design.

1. The Roofline Model: Prefill vs. Decode

LLM generation consists of two distinct phases, each hitting different architectural limits described by the Roofline Model, which correlates computational performance (FLOP/s) with operational intensity (FLOPs/byte).

  • Prefill Phase (Prompt Processing): The engine processes the entire input prompt simultaneously. This phase involves dense matrix multiplications (GEMMs) with large batch sizes (where batch size \(B\) effectively equals the prompt length \(L\)). The operational intensity is high because weights are loaded once and multiplied across many tokens. Consequently, the Prefill phase is compute-bound. Performance is dictated by the tensor core FLOP/s of the GPU.
  • Decode Phase (Token Generation): Tokens are generated auto-regressively, one by one. For each generated token, the entire model weight matrix must be loaded from HBM to SRAM to perform a matrix-vector multiplication (GEMV). The operational intensity is exceptionally low (e.g., 2 FLOPs per byte of weight loaded). Consequently, the Decode phase is severely memory-bandwidth bound. Performance is bottlenecked by the HBM bandwidth (e.g., ~3.2 TB/s on H100), not the compute units.

Optimization strategies inherently diverge based on this dichotomy: Prefill optimizations target compute efficiency (e.g., FlashAttention), while Decode optimizations target memory bandwidth utilization (e.g., Quantization, Speculative Decoding).

2. KV Cache Memory Footprint: Hardware Constraints

To avoid recomputing attention scores for past tokens during the Decode phase, inference engines cache the Key (K) and Value (V) tensors for all previously processed tokens. This KV Cache grows linearly with sequence length and batch size, dominating GPU memory capacity.

Exact Formula

For a single token in a Transformer layer, the KV cache size is: Size per token per layer = 2 (for K and V) * d_head * n_heads_kv * bytes_per_param

For the entire model and a sequence of length \(L\), the total KV cache memory footprint per request is:

\[ \text{Total KV Cache (Bytes)} = 2 \times b \times L \times n_{layers} \times n_{heads\_kv} \times d_{head} \times p_{bytes} \]

Where: * \(b\) = Batch size (number of sequences). * \(L\) = Sequence length (number of tokens). * \(n_{layers}\) = Number of transformer layers. * \(n_{heads\_kv}\) = Number of attention heads for Keys and Values (Note: in Multi-Query Attention (MQA) or Grouped-Query Attention (GQA), this is less than the number of query heads). * \(d_{head}\) = Dimension of each attention head. * \(p_{bytes}\) = Precision in bytes (e.g., 2 for FP16/BF16, 1 for INT8).

Example (Llama-3-70B, FP16): * \(n_{layers} = 80\) * \(n_{heads\_kv} = 8\) (GQA) * \(d_{head} = 128\) * \(p_{bytes} = 2\)

Memory per token = \(2 \times 80 \times 8 \times 128 \times 2 = 327,680\) bytes (~320 KB). For a context window of 8,192 tokens, a single request requires ~2.6 GB of KV Cache. A batch of 128 requests requires ~335 GB, exceeding the capacity of an 8xH100 node (80GB each) just for the cache, demonstrating why memory capacity limits batch size.

3. PagedAttention (vLLM): Virtualizing GPU Memory

Naïve KV cache allocation relies on contiguous memory blocks sized for the maximum possible sequence length. Because sequence lengths are unpredictable, this leads to massive internal fragmentation (over-provisioning) and external fragmentation (inability to fit new requests despite sufficient total free memory). Historically, memory waste in LLM serving reached up to 60-80%.

PagedAttention solves this by adopting operating system virtual memory concepts.

Mechanism

  1. Block Partitioning: The KV cache is divided into fixed-size "blocks" (pages), where each block contains the K and V vectors for a fixed number of tokens (e.g., 16 or 32 tokens).
  2. Virtual to Physical Mapping: Tokens are treated as logical memory, mapped to physical blocks via a Block Table.
  3. Non-Contiguous Allocation: Physical blocks do not need to be contiguous in GPU memory. When a sequence generates a new token, the engine writes to the current block. If the block is full, the engine dynamically allocates a new physical block and updates the Block Table.

Engineering Trade-offs & Implications

  • Zero Internal Fragmentation: Memory is allocated on-demand in small block increments, virtually eliminating wasted space reserved for ungenerated tokens.
  • Memory Sharing (Copy-on-Write): PagedAttention enables memory sharing at the block level for requests with identical prefixes (e.g., prompt templates, parallel sampling). Multiple logical blocks can point to the same physical block, with a reference count. Writes trigger a physical copy.
  • Overhead: Managing block tables introduces a slight scheduling overhead, and reading non-contiguous memory requires custom CUDA kernels that traverse the block table during the attention calculation, slightly degrading theoretical maximum L1/L2 cache locality compared to contiguous arrays.

4. Continuous Batching (Iteration-Level Scheduling)

Traditional DNN inference uses static batching: wait for \(N\) requests, execute them, and return. For LLMs, request lengths vary wildly. A batch is bound by the longest sequence; shorter sequences finish early, leaving GPU compute units idling (bubble time).

Continuous Batching (or iteration-level scheduling) resolves this by scheduling work at the granularity of a single generation step (iteration), rather than the entire request.

  1. State Machine: Requests exist in states: WAITING, RUNNING, SWAPPED, FINISHED.
  2. Dynamic Insertion/Eviction: After every single Decode step (producing one token for the batch), the scheduler evaluates the queue. Completed requests are immediately removed from the batch. New requests from the WAITING queue are immediately inserted into the batch.
  3. Preemption: If the batch size grows to the point where the required KV cache exceeds physical memory, the scheduler preempts active requests, either by swapping their KV cache to CPU RAM or by recomputing them later.

This maximizes batch size and keeps the GPU saturated, significantly improving throughput (requests/second) at the cost of a slightly more complex runtime and higher variance in Time-To-First-Token (TTFT).

5. Chunked Prefill & RadixAttention

Chunked Prefill

Mixing Prefill (compute-bound) and Decode (memory-bound) phases in continuous batching is problematic. A long Prefill request injected into an active Decode batch will stall the Decode tokens, causing high inter-token latency (jitter).

Chunked Prefill splits long Prefill prompts into smaller chunks (e.g., 512 tokens). * During an iteration, the scheduler can execute a 512-token chunk of a new prompt alongside the 1-token decode steps of existing requests. * This overlaps compute-bound work with memory-bound work, maximizing total GPU utilization while maintaining strict Service Level Agreements (SLAs) on inter-token latency.

RadixAttention (Prompt Caching)

In many production workloads (e.g., multi-turn chat, agentic loops), identical prompt prefixes are sent repeatedly. PagedAttention allows sharing, but only if the engine knows they are identical at request time.

RadixAttention (implemented in SGLang) structures the KV cache as a Radix Tree (Patricia trie). * The tree keys are sequences of tokens. The nodes map to physical PagedAttention blocks. * When a new request arrives, the engine performs a prefix match against the Radix Tree. * Matched prefixes are reused instantly (zero recomputation). Unmatched suffixes trigger normal Prefill. * Eviction policies (e.g., LRU) manage the Radix Tree when HBM is full, retaining frequently used prompt components (like system prompts or RAG context) in HBM indefinitely.