← Back to Table of Contents

Chapter 29 — Memory Hierarchy in Action: KV-Cache, Batching & Activations

Chapter 9 gave you the memory hierarchy — registers, SRAM, HBM, NVMe — as a stack of speed-versus-capacity tradeoffs. Chapter 11 gave you the roofline: the single number, arithmetic intensity, that tells you whether a kernel waits on the multipliers or on the wires. This chapter is where those two ideas stop being theory and start being the spreadsheet you keep open while you size a deployment. Every “will it fit?”, every “how many requests can I batch?”, every “why did my throughput collapse at 8k context?” is the memory hierarchy talking to you directly.

We will work through the three places the hierarchy bites hardest in practice: the KV-cache (the runtime memory that grows with context and batch, and routinely exceeds the weights), batching (the lever that drags memory-bound decode back toward compute-bound), and activation memory in training (the hidden tensor mountain that dwarfs your parameters at long sequence length). For each, we derive the byte math from first principles, plug in real model and hardware numbers, and write Python you can run to size your own case.

Throughout, our reference hardware is the NVIDIA H100 SXM: 80 GB HBM3 at ~3.35 TB/s, ~989 TFLOP/s dense BF16, ridge point ≈ 295 FLOP/byte (from Chapter 11). Substitute your own GPU’s two numbers — capacity (GB) and bandwidth (TB/s) — and every formula here still holds.

The one-sentence version: LLM serving is an exercise in spending a fixed HBM budget — weights, KV-cache, and activations all compete for the same gigabytes and the same terabytes-per-second, and almost every serving decision you make is really a decision about which of those three to shrink.

Recap: Why Decode Is Bandwidth-Bound (and the Token-Rate Ceiling)

Autoregressive decode generates one token per forward pass. To produce that one token, the GPU must read every weight of the model from HBM exactly once (at batch size 1). The arithmetic per token is tiny — a stack of matrix-vector products (GEMVs), arithmetic intensity ≈ 1 FLOP/byte — so the multipliers sit mostly idle and the HBM bus is the bottleneck. This is the memory wall in its purest form.

That gives a hard ceiling on token rate that depends only on two numbers:

The denominator is “bytes read per token,” which at batch=1 is essentially the model size in bytes (weights), plus the KV-cache bytes you must stream for attention. The weights dominate, so to first order the ceiling is HBM_BW / model_bytes.

Worked Example: A 70B Model at INT8

A 70-billion-parameter model quantized to INT8 (1 byte/weight) occupies ~70 GB of weights. Per generated token, all 70 GB must traverse the HBM bus:

That ~48 tok/s is a ceiling no amount of tensor-core horsepower can lift, because the tensor cores are not the constraint. The table below shows how the ceiling moves with precision and model size — read it as “the only two knobs that change decode speed at batch=1 are fewer bytes per weight and fewer weights.”

Note the immediate practical consequence: a 70B model in BF16 (140 GB) does not even fit in one 80 GB GPU, let alone run fast. Quantization here buys two things at once — it makes the model fit (capacity) and stream faster (bandwidth). This is why local/single-GPU LLM serving lives at INT4/INT8, a thread we pick up quantitatively in Chapter 25 and turn into a decision rule in Chapter 30.

The ceiling is optimistic — it ignores the KV-cache reads, attention, kernel-launch overhead, and the fact that you never hit 100% of peak bandwidth (70–85% is typical). Real throughput is below it. But it tells you the order of magnitude before you write a line of serving code.

The KV-Cache: The Memory That Grows With Your Conversation

What It Is and Why It Exists

In attention, each new token must attend to all previous tokens. Naively, generating token $t$ would require recomputing the keys ($K$) and values ($V$) for tokens $1 \ldots t-1$ on every step — an $O(t^2)$ blow-up that would make long generations quadratically slow. The KV-cache is the fix: after a token’s $K$ and $V$ vectors are computed once, they are stored in HBM and reused for every subsequent step. Decode then only computes $K, V$ for the single new token and reads the cached rest.

This trades compute for memory — the canonical hardware bargain. The cost is that the cache is runtime memory: it did not exist when you loaded the model, it grows with every token, and it grows with every concurrent request.

The KV-Cache Size Formula

For a single sequence, the cache stores a $K$ and a $V$ tensor at every layer, for every cached position. The size in bytes is:

where:

Note $n_{kv} \times d_{head}$ is the total KV width per token per layer. For a model with multi-head attention (MHA), $n_{kv} = n_{heads}$ and this equals $d_{model}$. For grouped-query attention (GQA) it is much smaller — which is the whole point (below).

