Module fp8_calibration

Module fp8_calibration 

Source
Expand description

Online FP8 KV cache scale calibration.

Tracks running max |K| and max |V| values during the first N tokens of inference to compute per-tensor scales: scale = max / 448.0 (mapping the observed dynamic range to FP8 E4M3 [-448, 448]).

The scale is frozen on the FIRST observation (which runs immediately before the first KV write) and is then constant for the entire lifetime of every cache entry. This is required for correctness: FP8 KV round-trips (write fp8=bf16/scale, read bf16=fp8*scale) only if the SAME scale quantizes and dequantizes an entry. Freezing after N warmup tokens (the old design) wrote early entries at a placeholder scale, then froze to a different value, silently invalidating all already-written / cached KV — a paged multi-query read spanning the freeze boundary then read history through the wrong scale and generation degenerated. See the freeze block in observe.

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.

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.