← AI

KV Cache Mechanics & Inference Memory Layout

ai autoregressive decode · KV cache · paged memory Round 4 ✓ math ✓ visual ✓

The optimization that turns the O(n²) attention wall into an O(n) memory wall: cache every token's Keys and Values once, and autoregressive generation stops recomputing the past — but the cache itself becomes the scarcest resource in LLM serving.

What · How · Why

What it is

During text generation a Transformer emits one token at a time, and each new token attends to all previous tokens. The KV cache stores the Key and Value vectors already computed for every past token so they are never recomputed. It converts generation from repeatedly re-reading the whole prompt into a single incremental step per token — at the cost of a memory buffer that grows linearly with sequence length.

How it works

Inference has two phases. Prefill runs the whole prompt through the model in one parallel pass and writes its K/V into the cache. Decode then generates token-by-token: for each new token only its query is computed, dotted against the cached keys, and used to blend the cached values. The cache is appended to each step. The compute per token is tiny; the bottleneck is reading the ever-growing cache out of memory.

Why it matters

Without the cache, generating token \(t\) would redo the attention of all \(t\) tokens — turning generation into \(O(n^3)\) work. With it, decode is memory-bandwidth bound, not compute-bound. That single fact drives the entire economics of LLM serving: batch size, context length, and hardware choice all reduce to "how much KV cache fits and how fast can you stream it."

Round 1 — Mental Model

Imagine writing a novel where, before adding each new sentence, you had to re-read the entire book from page one. That is attention without a cache. The KV cache is the running set of margin notes: as you read each page once, you jot down its two summaries — a Key ("what this page is about") and a Value ("what it says") — and pin them to a board. To write the next sentence you glance at your board instead of re-reading the book. The board only ever grows; you never erase, only append.

The catch is the board has finite wall space. A long book fills the wall, and serving many readers at once means many boards competing for the same wall. The whole art of LLM serving is packing boards onto that wall efficiently — which is why the cache, not the math, is the constraint.

The one idea to hold: attention is \(O(n^2)\) in compute for a full sequence, but autoregressive decode with a KV cache is \(O(n)\) compute per token and \(O(n)\) memory. The bottleneck migrates from FLOPs to bytes — decode is bandwidth-bound, and the cache is what you are streaming.
Decode step: generating token t+1 KV cache (grows each step, append-only) K₁,V₁tok 1 K₂,V₂tok 2 K₃,V₃tok 3 · · · Kₜ,Vₜtok t Kₜ₊₁new query qₜ₊₁ softmax(qₜ₊₁·Kᵀ)·V → output; append Kₜ₊₁,Vₜ₊₁ to cache Only ONE new query computed per step — past K/V reused, never recomputed.
Architecture diagram: each decode step computes a single new query, attends over the cached K/V of all past tokens, then appends its own K/V. Compute per step is constant; the cache grows by one slot.

Round 2 — Internal Mechanics & Mathematical Model

The recomputation avoided

In causal self-attention, generating token \(t\) requires \(\operatorname{softmax}(q_t K_{1:t}^\top/\sqrt{d_k})V_{1:t}\). The keys \(K_{1:t}\) and values \(V_{1:t}\) depend only on tokens \(1..t\), which are already fixed. Recomputing them each step would cost \(\sum_{t=1}^{n} O(t\,d) = O(n^2 d)\) just to rebuild K/V, on top of attention. Caching makes each step reuse \(K_{1:t-1}, V_{1:t-1}\) and compute only \(k_t, v_t\):

\[ K_{1:t} = \big[\,K_{1:t-1}\;;\;x_t W_K\,\big],\qquad V_{1:t} = \big[\,V_{1:t-1}\;;\;x_t W_V\,\big] \]

Cache size — the key equation

The total KV cache in bytes for a full context is:

\[ M_{\text{KV}} = 2 \cdot L \cdot n \cdot n_{\text{kv}} \cdot d_h \cdot b \cdot B \]

where the leading 2 is K and V, \(L\) = layers, \(n\) = sequence length, \(n_{\text{kv}}\) = key/value heads, \(d_h\) = head dimension, \(b\) = bytes per element (2 for fp16), and \(B\) = batch size (concurrent sequences). Worked example: Llama-2-70B has \(L{=}80\), \(n_{\text{kv}}{=}8\) (grouped-query), \(d_h{=}128\). At \(n{=}4096\), fp16, one sequence: \(2\cdot80\cdot4096\cdot8\cdot128\cdot2 \approx 2.7\) GB. Ten concurrent users → 27 GB of cache alone, before weights.

Why decode is memory-bandwidth bound

Define arithmetic intensity = FLOPs ÷ bytes moved. Prefill processes \(n\) tokens against the weights in one matmul-heavy pass — high intensity, compute-bound. Decode processes one token but must stream the entire cache and all weights from HBM each step — intensity \(\ll 1\), so the GPU's compute units idle waiting on memory. The per-token decode latency is approximately:

\[ t_{\text{tok}} \approx \frac{M_{\text{weights}} + M_{\text{KV}}}{\text{BW}_{\text{HBM}}} \]

This is why decode throughput scales with memory bandwidth, not FLOPS, and why the roofline for LLM inference sits on the bandwidth ceiling.

Complexity analysis

Compute: \(O(n\,d)\) per decode step, \(O(n^2 d)\) total — same asymptotic as recompute, but the constant collapses because K/V projections are done once. Memory: \(O(L\,n\,n_{\text{kv}}\,d_h)\), linear in context and batch. This linear-in-\(n\) memory is the new wall — where dense attention had a quadratic compute wall, cached decode has a linear memory wall that fills HBM.

