Skip to content

Memory Management

Jump to bottom
Richard Huang edited this page Aug 3, 2026 · 1 revision

Memory Management

Memory management decides where every tensor lives and for how long. In LLM inference this is not bookkeeping — it is the primary determinant of how many requests can be served concurrently and how long a context can grow before the system fails.

Decode is memory bound, so memory management is throughput management.

Objectives

  • Fit — the model, activations, and cache must fit in available memory
  • Maximize concurrency — memory freed is a request served
  • Avoid fragmentation — unusable free space is wasted capacity
  • Keep allocation cheap — allocation on the critical path costs latency
  • Degrade gracefully — running out of memory should throttle, not crash

Where It Fits

flowchart LR
    A[Graph Compiler] --> B[Memory Management]
    B --> C[Runtime]
    D[Batching and Scheduling] <--> B
    style B fill:#2d6a9f,color:#fff
Loading

The two-way arrow matters: the scheduler can only admit a request if memory exists for it, and memory availability changes as requests finish. These two subsystems are tightly coupled.

What Consumes Memory

Consumer Size Lifetime
Model weights Fixed, large Entire process
KV cache Grows per token, per request Duration of a request
Activations Transient, per step Within one forward pass
Workspace Kernel scratch space Within one kernel

Weights are a one-time cost. Activations are small in decode, since only one token is processed at a time. The KV cache is the variable that determines system capacity, and nearly all interesting memory management is about it.

The KV Cache

For each token processed, attention produces a key and a value vector per layer, which must be retained so future tokens can attend to them. Cache size scales as:

cache bytes ≈ 2 × layers × kv_heads × head_dim × seq_len × batch × bytes_per_element

Two properties make this hard:

  1. It grows unboundedly with generated length, and the final length is unknown when the request starts.
  2. It is read in full every decode step. Cache size doesn't just consume capacity — it directly sets decode speed.

The fragmentation problem

The naive approach reserves a contiguous block per request sized to the maximum possible length. This wastes catastrophically: a request that generates 100 tokens holds an allocation sized for thousands.

The standard solution is paged allocation — divide cache memory into fixed-size blocks, and give each request a block table mapping logical positions to physical blocks. Requests grow by acquiring blocks on demand.

Property Contiguous Paged
Internal waste Very high One partial block per request
External fragmentation Severe None — blocks are interchangeable
Sharing between requests Impossible Natural, via shared block references
Kernel complexity Simple addressing Requires indirection through a block table

The cost is real: kernels must gather through a block table rather than reading a contiguous span. The benefit — often several times more concurrent requests — is almost always worth it.

Sharing

Because paged blocks are referenced rather than owned, identical prefixes can share physical storage. When many requests share a common system prompt, this eliminates redundant copies entirely. Copy-on-write handles divergence: shared blocks stay shared until a request writes, at which point it gets its own copy.

Reducing Cache Size

Technique Mechanism Cost
Grouped-query attention Several query heads share one KV head Architectural; must be trained in
Cache quantization Store keys and values in fewer bits Small accuracy cost — see Quantization
Eviction Drop tokens judged unimportant Accuracy risk if the wrong tokens go
Windowing Retain only a recent window Loses long-range attention
Offloading Move cold cache to slower memory Bandwidth cost on recall

Eviction deserves care. The observation that attention concentrates on a small subset of tokens — a few early tokens plus recent ones — motivates dropping the rest. It works well on many workloads and fails badly on tasks needing specific mid-context recall. Treat any eviction policy as an accuracy-affecting change requiring measurement, not a free win.

Activation and Workspace Memory

Within a forward pass, intermediate tensors are allocated and freed constantly. Doing this through a general allocator on the critical path is too slow.

Standard approach: the Graph Compiler computes a static allocation plan offline. Knowing every tensor's lifetime, it assigns offsets in a pre-allocated arena so that non-overlapping lifetimes reuse the same space. Runtime allocation becomes a pointer offset — free.

This is a classic interval-packing problem: minimize peak memory subject to lifetime constraints. Fusion helps here too, since fused intermediates never need an allocation at all.

Metrics

Metric Why it matters
Memory utilization Fraction actually holding useful data
Fragmentation Free but unusable space
Max concurrent requests The capacity number that matters commercially
Cache hit rate on shared prefixes How much sharing is being exploited
Peak activation memory Determines headroom for cache
Allocation latency Should be negligible on the critical path
Preemption rate How often requests are evicted under pressure

Common Problems

Problem Root cause Fix
Low concurrency despite free memory Fragmentation from contiguous allocation Paged allocation
Out of memory mid-generation Cache grew past the reservation Admission control; preemption policy
Throughput collapses under load Thrashing between admission and eviction Hysteresis; reserve headroom
Shared prefixes duplicated No sharing mechanism Reference-counted blocks with copy-on-write
Allocation shows in the profile Dynamic allocation on the hot path Static plan; pre-allocated arena
Accuracy drops on long context Eviction discarding needed tokens Re-examine the policy; measure on recall tasks
Wasted space per request Block size too large Tune block size against typical lengths

Best Practices

  • Page the KV cache. The complexity cost is real and worth paying.
  • Pre-allocate everything at startup. Steady-state allocation should be near zero.
  • Plan activations statically. The compiler knows all lifetimes; use that.
  • Reserve headroom. Running at 100% capacity means the next request fails.
  • Make admission memory-aware. Accepting a request you can't finish wastes all the work done on it.
  • Share aggressively, copy lazily. Common prefixes are extremely common in practice.
  • Instrument capacity. Utilization, fragmentation, and preemption rate should be visible in production, not discovered during an incident.
  • Treat eviction as an accuracy change. Measure it like one.

See also: Batching and Scheduling · Quantization · Graph Compiler · Runtime · Inference Software

Clone this wiki locally