How do you calibrate FP8 KV-cache quantization scales using representative inference data?

FP8 KV-cache calibration estimates the numerical scales used to store Transformer attention key/value (KV) tensors in 8-bit floating-point format. Rather than assuming every tensor fits a default scale of 1.0, calibration runs representative prompts through the original model, observes KV activation ranges, and chooses scales that reduce clipping and quantization error. This is numerical quantization calibration, not the probability calibration of LLM answer confidence covered in Question 203.

1. Why calibration is needed

• KV states can have very different magnitudes across layers, attention heads, and keys versus values.
• FP8 stores fewer bits than BF16; a poorly chosen scale can cause clipping of large values or loss of precision for smaller values.
• vLLM supports FP8 KV cache with default scales of 1.0, or dataset-derived scales using llm-compressor. Per-tensor and, on a supported FlashAttention backend, per-attention-head calibration are available.
• Weights and KV cache are distinct targets: you can retain original BF16 model weights and change only the KV cache, subject to backend support. Some FP8 attention paths also quantize Q, so identify the complete experiment.

2. How scaling works

A simplified convention is to encode a floating-point value x using q = FP8(x / s) and reconstruct approximately x_hat = s × float(q), where s is the calibrated scale. The exact storage convention and kernel implementation vary by backend.

Beginner-friendly numerical example:

• Suppose a sampled key tensor contains values with maximum absolute magnitude 896.
• Assume an FP8 E4M3 representation whose largest finite magnitude is 448 (for the relevant encoding).
• A simple maximum-based illustrative scale is s = 896 / 448 = 2.
• A key activation of 600 is first divided by 2 → 300, then rounded to a representable FP8 number, and multiplied by 2 during reconstruction.
• Using s = 1 would put 600 above the illustrative FP8 maximum of 448, so a saturating quantizer would clip it.

Important: The formula is a teaching example of abs-max scaling, not a guarantee that production calibration simply uses the maximum; outlier-aware or percentile-based policies can behave differently. FP8 levels are not uniformly spaced, and scale estimates are subject to backend requirements.

3. Choose representative calibration inputs

1. Gather a small, authorized, privacy-reviewed sample of real workload patterns or realistic synthetic data.
2. Cover short and long prompts, chat templates, system instructions, document retrieval, code/math, multilingual inputs and uncommon formatting relevant to your service.
3. Match the token-length distribution and challenging context windows. Calibrating only on short, generic text can fail at 32k/128k contexts.
4. Keep a separate holdout evaluation dataset. Do not select scales by repeatedly tuning to the final test set.
5. Track tokenizer/model revision, attention backend, dtype, calibration seed and sample counts for reproducibility.

4. Practical calibration workflow using vLLM tools

• Start from an unquantized, compatible checkpoint and its tokenizer.
• Prepare the calibration prompts using the same chat template and tokenization as production; select samples and maximum sequence lengths deliberately.
• Use llm-compressor and its FP8 attention/KV-cache calibration recipe, selecting per-tensor scales initially and per-head scales where supported.
• Run calibration to observe attention activations and write the computed Q/K/V scales into the saved model's quantization metadata as supported.
• Load that calibrated artifact in vLLM with kv_cache_dtype="fp8"; verify that the calibrated scales, not default 1.0 scales, were loaded and that the selected attention backend supports the chosen granularity.
• The official commands and supported parameters evolve: follow the pinned vLLM/llm-compressor version's executable example rather than copying an old recipe blindly.

5. Per-tensor versus per-head example

Imagine two KV heads in the same layer:

• Head A observed maximum absolute value = 448 → illustrative scale 1.0.
• Head B observed maximum absolute value = 896 → illustrative scale 2.0.
• One shared tensor scale of 2.0 can accommodate both heads but reduces the representational resolution available for smaller Head A values.
• Per-head scales of 1.0 and 2.0 can preserve more useful detail for each head, at the cost of added scale metadata and backend constraints.

Separate K and V scales matter; some runtimes also calibrate Q scales when using quantized attention. Do not assume every serving kernel supports per-head scales.

6. Validate before deployment

Compare BF16 baseline, FP8 default scale, FP8 calibrated per-tensor, and FP8 calibrated per-head (if supported), keeping the model and workload fixed:

• Accuracy: long-context retrieval, needle-in-a-haystack tests, exact-match, reasoning/code correctness and task-specific failure rates.
• Numerics: saturation/clipping frequency, quantization error by layer/head and rare-activation outliers.
• Performance: KV bytes, p50/p95/p99 TTFT, inter-token latency, cache preemption rate, concurrency and actual wire-transfer bytes in distributed serving.
• Safety and release: predeclare acceptable accuracy degradation, use shadow/canary traffic, and revert to BF16 or selected higher-precision layers for sensitive regressions.

Calibration should be rerun or checked after changes to model weights, attention backend, chat-template/data distribution, supported context lengths, or quantization scheme.

7. Typical troubleshooting

• Long-context quality worsens: extend calibration samples to realistic long contexts and test per-head or mixed-precision options.
• FP8 metrics unexpectedly differ: ensure scales were actually loaded; check default-scale fallbacks and the selected kernel's Q/K/V quantization.
• KV memory is smaller but latency is not: inspect quantize/dequantize overhead, occupancy, attention kernels and network-transfer format.
• Rare prompt types fail: improve domain coverage; inspect head-level outliers and protected layers rather than simply lowering a global quality threshold.

Interview-ready answer: “I'd collect representative prompt traces, run the BF16 checkpoint through an FP8 KV calibration recipe, estimate compatible per-tensor or per-head Q/K/V scales, save them with the model and verify vLLM loads them. I'd compare BF16, default-scale FP8 and calibrated FP8 on long-context quality and p95 latency, with rollback if accuracy or backend compatibility fails.”

References: vLLM — Quantized KV Cache; LLM Compressor — KV Cache Quantization Example.