Worked Example: Llama-3-style 8B at FP16

Take an 8B model with $L = 32$ layers, $d_{model} = 4096$, $n_{heads} = 32$, $d_{head} = 128$, and GQA with $n_{kv} = 8$ key/value heads. Per token per layer the cache stores $2 \times n_{kv} \times d_{head} = 2 \times 8 \times 128 = 2048$ elements. Across 32 layers that is $32 \times 2048 = 65{,}536$ elements/token, or at FP16 (2 bytes) 128 KB per token.

Now scale by context and batch:

The headline: at 128k context, the KV-cache for one sequence nearly equals the entire model’s weights. At batch 64 and 8k context, the KV-cache is four times the size of the weights. This is why long context and high concurrency are “memory-hungry” — it is not vague; it is this formula. The KV-cache, not the weights, is usually what determines how many requests you can serve at once.

The Same Model With Multi-Head Attention (No GQA)

Had this model used full MHA ($n_{kv} = 32$ instead of 8), every number above multiplies by 4×. The 8,192-token, batch-32 case would need ~137 GB of KV-cache — impossible on an 80 GB GPU. This single architectural choice (GQA) is the difference between “serves 64 concurrent users” and “OOMs at 16.”

Reducing the KV-Cache: GQA/MQA, Quantization, Paging, Prefix Reuse

Because the cache so often dominates, there is a whole toolbox aimed squarely at the four factors in the formula:

GQA/MQA ([grouped/multi-query attention]) directly attack $n_{kv}$. With $n_{heads}$ query heads but only $n_{kv}$ key/value heads, several query heads share one K/V head. MQA ($n_{kv}=1$) is the extreme; GQA ($n_{kv}=4$ or $8$) is the sweet spot that keeps quality. The KV-cache shrinks by exactly $n_{heads}/n_{kv}$.

KV-cache quantization attacks the bytes factor. Keys and values are activations, so they are harder to quantize than weights (outliers — see Chapter 25), but per-token or per-channel scaling makes INT8/FP8 KV practical, halving the cache. It is the right lever precisely when context is long (the cache, not the weights, is your capacity problem).

PagedAttention (the idea behind vLLM) attacks waste. A naive serving system reserves a contiguous buffer sized for each request’s maximum length, so a request that generates 100 tokens but was sized for 8,192 wastes 98% of its reservation — internal fragmentation. Paging the cache into small fixed blocks (e.g., 16 tokens) and mapping them through a block table — exactly the virtual-memory trick from Chapter 10 — lets memory be allocated on demand, raising the number of concurrent sequences 2–4×.

Prefix caching attacks redundancy: many requests share a long system prompt or few-shot preamble. Compute its KV once, reuse it for every request that shares the prefix. Saves both the prefill FLOPs and the duplicate KV bytes.

Batching: Trading Latency for Throughput by Raising Arithmetic Intensity

Why Batching Works

We saw in Chapter 11 that decode arithmetic intensity at batch $B$ is approximately:

because $B$ sequences all multiply against the same weight matrix that was read from HBM only once. You read the weights once and do $B$ times the arithmetic. Below the ridge point (~295 on H100), throughput scales nearly linearly with $B$; above it, you become compute-bound and further batching stops helping per-token speed.

The Throughput vs Latency Tradeoff

Batching is free throughput but not free latency. Two metrics pull in opposite directions:

A single user wants low TPOT and TTFT → small batch. A serving fleet wants maximum aggregate tok/s per GPU-dollar → large batch. The same hardware serves both, just at different points on the roofline. This is the central tension every inference team negotiates, and Chapter 30 turns it into explicit scenarios.

A useful way to hold it: batching converts idle HBM bandwidth into throughput. At batch=1 you pay for 3.35 TB/s and use it to stream weights you immediately discard; batching makes each of those expensive byte-reads do $B$ tokens’ worth of work.

Continuous (In-Flight) Batching

Classic “static” batching waits to assemble $B$ requests, runs them lockstep until all finish, then starts the next batch — so one long generation stalls a whole batch of short ones, wrecking utilization. Continuous batching (a.k.a. in-flight or iteration-level batching) instead schedules at the granularity of a single decode step: finished sequences leave the batch and new ones join every iteration. Combined with PagedAttention, this keeps the GPU near its target batch size at all times and is why modern servers (vLLM, TensorRT-LLM, TGI) get multiples of the throughput of naive batching.

