Module fp8_calibration

Module fp8_calibration 

Source
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§

Fp8KvCalibration
Online FP8 KV cache scale calibration tracker for one attention layer.
Fp8KvWriteTarget
Everything the replay needs to re-run reshape_and_cache_fp8 for 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.