Model Optimization: Compression & KV Cache
Learn quantization, pruning, distillation, and the size and structure of KV Caches.
KV Parameters:
▶📚 Prerequisites(1)
🎓 Learning objectives
- •Compare PTQ and QAT quantization techniques
- •Explain KV Caching and calculate KV cache size: $O(layers imes heads imes seq imes dim)$
- •Describe Multi-Query Attention (MQA) and Grouped-Query Attention (GQA)
What is it?
LLM optimization encompasses the techniques for reducing inference cost and latency while preserving model quality: quantization (reducing weight precision), KV cache optimization, efficient attention mechanisms, batching strategies, and model distillation. These techniques make LLM deployment economically viable at production scale — without them, serving billions of requests is prohibitively expensive.
Why it exists
A naive GPT-4 deployment serving 1M requests/day at $0.03/request costs $30K/day. Optimization techniques reduce this by 5–20×. As LLMs move from experimental to mission-critical infrastructure, the difference between an optimized and unoptimized serving stack can be $5M+ per year for a mid-sized company.
Problem it solves
Two core problems: (1) memory bottleneck — LLM inference is bound by GPU memory bandwidth, not compute. A 70B model requires 140GB of memory just for weights, severely limiting batch sizes and throughput. (2) KV cache growth — each generated token requires reading all previous K/V matrices, making long-context generation O(n) in tokens and O(n²) in total memory across the sequence.
Intuition
Think of LLM inference as a pipeline factory. The bottleneck is not the assembly line speed (GPU FLOPS) — it is the warehouse capacity (GPU memory) and how fast workers can fetch parts from storage (memory bandwidth). Quantization reduces part size (4-bit vs 16-bit weights). KV caching means workers do not re-fetch the same parts every time. Continuous batching keeps the assembly line running even when individual orders arrive at different times.
If you come from Java/Spring Boot: LLM optimization is analogous to database optimization. Quantization is compression (smaller storage format for the same data). KV cache is query result caching (do not recompute what we already computed). Continuous batching is connection pooling (share resources across requests). Flash Attention is an optimized query execution plan.
If you come from React/Frontend: quantization is like compressing images (WebP vs PNG — 4× smaller, visually identical). KV cache is like browser caching — do not re-fetch resources already loaded. Continuous batching is like React's batched state updates — process multiple updates together rather than one by one.
Analogy
LLM optimization is like running an efficient restaurant kitchen. Quantization is using standard ingredient portions instead of custom measurements (simpler, faster, slightly less precise). KV cache is mise en place — prep all ingredients once, use them throughout service. Flash Attention is a more efficient cooking technique that produces the same dish faster. Continuous batching is coordinating multiple orders so the oven is never idle.
Technical explanation
Key optimization techniques:
-
QUANTIZATION FP32 → FP16: 2× memory reduction, ~0% quality loss on most tasks FP16 → INT8: 2× further reduction (4× vs FP32), <1% quality loss on most benchmarks INT8 → INT4: 2× further (8× vs FP32), 2–5% quality loss. Acceptable for many tasks. NF4 (QLoRA): 4-bit Normal Float — better quantization for LLM weights than INT4, ~1% quality loss Memory savings for Llama 3 70B: FP16=140GB → INT8=70GB → INT4=35GB
-
KV CACHE During autoregressive generation, K and V matrices for all previous tokens are cached. Without cache: token N requires recomputing all N-1 previous K/V pairs → O(n²) total FLOPs With cache: token N only computes its own K/V → O(n) total FLOPs, ~60–80% speedup vLLM PagedAttention: manages KV cache in pages (like OS virtual memory) → enables 24× higher throughput than naive implementations by eliminating fragmentation
-
FLASH ATTENTION Standard attention reads/writes O(n²) data to HBM (high-bandwidth memory). FA rewrites attention as a memory-efficient tiled computation — O(n) HBM I/O. 2–4× faster on long sequences, critical for >4K token contexts.
-
CONTINUOUS BATCHING (iteration-level scheduling) Naive: wait until batch is complete before accepting new requests. Long requests block short ones. Continuous batching: insert new requests mid-batch at iteration boundaries. GPU utilization goes from 30–50% to 80–90%.
Architecture
LLM serving stack (optimized, e.g. vLLM):
[Request Queue] ↓ [Scheduler] — continuous batching, priority queuing ↓ [PagedAttention KV Cache Manager] KV Cache: [Page 1][Page 2][Page 3]...[Page N] (non-contiguous physical memory, like OS pages) ↓ [Model Executor] Weights: INT4/INT8 quantized, loaded in GPU VRAM Flash Attention: tiled SRAM computation Tensor parallelism (multi-GPU): weights sharded across GPUs ↓ [Tokenizer + Detokenizer] ↓ [Response]
Multi-GPU deployment (tensor parallelism): 70B model weights sharded across 4× A100 40GB Each GPU holds 17.5GB of weights All-reduce synchronization at each attention layer 4× throughput vs single GPU, ~3× after comms overhead
Workflow
- Benchmark baseline: measure TTFT (time to first token), TPS (tokens per second), and GPU memory usage with unoptimized serving
- Apply quantization: INT8 for <1% quality loss, INT4 if memory-constrained and task tolerates it
- Enable KV caching: verify cache hit rates (should be >90% for repeated prefixes)
- Switch to Flash Attention: enabled by default in modern frameworks (transformers v4.34+, vLLM)
- Implement continuous batching: use vLLM or TGI (HuggingFace Text Generation Inference) instead of naive serving
- Profile bottlenecks: is inference memory-bound (not enough VRAM for batch) or compute-bound (GPU near 100% utilization)?
- Scale horizontally: multiple replicas behind a load balancer once single-GPU is optimized
Example
Deploying Llama 3 8B with vLLM (production-grade serving)
from vllm import LLM, SamplingParams
INT4 quantization via AWQ — 4× less VRAM than FP16
llm = LLM( model="meta-llama/Meta-Llama-3-8B-Instruct", quantization="awq", # AWQ INT4: 4GB vs 16GB FP16 max_model_len=8192, # max context length gpu_memory_utilization=0.90, # use 90% of GPU VRAM tensor_parallel_size=1 # increase for multi-GPU )
params = SamplingParams(temperature=0.7, max_tokens=512)
Continuous batching: submit multiple requests simultaneously
prompts = [ "Explain transformer attention in 2 sentences", "What is RAG?", "Write a Python hello world" ] outputs = llm.generate(prompts, params) # batched efficiently for output in outputs: print(output.outputs[0].text)
Real-world usage
Groq: custom hardware (LPU — Language Processing Unit) achieves 700+ tokens/second for Llama 3 70B — 10× faster than GPU-based serving. Focuses entirely on memory bandwidth optimization, which is the actual bottleneck.
vLLM (UC Berkeley): open-source LLM serving framework with PagedAttention. Used by companies including LinkedIn, IBM, and LG to serve open-source models. Reports 24× throughput improvement over naive serving, 3.5× over HuggingFace TGI.
TogetherAI: serves 50+ open-source models at commodity prices ($0.20/1M tokens for Llama 3 8B vs $1.00+ from closed API providers) through aggressive quantization and hardware optimization.
Trade-offs
Quantization level vs quality: INT8 is almost always safe (benchmark first). INT4 reduces quality by 2–5% on reasoning tasks — acceptable for classification/extraction, risky for complex code generation or math. GPTQ and AWQ are better INT4 methods than naive round-to-nearest quantization.
Batch size vs latency: larger batches improve GPU utilization and reduce cost-per-token but increase latency (requests wait longer for the batch to fill). Use smaller batches for interactive use cases (<500ms TTFT requirement), larger batches for batch processing jobs.
Model size vs capability: a 70B model at INT4 (35GB) costs 5× more to serve than a 7B model at FP16 (14GB) but is significantly better at complex reasoning. For simple tasks (classification, extraction, summarization), the 7B model is often sufficient at 5× lower cost.
Visual explanation
Dense Attention vs. GQA vs. MQA: DENSE ATTENTION GROUPED-QUERY (GQA) MULTI-QUERY (MQA) Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q Q │ │ │ │ │ │ │ │ / / / / \ / / / / / K K K K K K K K K K K K K V V V V V V V V V V V V V (1 KV Head per Q Head) (1 KV Head per Group) (1 KV Head for All)
Advantages
- —
Speeds up token generation by 10x-100x
- —
Shrinks model sizes by 4x using INT4 compression
Disadvantages
- —
Low-bit quantization can degrade logical reasoning capabilities in small models
Common mistakes
- —
Applying INT4 quantization without benchmarking quality first. Some tasks are highly sensitive to quantization (code generation, multi-step math) while others are not (summarization, classification). Always run your specific task benchmarks before deploying INT4 in production.
- —
Not monitoring KV cache eviction rates in production. When the KV cache fills up, vLLM evicts older sequences — this causes those requests to recompute from scratch (latency spike). Monitor cache hit rate and cache size. If eviction rate >5%, increase allocated cache memory or reduce max concurrent sequences.
- —
Using batch inference API for interactive applications. Batch APIs optimize for throughput — responses may take 5–60 minutes. Deploying this for a chatbot creates terrible UX. Use streaming real-time inference for interactive use cases, batch for background processing jobs.
- —
Not enabling Flash Attention when available. Flash Attention is a drop-in replacement for standard attention with no quality change, only speed improvement. It is disabled by default in some frameworks. Check that flash_attn is installed and enabled: model.config.use_flash_attention_2 = True.
- —
Sharding a model across too many GPUs. Tensor parallelism introduces all-reduce communication overhead at every layer. 2-GPU sharding for a 7B model achieves ~1.8× (not 2×) throughput. 4 GPUs achieve ~3× (not 4×). Beyond 4 GPUs for 70B models, communication overhead dominates — use pipeline parallelism instead.
🎤 Interview questions
Explain the difference between Post-Training Quantization (PTQ) and Quantization-Aware Training (QAT). When is QAT necessary?
Why is autoregressive LLM decoding memory-bandwidth bound rather than compute bound?
📂 Subtopics
Quantization: INT8, INT4, GPTQ and AWQ
Quantization shrinks a model by storing its weights in fewer bits — trading a little quality for large memory and speed gains. GPTQ and AWQ are the two dominant 4-bit methods.
~14 min
KV Cache Optimization and PagedAttention
The KV cache stores attention keys and values so the model doesn't recompute them every token — but it eats memory fast. vLLM's PagedAttention manages that memory like an operating system manages RAM.
~14 min
Batching Strategies: Static, Dynamic and Continuous
Batching runs multiple requests through the GPU together to keep it busy. Continuous batching (used by vLLM) is the key advance: it swaps finished requests out and new ones in at every step instead of waiting for the whole batch to finish.
~13 min
Model Pruning and Distillation
Pruning removes weights the model barely uses; distillation trains a small 'student' model to imitate a large 'teacher.' Both shrink the model, but distillation changes the architecture while pruning thins the existing one.
~13 min