The constraint on how large the running batch can grow is HBM capacity, and the thing that fills it is — once more — the KV-cache. Continuous batching and PagedAttention exist to pack as many concurrent KV-caches as possible into the gigabytes left over after the weights.

Activation Memory in Training: The Hidden Mountain

Inference is dominated by weights + KV-cache. Training adds a third, often largest, consumer: activations — the intermediate tensors saved during the forward pass so the backward pass can compute gradients. (We met the parameter/optimizer/gradient side of the training memory budget in Chapter 28; this is the fourth pillar.)

Why Activations Dominate at Long Sequence Length

For a transformer, the activations saved per layer scale with batch × sequence × hidden size, and the attention activations scale with batch × heads × sequence² before fusion. Crucially, activation memory grows with sequence length while parameter memory does not. So as context grows, activations come to dominate.

A rough per-layer activation estimate (forward, pre-checkpointing) is on the order of:

with $c \approx 10\text{–}30$ depending on how many intermediate tensors the implementation stores (QKV, attention output, MLP intermediate, norms, etc.).

Worked example. A model with $d_{model}=4096$, $L=32$, training in BF16 at $B=8$, $\text{seq}=4096$, with $c \approx 16$:

That 137 GB of activations dwarfs the ~16 GB of weights and would never fit. Double the sequence to 8,192 and it doubles to ~274 GB. This is why long-context training is an activation-memory problem first and a compute problem second.

Activation / Gradient Checkpointing: Recompute Instead of Store

The fix is the same bargain as the KV-cache, run in reverse: trade compute for memory by not storing most activations. With activation (gradient) checkpointing, you save only a few activations per layer (the “checkpoints”) and recompute the rest during the backward pass by re-running the forward of that segment.

The classic result: checkpointing every layer reduces activation memory from $O(L)$ to roughly $O(\sqrt{L})$ at the cost of one extra forward pass (~33% more compute). In PyTorch it is a one-line wrap with torch.utils.checkpoint.checkpoint. The recompute is comparatively cheap because — as Chapter 11 showed — much of the forward pass is memory-bound, so the “extra” FLOPs are partly hidden behind bandwidth you were paying for anyway.

Offloading: Spilling to CPU and NVMe

When even checkpointing is not enough, you can offload tensors down the hierarchy from Chapter 9: push optimizer states, gradients, or activations to CPU DRAM (~50–100 GB/s over PCIe) or even NVMe (~3–14 GB/s). This is what ZeRO-Offload / ZeRO-Infinity do (Chapter 28). It buys capacity at a steep bandwidth cost: moving a tensor to CPU and back over PCIe 5.0 (~64 GB/s) is ~50× slower than HBM, so offloading is a “fit it at all” tool, not a “make it fast” tool. The art is overlapping the transfers with compute so the PCIe latency hides behind the GPU’s work.

Prefill vs Decode: Two Workloads, Two Optimizations

A single LLM request has two phases with opposite roofline character, and conflating them is a common and expensive mistake.

Chunked prefill is the modern reconciliation: a long prompt’s prefill is broken into chunks and interleaved with ongoing decode steps of other requests in the same iteration. This stops a big prefill from blocking everyone’s decode (a TTFT-vs-TPOT fairness fix) and keeps both the tensor cores (prefill) and the bandwidth (decode) busy simultaneously — using the one request type to fill the roof the other leaves empty.

Python: Sizing the Memory Hierarchy

The three calculators below are the ones you actually reach for: KV-cache sizing across (model, seq, batch, dtype), the decode token-rate ceiling, and the batch-vs-arithmetic-intensity curve. They have no dependencies beyond the standard library.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from dataclasses import dataclass

# ---- Hardware (NVIDIA H100 SXM; swap for your GPU) ----
HBM_CAPACITY_GB = 80
HBM_BW_BYTES    = 3.35e12      # 3.35 TB/s
PEAK_FLOPS_BF16 = 989e12       # dense BF16 tensor-core FLOP/s
RIDGE = PEAK_FLOPS_BF16 / HBM_BW_BYTES   # ≈ 295 FLOP/byte

@dataclass
class Model:
    name: str
    n_params: float       # total parameters
    n_layers: int         # L
    d_model: int
    n_heads: int          # query heads
    n_kv_heads: int       # key/value heads (= n_heads for MHA, fewer for GQA/MQA)
    @property
    def d_head(self) -> int:
        return self.d_model // self.n_heads

# A few reference models
LLAMA3_8B  = Model("Llama-3-8B",  8.03e9, 32, 4096, 32, 8)    # GQA 8 kv-heads
LLAMA3_70B = Model("Llama-3-70B", 70.6e9, 80, 8192, 64, 8)    # GQA 8 kv-heads
MISTRAL_7B = Model("Mistral-7B",  7.24e9, 32, 4096, 32, 8)


