KV Cache Optimization and Memory Management: An Accounting Method for Long-Context Inference

Note: This is a review article (no first-hand experimental data), based on public architecture documentation and technical analysis.

Experiment Objective

Long-context inference (128K, 1M tokens) moves the bottleneck from "compute" to "KV Cache in memory." A 70B model needs about 10-20GB of KV Cache at 8K context; at 128K it easily exceeds 200GB — far more than the weights themselves. The goal: build an accounting framework for KV Cache (how to compute its size, where it grows, how to shrink it), and map the four mainstream optimizations — PagedAttention, continuous batching, quantization/pruning, and speculative decoding — with their benefits and boundaries, so you can answer quantitatively: "how much memory does KV take on this inference path, and how much can it still be compressed?"

Code Walkthrough

Part 1: KV Cache size accounting

def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len, batch, dtype_bytes=2):
    """KV Cache footprint for a single generation (current seq_len only)."""
    per_token_per_layer = 2 * n_kv_heads * head_dim * dtype_bytes  # one K + one V
    return n_layers * per_token_per_layer * seq_len * batch

# Example: 70B-class config (estimated)
print(kv_cache_bytes(80, 8, 128, 8192, 1) / 2**30, "GB @ 8K, bs=1")
print(kv_cache_bytes(80, 8, 128, 131072, 1) / 2**30, "GB @ 128K, bs=1")
# Outputs ≈ 12.6 GB / 8K; 201.3 GB / 128K

Key points:

1. KV Cache grows linearly with sequence length and multiplies with batch: seq_len × batch is a product — the higher the concurrency and the longer the context, the faster the cache balloons. This is why "long context + high concurrency" is the most memory-hungry inference combination.
2. 2 * n_kv_heads * head_dim: one copy each for K and V. Note it is KV heads, not Q heads — GQA (Grouped Query Attention) is the key to cutting cache here: 70B-class models commonly use 8 KV heads, versus 64 Q heads in MHA, so KV is 1/8 of Q.
3. dtype_bytes=2: 2 bytes per element in FP16/BF16. Quantizing to 8-bit halves it; 4-bit halves it again — KV quantization is a "multiplicative" memory win.

Part 2: Memory watermark budgeting

def memory_budget(model_gb, kv_gb_per_req, max_concurrent, total_gb, safety=0.05):
    usable = total_gb * (1 - safety)
    reserved = model_gb + 2.0            # weights + CUDA context/temp buffers
    budget_for_kv = usable - reserved
    max_requests = int(budget_for_kv // kv_gb_per_req)
    return max_requests, budget_for_kv

# Example: H100 80G single card, 70B model → 4-bit weights ~70GB
max_req, budget = memory_budget(70, 12.6, 0, 80)
print(f"4-bit 70B weights + 8K KV: {max_req} concurrent 8K requests per card")

Key understanding:

• Weight quantization is a prerequisite: 140GB of BF16 weights for a 70B model cannot fit in 80GB; quantize to 4-bit first (~70GB), then KV has budget. The inference memory ledger has four entries — weights + KV + temp buffers + CUDA overhead.
• safety=0.05 is the OOM floor: fragmentation, dynamic shapes, and CUDA-graph extra usage all need 5% headroom. A service packed to the brim will OOM on long-sequence requests.

Principles

Why KV Cache must exist

Transformer attention must "see" the K/V of all previous tokens. Recomputing history K/V (no cache) makes complexity grow quadratically with sequence length — 128K contexts would never run. So inference frameworks store the K/V of every generated token (the KV Cache), and each step only computes and appends the new token's K/V. Trading memory for time — that is the essence of KV Cache.

The four mainstream optimizations and their boundaries

Optimization	Principle	Typical gain	Boundary/cost
PagedAttention (vLLM)	allocate KV in pages, eliminate fragmentation, grow on demand	2-4× concurrency	needs framework support; page scheduling has CPU overhead
Continuous batching	request-level scheduling, no waiting for the whole batch	5-10× throughput	scheduler complexity; long-tail requests hold memory
KV quantization / GQA / pruning	store KV in 8/4-bit, reduce KV heads, sparse pruning	cache size 1/2-1/4	accuracy loss needs evaluation; pruning needs retraining
Speculative decoding	small model drafts, big model verifies in parallel	fewer generation steps → fewer KV writes	acceptance rate decides the gain (see the speculative-decoding article)

Design notes:

• PagedAttention "reduces waste"; quantization "reduces the total": the former fits more requests into existing memory, the latter makes each request cheaper — they are orthogonal and usually combined in production.
• Continuous batching solves the throughput ceiling: in static batching the slowest request drags the whole batch; continuous batching lets each request enter/exit independently, so the GPU always works on requests with work — throughput and memory utilization rise together.
• Long context is where quantization pays most: at 8K, KV is only 12GB and halving it saves 6GB — marginal; at 128K it saves 100GB, and quantization goes from "nice to have" to "necessary."

The memory ledger (accounting basis)

Inference of a 70B model on an 80GB card (estimated):

Item	Footprint	Note
Weights (4-bit)	~70GB	quantization required to fit
CUDA context/temp	~2GB	framework fixed cost
KV budget	~4GB	little left → single-card concurrency limited
Conclusion	—	70B long-context inference basically needs multi-card TP or bigger cards

Results

(Estimated, not measured)

Scenario	No optimization	+PagedAttention	+KV 8-bit	+Continuous batching
Concurrent 8K requests (70B 4-bit)	0-1	3-4	6-8	stacks with the above
128K requests feasible?	no (OOM)	barely	yes	yes

Conclusion: the long-context combination is weight quantization (make the model fit) + KV quantization (halve the cache) + PagedAttention (zero fragmentation) + continuous batching (max throughput) — miss any one and 128K scenarios hit the memory wall.

Engineering Takeaways

1. Account before optimizing: use the formula above to compute the three entries — weights + KV + fixed overhead — to find the true bottleneck. Many "slow inference" cases are actually KV filling memory, causing swapping/queuing.
2. Quantize KV first, not only weights: weight quantization usually gets priority, but in long-context scenarios KV quantization can pay more (halving cache directly doubles concurrency).
3. Add a KV dimension to monitoring: DCGM shows total memory, not "how much is KV." vLLM metrics (vllm:num_preemptions, gpu_cache_usage_perc) belong on the dashboard; KV hitting the reserved watermark is the scaling signal.
4. Include acceptance rate in KV decisions: speculative decoding (covered in chapter 12) reduces KV writes indirectly through fewer generation steps; combined with quantization, the gains multiply — evaluate them together.

Next article preview: this one completed the "inference memory ledger"; next we return to the application side of the AI-compute × supercomputing stack — "Surrogate Models for Scientific Computing: Order Reduction from High-Fidelity Simulation to AI Surrogates."