Module fp8_calibration

Module fp8_calibration 

Source
Expand description

2026-09-25: Online per-tensor FP8 KV scale calibration for one attention layer.

The running max of |K| and |V| accumulates across requests over the first window_tokens observed tokens (--fp8-kv-calibration-tokens, clamped to MAX_STAGED_TOKENS). The observe that reaches the window freezes scale = amax * headroom / 448, over every observed token including its own batch (state.rs).

Attention dequantizes a sequence’s whole history with one k_scale/v_scale, so an entry is read correctly only at the scale that wrote it. Entries written inside the window use PROVISIONAL_SCALE; their BF16 K/V and slot mappings are staged and, at the freeze, rewritten through ops::reshape_and_cache_fp8 at the frozen scale (staging.rs).

Owner: model-layers (FP8 KV cache). Invariants:

  • The scales change once, at the freeze, and afterwards only on the opt-in METRALE_FP8_KV_EMA_RECAL path.
  • The scale set at the freeze is at least amax / 448, where amax covers every token observed up to the freeze (headroom is clamped to >= 1.0 in Fp8KvCalibration::new).

Structs§

Fp8KvCalibration
2026-09-25: The calibrator of one attention layer. observe takes &self; the state sits behind a Mutex.
Fp8KvWriteTarget
2026-09-25: Where the caller’s live FP8 write goes, so the replay can repeat it (qwen3_attention/decode/write_kv_cache_fp8.rs builds it).

Functions§

dtype_runs_online_fp8_kv_calibration
2026-09-25: Whether this KV dtype’s write path calls Fp8KvCalibration::observe.
graphs_ready_after_fp8_kv_cal
2026-09-25: True when every calibrating layer has frozen, which is when the model lifts its CUDA-graph suppression (trait_impl/decode_a.rs).