def kv_cache_bytes(m: Model, seq: int, batch: int, bytes_per_elem: int = 2) -> int:
    """KV-cache size = 2 * L * n_kv * d_head * seq * batch * bytes."""
    return 2 * m.n_layers * m.n_kv_heads * m.d_head * seq * batch * bytes_per_elem


def weight_bytes(m: Model, bytes_per_weight: float = 2.0) -> float:
    """Model weight footprint in bytes (2=BF16, 1=INT8, 0.5=INT4)."""
    return m.n_params * bytes_per_weight


def decode_tokens_per_sec_ceiling(m: Model, bytes_per_weight: float = 2.0,
                                  seq: int = 0, batch: int = 1,
                                  kv_bytes_per_elem: int = 2) -> float:
    """
    Upper bound on decode tok/s at batch=1: HBM_BW / bytes_read_per_token.
    bytes_read = all weights + (optionally) the KV-cache streamed for attention.
    """
    wb = weight_bytes(m, bytes_per_weight)
    kv = kv_cache_bytes(m, seq, batch, kv_bytes_per_elem) if seq else 0
    bytes_per_token = wb + kv
    return HBM_BW_BYTES / bytes_per_token * batch  # ×batch: each step emits `batch` tokens


def decode_arith_intensity(batch: int) -> float:
    """AI of decode ≈ batch (weights read once, reused across the batch)."""
    return float(batch)


def fits(m: Model, seq: int, batch: int,
         bytes_per_weight: float = 2.0, kv_bytes_per_elem: int = 2,
         capacity_gb: float = HBM_CAPACITY_GB) -> bool:
    total = weight_bytes(m, bytes_per_weight) + kv_cache_bytes(m, seq, batch, kv_bytes_per_elem)
    return total <= capacity_gb * 1e9


# ---------- 1) KV-cache across context & batch ----------
print(f"{'='*64}\nKV-cache for {LLAMA3_8B.name} (FP16 KV), weights={weight_bytes(LLAMA3_8B)/1e9:.1f} GB")
print(f"{'context':>9} {'batch':>6} {'KV (GB)':>10} {'% of weights':>14}")
for seq in (2048, 8192, 32768, 128000):
    for batch in (1, 32, 64):
        kv = kv_cache_bytes(LLAMA3_8B, seq, batch) / 1e9
        pct = 100 * kv / (weight_bytes(LLAMA3_8B) / 1e9)
        print(f"{seq:>9} {batch:>6} {kv:>10.2f} {pct:>13.0f}%")

# ---------- 2) Decode token-rate ceiling ----------
print(f"\n{'='*64}\nDecode tok/s ceiling (weights only, batch=1)")
print(f"{'model':>14} {'BF16':>8} {'INT8':>8} {'INT4':>8}")
for m in (MISTRAL_7B, LLAMA3_8B, LLAMA3_70B):
    bf16 = decode_tokens_per_sec_ceiling(m, 2.0)
    int8 = decode_tokens_per_sec_ceiling(m, 1.0)
    int4 = decode_tokens_per_sec_ceiling(m, 0.5)
    print(f"{m.name:>14} {bf16:>8.0f} {int8:>8.0f} {int4:>8.0f}")

# ---------- 3) Batch vs arithmetic intensity ----------
print(f"\n{'='*64}\nBatch vs arithmetic intensity (ridge ≈ {RIDGE:.0f} FLOP/byte)")
print(f"{'batch':>6} {'AI':>8} {'regime':>14}")
for b in (1, 8, 32, 128, 256, 295, 512, 1024):
    ai = decode_arith_intensity(b)
    regime = "memory-bound" if ai < RIDGE * 0.9 else ("compute-bound" if ai > RIDGE * 1.1 else "balanced")
    print(f"{b:>6} {ai:>8.0f} {regime:>14}")

Sample output (abridged):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
================================================================
KV-cache for Llama-3-8B (FP16 KV), weights=16.1 GB
  context  batch    KV (GB)   % of weights
     2048      1       0.27             2%
     8192      1       1.07             7%
     8192     32      34.36           214%
     8192     64      68.72           428%
   128000      1      16.78           104%
================================================================
Decode tok/s ceiling (weights only, batch=1)
         model     BF16     INT8     INT4
    Mistral-7B      231      463      925
    Llama-3-8B      209      417      834
   Llama-3-70B       24       47       95
