LLM Inference Optimization — Quantization, Speculative Decoding, and KV Cache

Introduction

LLM inference bottlenecks are memory bandwidth, not compute. Every token generated requires reading gigabytes of model weights from GPU memory, making it fundamentally different from training. Modern optimization techniques — quantization, speculative decoding, paged attention, and flash attention — stack multiplicatively to deliver 5-10x speedups with minimal quality loss.

Why LLM Inference Is Slow

The root cause is the memory wall. A 70B parameter model at FP16 requires 140 GB of GPU memory just to load. Each forward pass reads all those weights once per token. At 2 TB/s of GPU memory bandwidth, that limits you to roughly 14 tokens per second at batch size 1 — no matter how many FLOPS your GPU has. Optimization strategies attack this constraint from different angles.

Key insight: inference is memory-bandwidth-bound, not compute-bound. A GPU with 4x more FLOPS but the same memory bandwidth produces the same token throughput for large models.

Quantization: Shrink the Model, Speed Up the Read

Quantization reduces the bits used to represent each weight, shrinking memory footprint and speeding up reads:

# Quantization comparison for a 7B parameter model
quantization_tiers = {
    "FP32": {"vram_gb": 28, "speedup": 1.0, "quality_loss_pct": 0.0},
    "FP16": {"vram_gb": 14, "speedup": 1.3, "quality_loss_pct": 0.1},
    "BF16": {"vram_gb": 14, "speedup": 1.2, "quality_loss_pct": 0.2},
    "INT8": {"vram_gb":  7, "speedup": 2.0, "quality_loss_pct": 0.5},
    "INT4 (GPTQ)": {"vram_gb": 3.5, "speedup": 3.5, "quality_loss_pct": 1.0},
    "INT4 (NF4)":  {"vram_gb": 3.5, "speedup": 3.2, "quality_loss_pct": 0.8},
}
 
for name, stats in quantization_tiers.items():
    ratio = stats["speedup"] / max(stats["quality_loss_pct"], 0.1)
    print(
        f"{name:15} | {stats['vram_gb']:5.1f} GB | "
        f"{stats['speedup']:.1f}x faster | "
        f"{stats['quality_loss_pct']:.1f}% quality loss | "
        f"efficiency ratio: {ratio:.1f}"
    )

Decision guide:

FP16 is the production default — 1.3x faster, negligible quality loss.
INT8 is the sweet spot for memory-constrained deployments — 2x faster, under 1% quality loss.
INT4 GPTQ for extreme compression (3.5x faster) — acceptable for tasks like summarization where perfect fidelity is not critical.

Implementing Quantization With vLLM

from vllm import LLM, SamplingParams
 
# INT4 GPTQ — load a pre-quantized model
llm_int4 = LLM(
    model="TheBloke/Llama-2-70B-GPTQ",
    quantization="gptq",
)
 
# AWQ — an alternative 4-bit format with better quality retention
llm_awq = LLM(
    model="TheBloke/Llama-2-70B-AWQ",
    quantization="awq",
)
 
sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
 
prompts = ["Explain transformer attention mechanisms"] * 10
outputs = llm_int4.generate(prompts, sampling_params)
 
for output in outputs:
    print(output.outputs[0].text[:200])

Pre-quantized GPTQ models are faster than dynamic INT8 quantization because the weight conversion is done once at export time rather than on every forward pass.

Speculative Decoding: Parallelize Token Generation

Standard generation processes one token at a time sequentially. Speculative decoding uses a small draft model to predict multiple tokens ahead, then verifies them in a single main model forward pass:

from vllm import LLM, SamplingParams
 
# vLLM has native speculative decoding support
llm_speculative = LLM(
    model="meta-llama/Llama-2-70b-hf",         # Main model
    speculative_model="meta-llama/Llama-2-7b-hf", # Draft model
    num_speculative_tokens=5,                    # Tokens to draft per step
    tensor_parallel_size=4,
)
 
# The speedup depends on draft acceptance rate
# For typical conversational text: 2-3x speedup
# For code generation: 1.5-2x speedup (less predictable patterns)
 
result = llm_speculative.generate(
    ["Write a Python function to merge two sorted lists"],
    SamplingParams(temperature=0.0, max_tokens=256),
)
print(result[0].outputs[0].text)

Speculative decoding works best when the output is predictable (e.g., the draft model gets 60-80% of tokens right). It shines for long outputs above 256 tokens and loses its advantage for very short responses.

KV Cache Management With Paged Attention

Attention mechanisms must cache key-value tensors for all previous tokens. For a 70B model with a 4K context, this KV cache alone requires roughly 128 GB — exceeding what fits in GPU memory for large batches. vLLM's PagedAttention solves this:

from vllm import LLM, SamplingParams
 
# vLLM manages KV cache using OS-style virtual memory paging
llm = LLM(
    model="meta-llama/Llama-2-13b-hf",
    max_num_batched_tokens=32768,   # Total tokens across all sequences in a batch
    max_model_len=4096,             # Max sequence length
    gpu_memory_utilization=0.90,    # Use 90% of GPU memory for KV cache
)
 
# Without PagedAttention: batch size 8, latency 1000ms
# With PagedAttention:     batch size 256, latency 50ms (20x throughput)
 
sampling_params = SamplingParams(temperature=0.7, max_tokens=128)
outputs = llm.generate(["Tell me about quantum computing"] * 64, sampling_params)
print(f"Processed {len(outputs)} requests")

