pub struct Fp8KvCalibration { /* private fields */ }Expand description
Online FP8 KV cache scale calibration tracker for one attention layer.
Wraps calibration state in a Mutex so it can live inside a Send + Sync
struct (required by TransformerLayer trait).
Implementations§
Source§impl Fp8KvCalibration
impl Fp8KvCalibration
Sourcepub fn new(
attn_layer_idx: usize,
window_tokens: usize,
headroom: f32,
gpu: &dyn GpuBackend,
) -> Result<Self>
pub fn new( attn_layer_idx: usize, window_tokens: usize, headroom: f32, gpu: &dyn GpuBackend, ) -> Result<Self>
Create a new calibration tracker.
attn_layer_idx: this layer’s index, for the freeze log line.
window_tokens: --fp8-kv-calibration-tokens. The amax accumulates
over this many observed tokens, across requests, before the scale
freezes (#919). Clamped to MAX_STAGED_TOKENS so the BF16 staging
cannot reserve gigabytes per layer; 1 reproduces the pre-#919
freeze-on-first-observe behaviour, and 0 never gets here (the
attention initializer does not build a calibrator at all).
headroom: multiplier on the accumulated amax when freezing
(--fp8-kv-headroom, CLI-validated ≥ 1.0; clamped here as defense in
depth because a sub-1.0 value guarantees clipping).
gpu: GPU backend for allocating the absmax reduction buffer.
Sourcepub fn is_calibrating(&self) -> bool
pub fn is_calibrating(&self) -> bool
Whether calibration is still in warmup phase (scales not yet frozen).
Sourcepub fn scales(&self) -> (f32, f32)
pub fn scales(&self) -> (f32, f32)
Get current scales. Returns (k_scale, v_scale).
Inside the window: PROVISIONAL_SCALE (private to this module), which
every read during the window also uses — consistent, just coarse. After
the freeze: the data-derived scale the whole window was requantized to
(constant thereafter).
Sourcepub fn observe(
&self,
gpu: &dyn GpuBackend,
k_data: DevicePtr,
v_data: DevicePtr,
num_tokens: u32,
num_kv_heads: u32,
head_dim: u32,
stream: u64,
target: &Fp8KvWriteTarget,
) -> Result<()>
pub fn observe( &self, gpu: &dyn GpuBackend, k_data: DevicePtr, v_data: DevicePtr, num_tokens: u32, num_kv_heads: u32, head_dim: u32, stream: u64, target: &Fp8KvWriteTarget, ) -> Result<()>
Observe K/V projection outputs and update the running max.
Launches absmax reductions on the K and V buffers, reads them back after
a sync, then either stages this batch (still inside the window) or
freezes and requantizes the staged window. Call this AFTER the K/V
projections and BEFORE writing to the KV cache — the caller’s write then
uses Self::scales, which is exactly what this call just decided.
k_data/v_data: device BF16 K/V projection outputs.
num_tokens/num_kv_heads/head_dim: this batch’s shape.
target: where the caller is about to write, so the freeze can replay
the window into the same pools.