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(
_warmup_tokens: usize,
headroom: f32,
gpu: &dyn GpuBackend,
) -> Result<Self>
pub fn new( _warmup_tokens: usize, headroom: f32, gpu: &dyn GpuBackend, ) -> Result<Self>
Create a new calibration tracker.
_warmup_tokens: retained for API/CLI compat but no longer gates the
freeze — the scale is now frozen on the FIRST observe (before any KV is
persisted), so a warmup window would only reintroduce the write/read
scale mismatch. Any value > 0 simply enables online calibration.
headroom: multiplier on the first-observe absmax 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).
Before the first observe: the conservative construction default (2.0, covering ±896) — never used for a persisted write, since observe() runs before every write and freezes on its first call. After the first observe: the frozen, data-derived scale (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,
) -> 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, ) -> Result<()>
Observe K/V projection outputs and update running max.
Launches absmax reduction kernels on the K and V buffers, then reads the results back to CPU after a sync. Call this AFTER K/V projections and BEFORE writing to the KV cache.
k_data: device BF16 buffer of K projection output
v_data: device BF16 buffer of V projection output
num_tokens: number of tokens in the batch
num_kv_heads: number of KV heads
head_dim: dimension per head