KV cache quantization is one of the highest-leverage optimizations for LLM inference serving. The key-value cache grows linearly with sequence length and batch size, often consuming more VRAM than model weights at production context lengths. Quantizing the cache to 8-bit or 4-bit integers can reclaim 50-75% of that memory with minimal quality degradation. This guide walks through the quantization approaches, implementation details, and verification steps you need to deploy it reliably.
Step 1: Understand the memory profile
Before quantizing, measure your baseline. The KV cache stores two tensors per layer per token: keys and values, each shaped [batch, num_heads, seq_len, head_dim]. For a 70B model with 80 layers, 64 heads, and 128 head dimension, each token consumes roughly 1.6 MB per layer in FP16 — about 128 MB per token across the full model. At 4K context with batch 8, that’s 4 GB just for the cache.
Run this profiling snippet on your target hardware:
import torch
from transformers import AutoModelForCausalLM, AutoConfig
model_id = "meta-llama/Llama-2-70b-hf"
config = AutoConfig.from_pretrained(model_id)
# Calculate per-token KV cache size in bytes
head_dim = config.hidden_size // config.num_attention_heads
kv_per_token_per_layer = 2 * config.num_attention_heads * head_dim * 2 # 2 for K+V, 2 for FP16 bytes
total_kv_per_token = kv_per_token_per_layer * config.num_hidden_layers
print(f"Per-token KV cache: {total_kv_per_token / 1e6:.1f} MB")
print(f"4K context, batch 8: {total_kv_per_token * 4096 * 8 / 1e9:.2f} GB")Verify the numbers match torch.cuda.memory_allocated() during a real forward pass with use_cache=True.
Step 2: Choose your quantization scheme
Three schemes dominate production deployments:
Per-tensor symmetric INT8 — Single scale per tensor. Simplest to implement, ~2x compression. Works well for keys; values are more sensitive.
Per-channel asymmetric INT8 — Separate scale and zero-point per output channel (head dimension). ~2x compression with better accuracy retention, especially for values.
Per-group INT4 — Group size 32-128 along the head dimension. ~4x compression. Requires careful calibration; quality drops sharply below group size 32.
Start with per-channel INT8 for values and per-tensor INT8 for keys. This hybrid approach captures most memory savings while preserving quality. Reserve INT4 for memory-constrained edge deployments where you can tolerate additional calibration effort.
Step 3: Implement quantization hooks
Most inference engines (vLLM, TGI, TensorRT-LLM) expose KV cache quantization via configuration flags. If you’re building a custom engine or need fine-grained control, insert quantization at the attention output projection. Here’s a minimal PyTorch reference:
import torch
import torch.nn.functional as F
class QuantizedKVCache:
def __init__(self, num_layers, num_heads, head_dim, dtype=torch.int8, group_size=128):
self.num_layers = num_layers
self.num_heads = num_heads
self.head_dim = head_dim
self.dtype = dtype
self.group_size = group_size
self.k_cache = []
self.v_cache = []
self.k_scales = []
self.v_scales = []
self.v_zero_points = []
def quantize(self, x, scales=None, zero_points=None):
"""Quantize FP16 tensor to INT8/INT4 with per-channel or per-group scaling."""
if self.dtype == torch.int8:
if scales is None: # per-tensor
scale = x.abs().max() / 127
x_q = (x / scale).round().clamp(-128, 127).to(torch.int8)
return x_q, scale
else: # per-channel
x_q = (x / scales.unsqueeze(-1)).round().clamp(-128, 127).to(torch.int8)
return x_q, scales
else: # int4 per-group
orig_shape = x.shape
x = x.view(-1, self.group_size)
max_val = x.abs().max(dim=-1, keepdim=True).values
scale = max_val / 7 # 4-bit signed: -8..7
x_q = (x / scale).round().clamp(-8, 7).to(torch.int8)
return x_q.view(orig_shape), scale.view(-1)
def dequantize(self, x_q, scales, zero_points=None):
if self.dtype == torch.int8:
if zero_points is not None:
return (x_q.float() - zero_points.unsqueeze(-1)) * scales.unsqueeze(-1)
return x_q.float() * scales.unsqueeze(-1)
else:
return x_q.float() * scales.view(-1, 1)
def append(self, layer_idx, k, v):
"""k, v: [batch, num_heads, seq_len, head_dim] in FP16"""
if layer_idx >= len(self.k_cache):
self.k_cache.append(None)
self.v_cache.append(None)
self.k_scales.append(None)
self.v_scales.append(None)
self.v_zero_points.append(None)
# Keys: per-tensor symmetric
k_q, k_scale = self.quantize(k)
self.k_scales[layer_idx] = k_scale if self.k_scales[layer_idx] is None else \
torch.max(self.k_scales[layer_idx], k_scale)
self.k_cache[layer_idx] = k_q if self.k_cache[layer_idx] is None else \
torch.cat([self.k_cache[layer_idx], k_q], dim=2)
# Values: per-channel asymmetric
v_reshaped = v.transpose(1, 2).reshape(-1, self.num_heads * self.head_dim)
v_max = v_reshaped.max(dim=0).values
v_min = v_reshaped.min(dim=0).values
v_scale = (v_max - v_min) / 255
v_zp = (-v_min / v_scale).round().clamp(0, 255).to(torch.uint8)
v_q = ((v_reshaped - v_min) / v_scale).round().clamp(0, 255).to(torch.uint8)
v_q = v_q.view(v.shape[0], v.shape[2], self.num_heads, self.head_dim).transpose(1, 2)
self.v_scales[layer_idx] = v_scale if self.v_scales[layer_idx] is None else \
torch.max(self.v_scales[layer_idx], v_scale)
self.v_zero_points[layer_idx] = v_zp if self.v_zero_points[layer_idx] is None else \
torch.max(self.v_zero_points[layer_idx], v_zp)
self.v_cache[layer_idx] = v_q if self.v_cache[layer_idx] is None else \
torch.cat([self.v_cache[layer_idx], v_q], dim=2)
def get(self, layer_idx):
k = self.dequantize(self.k_cache[layer_idx], self.k_scales[layer_idx])
v = self.dequantize(self.v_cache[layer_idx], self.v_scales[layer_idx], self.v_zero_points[layer_idx])
return k, vThis implementation prioritizes clarity over speed. Production kernels fuse quantization with the attention output projection and use custom CUDA kernels for the dequantize+matmul path.
Step 4: Calibrate with representative data
Static quantization requires calibration data that matches your production distribution. Run 100-500 sequences through the model at your target context lengths and collect per-channel min/max for values, per-tensor max for keys. Save these statistics alongside your model artifacts.
def calibrate_kv_cache(model, dataloader, num_layers, num_heads, head_dim, max_samples=200):
k_max = torch.zeros(num_layers)
v_min = torch.zeros(num_layers, num_heads, head_dim)
v_max = torch.zeros(num_layers, num_heads, head_dim)
model.eval()
with torch.no_grad():
for i, batch in enumerate(dataloader):
if i >= max_samples:
break
outputs = model(**batch, use_cache=True, output_attentions=False)
past_kv = outputs.past_key_values
for layer_idx, (k, v) in enumerate(past_kv):
k_max[layer_idx] = max(k_max[layer_idx], k.abs().max().item())
v_min[layer_idx] = torch.minimum(v_min[layer_idx], v.min(dim=2).values.min(dim=0).values)
v_max[layer_idx] = torch.maximum(v_max[layer_idx], v.max(dim=2).values.max(dim=0).values)
return {"k_max": k_max, "v_min": v_min, "v_max": v_max}Use your actual request payloads for calibration — synthetic data produces misleading ranges. If your traffic mixes short and long contexts, calibrate on the longest 10% of requests; shorter sequences naturally fall within the same ranges.
Step 5: Integrate with your inference engine
VLLM
Set kv_cache_dtype="fp8" or "int8" in EngineArgs. vLLM 0.4+ handles calibration automatically for FP8; for INT8 you provide a calibration file via kv_cache_dtype="int8:path/to/calib.json".
vllm serve meta-llama/Llama-2-70b-hf \
--kv-cache-dtype int8:calib.json \
--gpu-memory-utilization 0.85 \
--max-model-len 8192TensorRT-LLM
Enable kv_cache_quant_mode=INT8 in the build config. The trtllm-build step runs calibration if you pass --calibrate_kv_cache.
# build_config.py
from tensorrt_llm.builder import BuilderConfig
config = BuilderConfig(
max_batch_size=8,
max_input_len=4096,
max_output_len=4096,
kv_cache_quant_mode="INT8",
strongly_typed=True,
)Custom engine
Replace the standard past_key_values tuple with your QuantizedKVCache instance. Modify the attention forward pass to call cache.append(layer_idx, k, v) after the output projection and cache.get(layer_idx) before the attention scores. Ensure the dequantize path fuses with the torch.bmm or flash attention kernel — separate dequantize+matmul kills the memory bandwidth savings.
Step 6: Verify correctness and quality
Run three validation passes before deploying:
1. Numerical parity — Compare logits between quantized and FP16 cache on a fixed seed. Maximum absolute difference should stay below 1e-3 for INT8, 5e-3 for INT4.
def verify_numerical_parity(model, quantized_cache, test_inputs, tol=1e-3):
model.eval()
with torch.no_grad():
# FP16 baseline
out_fp16 = model(**test_inputs, use_cache=True)
logits_fp16 = out_fp16.logits[:, -1, :]
# Quantized
# (replace model's cache with quantized_cache instance)
out_q = model(**test_inputs, use_cache=True)
logits_q = out_q.logits[:, -1, :]
max_diff = (logits_fp16 - logits_q).abs().max().item()
print(f"Max logit diff: {max_diff:.6f}")
assert max_diff < tol, f"Parity failed: {max_diff} > {tol}"2. Perplexity evaluation — Run your standard eval harness (e.g., WikiText-2, C4, or your domain-specific eval set). INT8 should stay within 0.1-0.3 perplexity points of FP16. INT4 typically adds 0.5-1.5 points.
# Example with lm-evaluation-harness
lm_eval --model hf \
--model_args pretrained=meta-llama/Llama-2-70b-hf,kv_cache_dtype=int8 \
--tasks wikitext \
--device cuda:0 \
--batch_size 43. Latency and memory profiling — Measure peak VRAM and decode-step latency at your target batch and context. KV cache quantization should reduce peak memory by 40-60% for INT8, 70-80% for INT4. Decode latency may improve slightly (less memory bandwidth) or stay flat (dequantize overhead).
def profile_memory_and_latency(model, batch_size=8, seq_len=4096, num_steps=50):
import time
torch.cuda.reset_peak_memory_stats()
input_ids = torch.randint(0, 32000, (batch_size, seq_len), device="cuda")
# Prefill
with torch.no_grad():
_ = model(input_ids, use_cache=True)
prefill_mem = torch.cuda.max_memory_allocated() / 1e9
torch.cuda.reset_peak_memory_stats()
# Decode steps
next_token = torch.randint(0, 32000, (batch_size, 1), device="cuda")
latencies = []
for _ in range(num_steps):
torch.cuda.synchronize()
t0 = time.perf_counter()
with torch.no_grad():
_ = model(next_token, use_cache=True)
torch.cuda.synchronize()
latencies.append(time.perf_counter() - t0)
decode_mem = torch.cuda.max_memory_allocated() / 1e9
avg_latency = sum(latencies) / len(latencies) * 1000
print(f"Prefill peak VRAM: {prefill_mem:.2f} GB")
print(f"Decode peak VRAM: {decode_mem:.2f} GB")
print(f"Avg decode latency: {avg_latency:.2f} ms")
return prefill_mem, decode_mem, avg_latencyStep 7: Handle edge cases in production
Dynamic sequence lengths — The cache grows incrementally. Your quantization scales must be monotonic (only increase) or recomputed periodically. Monotonic scales avoid re-quantizing existing tokens but may waste dynamic range. A practical compromise: recompute scales every 1024 new tokens using the current cache window.
Mixed-precision attention — If you use flash attention with FP8 or BF16 accumulation, ensure the dequantized KV tensors match the kernel’s expected dtype. Mismatched dtypes trigger silent upcasts that negate memory savings.
Prefix caching — When sharing KV cache across requests (common in multi-turn chat), quantize the shared prefix once at ingestion time. Store the quantized tensors and scales in your cache store; avoid re-quantizing on every request.
Provider fallback — If you route requests across multiple model providers (e.g., via an inference gateway like n4n.ai), ensure quantization parameters travel with the model artifact or are recomputed per-provider. Different GPU architectures (H100 vs A100 vs L4) have different tensor core behaviors that can shift the quality-latency tradeoff.
Step 8: Monitor regression in production
Add these metrics to your inference observability stack:
kv_cache_memory_bytes— Track per-request and aggregate. Alert if it exceeds 90% of the expected quantized budget.quantization_scale_drift— Log the max scale per layer per 1000 requests. Sudden increases indicate distribution shift; trigger recalibration.decode_latency_p99— Quantization should not increase tail latency. If it does, the dequantize kernel is likely not fused properly.perplexity_sampled— Periodically sample 0.1% of requests for offline perplexity evaluation against a held-out set.
A simple Prometheus rule for scale drift:
- alert: KVCacheScaleDrift
expr: |
max_over_time(kv_cache_scale_max[1h])
/ max_over_time(kv_cache_scale_max[24h:1h]) > 1.5
for: 10m
labels:
severity: warning
annotations:
summary: "KV cache quantization scale increased >50% in 1h"Verification checklist
Before marking the rollout complete, confirm:
- Peak VRAM at target batch/context matches or beats your capacity plan
- Logit parity test passes at your chosen tolerance
- Perplexity delta is within your SLA (typically <0.3 for INT8)
- Decode latency p99 is at or below FP16 baseline
- Scale drift alerts are configured and tested
- Recalibration pipeline runs automatically on drift detection
- Prefix caching works correctly with quantized shared prefixes
- Fallback providers serve quantized cache without manual intervention
KV cache quantization is a rare optimization that delivers outsized returns for modest engineering investment. The memory headroom it unlocks lets you serve larger batches, longer contexts, or bigger models on the same hardware — often the difference between a viable deployment and one that OOMs under load.