================================================================
Batch vs arithmetic intensity (ridge ≈ 295 FLOP/byte)
 batch       AI         regime
     1        1   memory-bound
    32       32   memory-bound
   256      256   memory-bound
   295      295       balanced
   512      512  compute-bound

Read the three tables together and the whole chapter is in front of you: KV-cache can exceed the weights (table 1), so it caps your batch; decode speed is set by weight bytes alone at batch=1 (table 2), so precision is your speed knob; and you must batch into the hundreds to escape the memory-bound regime (table 3), which is exactly what the KV-cache capacity in table 1 is fighting you on.

Putting It Together: The HBM Budget Equation

Everything in this chapter reduces to one accounting identity. The GPU’s HBM must simultaneously hold:

1
2
3
4
5
┌──────────────────────── 80 GB HBM (H100) ────────────────────────┐
│ Weights (fixed)      │ KV-cache (grows with seq × batch) │ work │
│ e.g. 70B INT8 = 70GB │  ◄── this is what caps your batch ──►│ space│
└───────────────────────────────────────────────────────────────────┘
        capacity problem            capacity + bandwidth problem

• Quantizing weights (INT8/INT4) frees capacity and raises the decode ceiling (bandwidth). Pick this for “won’t fit” and “too slow at batch=1.”
• Quantizing the KV-cache (INT8/FP8) and using GQA free the capacity that limits your batch. Pick these for “OOMs at high concurrency / long context.”
• PagedAttention + continuous batching reclaim wasted capacity, letting you run closer to the true budget. Pick these always (they are free).
• Checkpointing / offloading trade compute or PCIe bandwidth for capacity in training. Pick these for “won’t fit” during training.

Each lever buys back a specific scarce resource — capacity, bandwidth, or compute. Naming which one you are short of is the entire skill, and it is the subject of the next chapter.

Why This Matters for Model Optimization

The practical workflow this chapter hands you:

1. Before deploying, run the HBM budget equation. Sum weights + KV-cache (at your target seq and batch) + ~10% workspace. If it exceeds capacity, you have a capacity problem → quantize weights, quantize KV, GQA, or shard (Chapter 28).
2. Compute the decode ceiling (HBM_BW / model_bytes). If it is below your latency target, you have a bandwidth problem → lower precision or a smaller model. No amount of tuning beats this ceiling.
3. Find your batch target on the roofline. If throughput per GPU is too low, you are under-batching (memory-bound) → enable continuous batching + PagedAttention and raise the batch until either you hit the ridge or KV-cache fills HBM.
4. For training, account for activations separately and reach for checkpointing before offloading before sharding — cheapest resource first.

Key Takeaways

• Decode is memory-bandwidth-bound. At batch=1 the token-rate ceiling is HBM_BW / model_bytes: a 70B INT8 model on an H100 (70 GB ÷ 3.35 TB/s) tops out near ~48 tok/s, regardless of tensor-core peak.
• The KV-cache trades compute for memory and is sized by $2 \times L \times n_{kv} \times d_{head} \times \text{seq} \times \text{batch} \times \text{bytes}$. It grows with context and batch and routinely exceeds the weights (≈98% at 128k context; 4× the weights at batch 64 / 8k).
• GQA/MQA cut $n_{kv}$ (a 32→8 head reduction is a 4× KV cut), KV quantization cuts the byte width (2–4×), PagedAttention kills fragmentation, and prefix caching removes redundant prefill — all to fit more concurrent sequences.
• Batching raises arithmetic intensity (AI ≈ batch). It buys throughput by amortizing weight reads, at the cost of per-token latency — the throughput-vs-latency tradeoff. Continuous batching keeps the GPU full; KV-cache capacity caps how far you can go.
• Prefill is compute-bound (big GEMM, optimize FLOPs, governs TTFT); decode is memory-bound (GEMV, optimize bytes, governs TPOT). Chunked prefill interleaves them.
• In training, activation memory grows with sequence length and dwarfs weights at long context; gradient checkpointing trades ~33% extra compute for $O(\sqrt{L})$ memory, and offloading spills to CPU/NVMe to fit at all.
• Every serving decision is the HBM budget equation: weights + KV-cache + workspace ≤ capacity, streamed within the bandwidth ceiling. Name which of capacity / bandwidth / compute you are short of, and the right lever follows.

Next: Chapter 30 — How Datatypes Drive Your Daily Optimization Choices, where these memory and bandwidth budgets become a single decision framework: given your model, GPU, and regime, exactly which precision and technique to reach for.

← Back to Table of Contents