The tricks that shrink it

Multi-Query (MQA) and Grouped-Query Attention (GQA) cut \(n_{\text{kv}}\): instead of one K/V head per query head (\(n_{\text{kv}} = h\)), share K/V across groups (\(n_{\text{kv}} = h/g\)). GQA-8 on a 64-head model shrinks the cache 8×. KV quantization drops \(b\) from 2 to 1 or below. PagedAttention (vLLM, 2023) attacks fragmentation, not size — see Round 3.

Invariants & limiting cases

Invariant: the cache is append-only and causal — token \(t\)'s entries never change once written, so a prompt shared by many requests can share its prefill cache (prefix caching). Limiting cases: as \(n \to 1\) the cache vanishes and decode ≈ prefill; as \(n \to \infty\) memory dominates everything and context length is capped purely by HBM; as \(B \to \infty\) (huge batch) you regain arithmetic intensity but hit the cache-capacity ceiling first — the fundamental batch-vs-context tension.

Round 3 — Where It Breaks & Expert Debates

Memory fragmentation was the hidden killer. Naïve serving pre-allocates a contiguous cache buffer for the model's maximum context per request. Real requests are variable-length, so most of that buffer is reserved-but-unused — measured internal + external fragmentation wasted 60–80% of KV memory in pre-2023 systems. PagedAttention borrowed the OS virtual-memory idea: split the cache into fixed-size blocks, keep a block table per sequence, allocate on demand. It raised achievable batch size several-fold — the single biggest serving-throughput jump of the era, and the reason vLLM took over.

The batch-vs-latency tension is unresolvable, only tradeable. Bigger batches amortize weight reads across more tokens (higher throughput) but each user waits longer per token and the cache footprint balloons. Continuous/in-flight batching (Orca, vLLM) helps by admitting and retiring sequences mid-batch, but the Pareto frontier between throughput and per-token latency is fundamental — you pick a point, you don't escape the curve.

Long context breaks the linear assumption in practice. At 128k–1M tokens the cache dwarfs the weights, and offloading it to CPU/NVMe reintroduces a bandwidth cliff. Whether the answer is architectural (linear/state-space attention, Mamba), lossy (KV eviction / H2O / attention-sink retention), or systems (tiered cache) is actively contested — no consensus as of 2026.

KV quantization vs quality. Dropping keys/values to int8 or int4 saves memory linearly but keys are more sensitive than values (they set the softmax logits). How far you can push before measurable degradation is model- and task-dependent, and benchmarks disagree.

Failure mode to remember: serving systems can thrash — if admitted sequences' combined cache exceeds HBM mid-generation, the scheduler must preempt and either recompute or swap out a sequence's cache, causing latency spikes. Over-admitting on optimistic length estimates is the classic outage cause.

Round 4 — AI × Networks Connection

The KV cache is the exact reason "inference at the edge" is a memory problem, not a FLOPs problem. An O-RAN near-RT RIC node has bounded HBM and a hard ~10 ms–1 s loop; the admissible context length and batch of any LLM/Transformer xApp is set by the cache-size equation from Round 2, not by the model's parameter count. When you ask "can a config-generation or traffic-forecasting model run on this RIC," you are really asking "does its KV cache fit and can it be streamed within the loop budget."

The batch-vs-latency curve maps one-to-one onto the RIC latency gradient. Centralized non-RT RIC inference can batch aggressively for throughput (offline training-data generation, policy synthesis); near-RT inference must stay small-batch for latency, paying worse hardware utilization for responsiveness. This is the same tradeoff as placing an inference xApp on the latency gradient — the serving intuition transfers directly to placement decisions.

Cross-links

AI · Transformer attention internals → the O(n²) source fact; the KV cache is the optimization that reshapes its cost from compute to memory.

AI · LLM serving on K8s → the cache-size equation here is the input to the autoscaling / bin-packing problem there.

Networks · O-RAN architecture → the RIC's HBM and loop budget bound the admissible cache, hence the deployable model.

Pending intersection / AI nodes this unblocks: LLM serving on K8s, LLM for network config generation, inference at the edge (constraints, architecture).

KV-cache budget across the RIC latency gradient near-RT RIC (~10ms) non-RT RIC (>1s) Small batch B, short n M_KV small → fits HBM low per-token latency poor GPU utilization use: live inference xApp Large batch B, long n M_KV large → HBM-bound high throughput high per-token latency use: offline policy synthesis M_KV = 2·L·n·n_kv·d_h·b·B sets what is deployable where
Intersection diagram: the cache-size equation decides where an LLM xApp can live. Near-RT placement forces small batch / short context for latency; non-RT placement can batch for throughput — the same Pareto curve as decode serving.

Open questions this raises

  • For a config-generation LLM on the near-RT RIC, what context length is admissible once the RIC's HBM is shared with other xApps — and does GQA/quantization buy enough headroom, or must the model be distilled?
  • Does prefix caching pay off for RAN config templates (long shared system prompt, short variable tail), and how much cache can realistically be shared across xApp requests?
  • Is lossy KV eviction (H2O-style) safe for network-control outputs where a dropped early token could change a generated command, or does control-plane use demand exact caches?
  • At what context length does a linear-attention / state-space model beat cached softmax attention for RAN KPI sequences — i.e. where does the memory wall justify changing architecture?

← Back to AI · Home