PagedAttention allocates KV cache in fixed-size pages and reuses them across requests. This eliminates GPU memory fragmentation and enables batch sizes an order of magnitude larger than naive implementations.

Prefix Caching for Shared System Prompts

When many requests share the same system prompt, recomputing its KV cache for every request is wasteful. Prefix caching stores these shared prefixes once:

from vllm import LLM, SamplingParams
 
llm = LLM(
    model="meta-llama/Llama-2-13b-hf",
    enable_prefix_caching=True,
)
 
system_prompt = """You are an expert technical assistant.
Answer with precise, production-ready code examples.
Always cite relevant documentation."""
 
user_queries = [
    "How do I handle database connection pooling in Python?",
    "What is the best way to implement retry logic?",
    "Explain asyncio event loops.",
] * 100
 
prompts = [f"{system_prompt}\n\nUser: {q}" for q in user_queries]
 
# First request: computes and caches the 80-token system prompt
# Requests 2-300: reuse cached KV tensors — 15-20% faster
outputs = llm.generate(prompts, SamplingParams(temperature=0.0, max_tokens=256))

Prefix caching is particularly valuable for multi-turn chat (accumulated conversation history), few-shot classification (shared examples in the prompt), and RAG pipelines (shared retrieval context).

Flash Attention: Faster and Memory-Efficient Attention

Standard attention computes the full N x N attention matrix and stores it in GPU memory, which is O(n^2) in sequence length. Flash Attention rewrites this computation to use GPU SRAM tiling, reducing memory usage to O(n) and improving throughput by 2-4x:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
 
# Enable Flash Attention 2 explicitly
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-13b-hf",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",  # Drop-in replacement
    device_map="auto",
)
 
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-13b-hf")
 
inputs = tokenizer("Explain LLM inference optimization", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=200)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 
# Flash Attention 2 benchmarks (A100 GPU):
# Standard attention:    50 tokens/sec at 4K context
# Flash Attention 2:    150 tokens/sec at 4K context  (3x speedup)
# Flash Attention 2:    200 tokens/sec at 32K context (even larger gains at longer context)

Flash Attention is now enabled by default in vLLM and Hugging Face Transformers. If you are not using it, you are leaving free speedup on the table.

Continuous Batching vs Static Batching

Static batching holds a batch until the longest request finishes, wasting GPU cycles on requests that completed early. Continuous batching releases completed sequences immediately and fills the freed slot with the next queued request:

# Continuous batching configuration in vLLM
llm = LLM(
    model="meta-llama/Llama-2-13b-hf",
    max_num_seqs=256,              # Maximum concurrent sequences
    max_num_batched_tokens=16384,  # Maximum total tokens per iteration
)
 
# Static batching (Hugging Face default):
#   Batch of 8 requests, max length 512 tokens
#   Short requests wait for the longest one to finish
#   GPU utilization: ~40%
 
# Continuous batching (vLLM default):
#   Batch dynamically filled as requests complete
#   GPU utilization: ~80-90%
#   2-4x higher throughput at identical hardware cost

Latency vs Throughput Configuration

Choose your optimization target based on the use case:

# Chat interface — optimize for low latency
chat_config = {
    "max_num_seqs": 32,           # Few concurrent sequences
    "max_num_batched_tokens": 4096,
    "gpu_memory_utilization": 0.7,
    # Expected: p50 latency < 100ms, p99 < 500ms
}
 
# Batch processing — optimize for throughput
batch_config = {
    "max_num_seqs": 256,
    "max_num_batched_tokens": 32768,
    "gpu_memory_utilization": 0.95,
    # Expected: 500+ requests/minute, latency 2-10s acceptable
}
 
# Embeddings — maximize throughput, latency irrelevant
embedding_config = {
    "max_num_seqs": 512,
    "max_num_batched_tokens": 65536,
    "gpu_memory_utilization": 0.95,
    # Expected: 5000+ embeddings/minute
}

Key Takeaways

LLM inference is memory-bandwidth-bound, not compute-bound — optimizations that reduce memory reads deliver direct throughput gains.
FP16 quantization is a free 1.3x speedup with negligible quality loss; INT8 doubles throughput with under 1% quality degradation on most benchmarks.
Flash Attention 2 reduces attention memory from O(n^2) to O(n) and delivers 2-4x speedup — it should always be enabled.
vLLM's PagedAttention enables batch sizes 10-20x larger than naive implementations by eliminating KV cache memory fragmentation.
Speculative decoding delivers 2-3x speedup on long outputs by parallelizing token verification; it is most effective above 256 tokens.
Prefix caching eliminates redundant computation for shared system prompts and saves 15-20% on requests with common prefixes.
Continuous batching achieves 80-90% GPU utilization versus 40% for static batching — never deploy vLLM without it.
Benchmark quality (MMLU, HellaSwag) after every quantization change — speed without quality verification is not a production optimization.

Introduction#

Why LLM Inference Is Slow#

Quantization: Shrink the Model, Speed Up the Read#

Implementing Quantization With vLLM#

Speculative Decoding: Parallelize Token Generation#

KV Cache Management With Paged Attention#

Prefix Caching for Shared System Prompts#

Flash Attention: Faster and Memory-Efficient Attention#

Continuous Batching vs Static Batching#

Latency vs Throughput Configuration#

Key Takeaways#

Related reading