Expand description
Online FP8 KV cache scale calibration.
Tracks running max |K| and max |V| during the first
--fp8-kv-calibration-tokens tokens of inference to compute per-tensor
scales: scale = amax * headroom / 448.0 (mapping the observed dynamic
range onto FP8 E4M3 [-448, 448]).
§The invariant
FP8 KV round-trips (write fp8 = bf16/scale, read bf16 = fp8*scale) only
if the SAME scale quantizes and dequantizes an entry. Paged / multi-query
attention reads a sequence’s whole history in one pass with ONE
k_scale/v_scale, so a scale that changes under live cache entries
dequantizes them through the wrong basis (~6x error → generation garbage:
loops, empty completions). That is why the 2026-07-25 hardening froze the
scale on the FIRST observe.
§Atlas #919
Freezing on the first observe made --fp8-kv-calibration-tokens 256 a lie.
Every H100 serve log showed
FP8 KV cache with online calibration (checkpoint ships no k/v scales): freezing per-tensor scales on the first observed tokens.
and the thing being observed was the readiness probe — so a 24k-context
serve ran on scales derived from ~13 tokens.
The window is now real: the amax accumulates ACROSS requests until
window_tokens have been observed, and the batch that reaches the window
freezes on the amax of everything seen, itself included. The invariant is
preserved by REWRITING the entries written inside the window — the window’s
BF16 K/V and slot mappings are staged aside and replayed through the
existing reshape_and_cache_fp8 kernel at the frozen scale (see the
private staging submodule of this module for the full tradeoff). A
readiness probe therefore counts toward the window and can never end it on
its own.
Thread safety: uses parking_lot::Mutex for interior mutability. The lock
is uncontended (single inference thread) so lock overhead is negligible.
Structs§
- Fp8Kv
Calibration - Online FP8 KV cache scale calibration tracker for one attention layer.
- Fp8Kv
Write Target - Everything the replay needs to re-run
reshape_and_cache_fp8for one layer.
Functions§
- dtype_
runs_ online_ fp8_ kv_ calibration - Whether this KV dtype’s write path calls
Fp8KvCalibration::observe. - graphs_
ready_ after_ fp8_ kv_ cal - Lift CUDA-graph suppression once every calibrating layer has frozen.