spark_runtime/

kv_cache.rs

// SPDX-License-Identifier: AGPL-3.0-only

//! Paged KV cache block allocator.
//!
//! Manages a pool of fixed-size blocks for attention KV storage.
//! Each block holds `block_size` token positions for all KV heads.

use crate::gpu::DevicePtr;
use anyhow::{Result, bail};

pub(crate) const NVFP4_GROUP_SIZE: usize = 16;

/// KV cache quantization dtype.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KvCacheDtype {
    /// 2 bytes per element.
    Bf16,
    /// 1 byte per element (FP8 E4M3 with per-tensor scale).
    Fp8,
    /// 0.5 bytes data + per-group FP8 scale (E2M1 packed nibbles).
    Nvfp4,
    /// 4-bit WHT + Lloyd-Max quantization (TurboQuant). Same byte layout as NVFP4
    /// but with Walsh-Hadamard rotation and optimal Gaussian codebook for ~2x
    /// lower MSE at the same bit rate.
    Turbo4,
    /// 3-bit WHT + Lloyd-Max (8 levels). 22% smaller than turbo4.
    Turbo3,
    /// 2-bit WHT + Lloyd-Max (4 levels). 6.4x compression vs bf16 (3 bits/elem
    /// total: 2 b data + 0.5 b scale + 0.5 b layout overhead). Full write +
    /// paged-decode + chunked-prefill kernel coverage. 2-bit keys cannot
    /// sustain tool-grammar constrained decoding with the standard boundary
    /// policy; requires the higher auto high-precision-layer default (see
    /// `auto_high_precision_layers`) validated on the GB10 flagship.
    Turbo2,
    /// WHT + FP8 E4M3. Same memory as FP8 but with outlier suppression.
    /// Enables FP8-level memory for models with large RMS norm weights.
    Turbo8,
    /// TurboQuant+ asymmetric: K stored at turbo4 (4-bit), V at turbo3 (3-bit).
    /// K dominates attention score precision; V tolerates lower precision per
    /// turboquant_plus/docs/papers/asymmetric-kv-compression.md. Saves ~14%
    /// bandwidth at decode (4.5 b/elem K + 3.375 b/elem V vs 4.5 + 4.5
    /// symmetric turbo4). Decode kernel dispatch needs a new
    /// `paged_decode_attn_turbo4k_turbo3v` variant; write kernel forks
    /// `reshape_and_cache_flash_turbo4` for K and `..._turbo3` for V on the
    /// same launch.
    Turbo4KTurbo3V,
    /// K=turbo4, V=turbo8. K=4-bit codebook; V=FP8. ~11% bandwidth saving vs
    /// pure turbo8 symmetric. Same dispatch-table follow-up applies.
    Turbo4KTurbo8V,
    /// K=turbo3, V=turbo8. Smallest K (3-bit) with V=FP8 retention.
    Turbo3KTurbo8V,
    /// TurboQuant+ safer-asym: K stored at BF16 baseline (full precision), V
    /// compressed to turbo4 4-bit codebook. Preserves K's attention-score
    /// fidelity completely while compressing V which dominates KV bandwidth
    /// at long context.
    Bf16KTurbo4V,
    /// K=bf16, V=turbo3 (3-bit). Aggressive V compression with full-precision K.
    Bf16KTurbo3V,
    /// K=fp8 (1 byte/elem with per-tensor scale), V=turbo4. K kept at the
    /// usual fp8 quality; V at 4-bit codebook. Middle ground between bf16/turbo
    /// and pure turbo8.
    Fp8KTurbo4V,
    /// K=fp8, V=turbo3. Smallest combo retaining fp8 K precision.
    Fp8KTurbo3V,
    /// K=bf16 baseline, V=turbo2 (2-bit). Most aggressive V compression with
    /// full-precision K. Per asymmetric-kv-compression.md: symmetric turbo2/
    /// turbo2 collapses quality (+58.5% PPL); this asym preserves K and only
    /// pays the +9.5% V-side cost — 6× better quality at the same V compression.
    Bf16KTurbo2V,
    /// K=fp8, V=turbo2. The canonical "asymmetric rescue" config (analog of
    /// llama-cpp-turboquant's `q8_0/turbo2`). Best compression-to-quality
    /// ratio for turbo2 V on tested models.
    Fp8KTurbo2V,
}

impl std::fmt::Display for KvCacheDtype {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Delegates to the catalogue so the canonical spelling exists in ONE
        // match — see `catalog::name` for why that match must stay exhaustive.
        f.write_str(self.name())
    }
}

impl KvCacheDtype {
    /// Returns the (K_dtype, V_dtype) pair. For symmetric variants both
    /// elements are identical. For asymmetric variants the pair differs.
    pub fn kv_pair(self) -> (KvCacheDtype, KvCacheDtype) {
        match self {
            KvCacheDtype::Turbo4KTurbo3V => (KvCacheDtype::Turbo4, KvCacheDtype::Turbo3),
            KvCacheDtype::Turbo4KTurbo8V => (KvCacheDtype::Turbo4, KvCacheDtype::Turbo8),
            KvCacheDtype::Turbo3KTurbo8V => (KvCacheDtype::Turbo3, KvCacheDtype::Turbo8),
            KvCacheDtype::Bf16KTurbo4V => (KvCacheDtype::Bf16, KvCacheDtype::Turbo4),
            KvCacheDtype::Bf16KTurbo3V => (KvCacheDtype::Bf16, KvCacheDtype::Turbo3),
            KvCacheDtype::Fp8KTurbo4V => (KvCacheDtype::Fp8, KvCacheDtype::Turbo4),
            KvCacheDtype::Fp8KTurbo3V => (KvCacheDtype::Fp8, KvCacheDtype::Turbo3),
            KvCacheDtype::Bf16KTurbo2V => (KvCacheDtype::Bf16, KvCacheDtype::Turbo2),
            KvCacheDtype::Fp8KTurbo2V => (KvCacheDtype::Fp8, KvCacheDtype::Turbo2),
            other => (other, other),
        }
    }

    /// True for the symmetric turbo dtypes whose cache contents are stored
    /// in the WHT-rotated basis (the write path applies `wht_bf16_inplace`
    /// before quantizing). Gates the WHT(Q) / iWHT(out) attention bookends —
    /// call on the K or V side of `kv_pair()`, not on the combined variant.
    /// Turbo2 is rotated by the write path like the rest; omitting it here
    /// is what desynced the decode bookends from the write path.
    pub fn is_wht_rotated(self) -> bool {
        matches!(
            self,
            KvCacheDtype::Turbo2
                | KvCacheDtype::Turbo3
                | KvCacheDtype::Turbo4
                | KvCacheDtype::Turbo8
        )
    }

    /// True if K and V use different storage layouts.
    pub fn is_asymmetric(self) -> bool {
        matches!(
            self,
            KvCacheDtype::Turbo4KTurbo3V
                | KvCacheDtype::Turbo4KTurbo8V
                | KvCacheDtype::Turbo3KTurbo8V
                | KvCacheDtype::Bf16KTurbo4V
                | KvCacheDtype::Bf16KTurbo3V
                | KvCacheDtype::Fp8KTurbo4V
                | KvCacheDtype::Fp8KTurbo3V
                | KvCacheDtype::Bf16KTurbo2V
                | KvCacheDtype::Fp8KTurbo2V
        )
    }
}

impl std::str::FromStr for KvCacheDtype {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self> {
        match s {
            "bf16" => Ok(KvCacheDtype::Bf16),
            "fp8" => Ok(KvCacheDtype::Fp8),
            "nvfp4" => Ok(KvCacheDtype::Nvfp4),
            "turbo4" => Ok(KvCacheDtype::Turbo4),
            "turbo3" => Ok(KvCacheDtype::Turbo3),
            "turbo2" => Ok(KvCacheDtype::Turbo2),
            "turbo8" => Ok(KvCacheDtype::Turbo8),
            "turbo4k_turbo3v" | "turbo4k3v" => Ok(KvCacheDtype::Turbo4KTurbo3V),
            "turbo4k_turbo8v" | "turbo4k8v" => Ok(KvCacheDtype::Turbo4KTurbo8V),
            "turbo3k_turbo8v" | "turbo3k8v" => Ok(KvCacheDtype::Turbo3KTurbo8V),
            "bf16k_turbo4v" | "bf16k4v" => Ok(KvCacheDtype::Bf16KTurbo4V),
            "bf16k_turbo3v" | "bf16k3v" => Ok(KvCacheDtype::Bf16KTurbo3V),
            "fp8k_turbo4v" | "fp8k4v" => Ok(KvCacheDtype::Fp8KTurbo4V),
            "fp8k_turbo3v" | "fp8k3v" => Ok(KvCacheDtype::Fp8KTurbo3V),
            "bf16k_turbo2v" | "bf16k2v" => Ok(KvCacheDtype::Bf16KTurbo2V),
            "fp8k_turbo2v" | "fp8k2v" => Ok(KvCacheDtype::Fp8KTurbo2V),
            other => bail!(
                "Unsupported --kv-cache-dtype '{other}'. Symmetric: 'bf16', 'fp8', 'nvfp4', 'turbo4', 'turbo3', 'turbo8'. \
                Asymmetric (TQ+): turbo*_turbo*v, bf16k_turbo[34]v (safer asym: K baseline, V compressed), fp8k_turbo[34]v."
            ),
        }
    }
}

/// Configuration for the paged KV cache.
pub struct KvCacheConfig {
    /// Tokens per block.
    pub block_size: usize,
    /// Number of KV heads.
    pub num_kv_heads: usize,
    /// Dimension per head.
    pub head_dim: usize,
    /// Number of attention layers (only full_attention layers have KV cache).
    pub num_layers: usize,
    /// Quantization dtype for cache storage (uniform fallback).
    pub dtype: KvCacheDtype,
    /// Per-layer KV cache dtype override. When non-empty, `layer_dtypes[i]`
    /// specifies the dtype for attention layer `i`. When empty, all layers
    /// use the uniform `dtype` field (backward compatible).
    pub layer_dtypes: Vec<KvCacheDtype>,
    /// Per-layer (num_kv_heads, head_dim) overrides for heterogeneous
    /// attention models (e.g. Gemma-4 with sliding 16×256 vs full 4×512).
    /// When non-empty, `layer_dims[i]` specifies the (nkv, hd) for layer
    /// `i`; allocation and kernel stride computations use these per-layer
    /// values so writes/reads land at the correct offsets. When empty,
    /// all layers use the uniform `num_kv_heads`/`head_dim` (backward
    /// compatible — homogeneous models need no change).
    pub layer_dims: Vec<(usize, usize)>,
    /// `--high-speed-swap` HBM-shrink knob (Phase 6.1). When `Some(N)`,
    /// each sequence is capped at `N` HBM-resident blocks; older blocks
    /// are evicted to disk via `HighSpeedSwap` and read back on demand.
    /// `None` (default) preserves the existing behavior — sequences hold
    /// every block in HBM forever and rely on `--swap-space-gb` for
    /// admission control. The `try_evict_oldest_for_seq` helper below is
    /// only valid when this is `Some`.
    pub cache_blocks_per_seq: Option<u32>,
}

impl KvCacheConfig {
    /// Resolve the effective dtype for a given attention layer index.
    pub fn dtype_for_layer(&self, layer_idx: usize) -> KvCacheDtype {
        if layer_idx < self.layer_dtypes.len() {
            self.layer_dtypes[layer_idx]
        } else {
            self.dtype
        }
    }

    /// Bytes per block (K or V, not both) for a specific dtype and
    /// (num_kv_heads, head_dim) pair. Per-layer callers pass their layer's
    /// actual dimensions; homogeneous callers pass the global values.
    fn block_bytes_dims(&self, dtype: KvCacheDtype, nkv: usize, hd: usize) -> usize {
        let elems = self.block_size * nkv * hd;
        match dtype {
            KvCacheDtype::Bf16
            | KvCacheDtype::Bf16KTurbo4V
            | KvCacheDtype::Bf16KTurbo3V
            | KvCacheDtype::Bf16KTurbo2V => elems * 2,
            KvCacheDtype::Fp8
            | KvCacheDtype::Fp8KTurbo4V
            | KvCacheDtype::Fp8KTurbo3V
            | KvCacheDtype::Fp8KTurbo2V => elems,
            KvCacheDtype::Nvfp4
            | KvCacheDtype::Turbo4
            | KvCacheDtype::Turbo4KTurbo3V
            | KvCacheDtype::Turbo4KTurbo8V => {
                // Both NVFP4 and Turbo4 use 4-bit data + FP8 per-group scales.
                // Same byte layout, different codebook (E2M1 vs Lloyd-Max).
                let data = elems / 2; // 2 nibbles per byte
                let num_groups = elems / NVFP4_GROUP_SIZE;
                data + num_groups // +1 FP8 scale byte per group
            }
            KvCacheDtype::Turbo3 | KvCacheDtype::Turbo3KTurbo8V => {
                // 3-bit WHT + Lloyd-Max (8 levels). Packed: 8 values in 3 bytes.
                let data = elems * 3 / 8;
                let num_groups = elems / NVFP4_GROUP_SIZE;
                data + num_groups
            }
            KvCacheDtype::Turbo2 => {
                // 2-bit WHT + Lloyd-Max (4 levels). Packed: 4 values per byte.
                let data = elems / 4;
                let num_groups = elems / NVFP4_GROUP_SIZE;
                data + num_groups
            }
            KvCacheDtype::Turbo8 => {
                // WHT + FP8 E4M3 data + per-group BF16 scales.
                // 2026-04-28: scales upgraded from FP8 (1 byte) to BF16 (2 bytes)
                // because FP8's ~12% per-scale relative error compounds
                // catastrophically across MiniMax M2.7's 58 Turbo8 layers
                // (gibberish output). BF16 scales (~0.4% relative error)
                // keep compounding tractable. ~6% extra cache memory for
                // a 256× precision improvement on the per-group scaling.
                let num_groups = elems / NVFP4_GROUP_SIZE;
                elems + num_groups * 2 // 1 byte data + BF16 scale per group
            }
        }
    }

    /// V-side bytes per block for asymmetric dtypes; equals block_bytes_dims
    /// for symmetric dtypes. Use this when allocating the V pool separately
    /// from K. Once real asym kernels land, callers should switch to:
    ///   k_size = block_bytes_dims(kv_pair().0, ...)
    ///   v_size = block_bytes_dims(kv_pair().1, ...)
    /// and allocate the two pools independently.
    #[allow(dead_code)]
    pub fn v_block_bytes_dims(&self, dtype: KvCacheDtype, nkv: usize, hd: usize) -> usize {
        let (_, v) = dtype.kv_pair();
        self.block_bytes_dims(v, nkv, hd)
    }

    /// K-side bytes per block for a specific (asym-aware) dtype and dims.
    /// For symmetric dtypes, returns the same value as `block_bytes_dims`.
    /// For asymmetric, returns the K-component (e.g. Bf16KTurbo3V → bf16 bytes).
    #[allow(dead_code)]
    pub fn k_block_bytes_dims(&self, dtype: KvCacheDtype, nkv: usize, hd: usize) -> usize {
        let (k, _) = dtype.kv_pair();
        self.block_bytes_dims(k, nkv, hd)
    }

    /// K-side bytes per block for a specific attention layer.
    /// Replaces the legacy single-stride view for asym dtypes by routing
    /// through the K component of the dtype pair.
    pub fn k_block_bytes_for_layer(&self, layer_idx: usize) -> usize {
        let (nkv, hd) = self.dims_for_layer(layer_idx);
        let (k, _) = self.dtype_for_layer(layer_idx).kv_pair();
        self.block_bytes_dims(k, nkv, hd)
    }

    /// V-side bytes per block for a specific attention layer.
    /// For symmetric dtypes this equals `k_block_bytes_for_layer`.
    pub fn v_block_bytes_for_layer(&self, layer_idx: usize) -> usize {
        let (nkv, hd) = self.dims_for_layer(layer_idx);
        let (_, v) = self.dtype_for_layer(layer_idx).kv_pair();
        self.block_bytes_dims(v, nkv, hd)
    }

    /// Legacy name: bytes per block using global dims and a given dtype.
    /// Used when the caller has a uniform KV geometry; prefer
    /// `block_bytes_for_layer` for layer-aware paths.
    fn block_bytes_for_dtype(&self, dtype: KvCacheDtype) -> usize {
        self.block_bytes_dims(dtype, self.num_kv_heads, self.head_dim)
    }

    /// Bytes per block per layer (K or V, not both), using the uniform dtype.
    pub fn block_bytes(&self) -> usize {
        self.block_bytes_for_dtype(self.dtype)
    }

    /// (num_kv_heads, head_dim) for a specific attention layer.
    /// Falls back to the global values when no per-layer override is set.
    pub fn dims_for_layer(&self, layer_idx: usize) -> (usize, usize) {
        if layer_idx < self.layer_dims.len() {
            self.layer_dims[layer_idx]
        } else {
            (self.num_kv_heads, self.head_dim)
        }
    }

    /// Bytes per block for a specific attention layer (K or V, not both).
    /// Uses per-layer (nkv, hd) from `layer_dims` when set — this lets
    /// heterogeneous layers (Gemma-4 sliding vs full) allocate tight,
    /// correctly-sized pools so the kernel's per-layer stride matches
    /// the allocation layout.
    pub fn block_bytes_for_layer(&self, layer_idx: usize) -> usize {
        let (nkv, hd) = self.dims_for_layer(layer_idx);
        self.block_bytes_dims(self.dtype_for_layer(layer_idx), nkv, hd)
    }

    /// Bytes per block per layer (K + V combined).
    pub fn block_bytes_kv(&self) -> usize {
        self.block_bytes() * 2
    }

    /// Sum of K+V block bytes across all layers for one block slot.
    /// Accounts for mixed dtypes when layer_dtypes is set AND asym K/V splits.
    pub fn block_bytes_kv_all_layers(&self) -> usize {
        (0..self.num_layers)
            .map(|i| self.k_block_bytes_for_layer(i) + self.v_block_bytes_for_layer(i))
            .sum()
    }

    /// Cache stride in elements (for FP8/BF16 kernels).
    /// Elements per block = block_size * num_kv_heads * head_dim.
    pub fn cache_stride_elements(&self) -> usize {
        self.block_size * self.num_kv_heads * self.head_dim
    }

    /// NVFP4 data section bytes per block (packed E2M1 nibbles).
    pub fn nvfp4_data_bytes(&self) -> usize {
        self.block_size * self.num_kv_heads * self.head_dim / 2
    }

    /// NVFP4 scale section bytes per block (FP8 per-group scales).
    pub fn nvfp4_scale_bytes(&self) -> usize {
        self.block_size * self.num_kv_heads * self.head_dim / NVFP4_GROUP_SIZE
    }

    /// Turbo4 data section bytes (same layout as NVFP4: 4-bit packed).
    pub fn turbo4_data_bytes(&self) -> usize {
        self.nvfp4_data_bytes()
    }

    /// Turbo4 scale section bytes (same layout as NVFP4: FP8 per-group).
    pub fn turbo4_scale_bytes(&self) -> usize {
        self.nvfp4_scale_bytes()
    }

    /// Turbo3 data section bytes (3-bit packed: 8 values in 3 bytes).
    pub fn turbo3_data_bytes(&self) -> usize {
        let elems = self.block_size * self.num_kv_heads * self.head_dim;
        elems * 3 / 8
    }

    /// Turbo3 scale section bytes (FP8 per-group, same as turbo4).
    pub fn turbo3_scale_bytes(&self) -> usize {
        self.nvfp4_scale_bytes()
    }

    /// Turbo2 data section bytes (2-bit packed: 4 values per byte).
    pub fn turbo2_data_bytes(&self) -> usize {
        let elems = self.block_size * self.num_kv_heads * self.head_dim;
        elems / 4
    }

    /// Turbo2 scale section bytes (FP8 per-group, same as turbo3/turbo4).
    pub fn turbo2_scale_bytes(&self) -> usize {
        self.nvfp4_scale_bytes()
    }

    /// Turbo8 data section bytes (FP8 E4M3: 1 byte per element).
    pub fn turbo8_data_bytes(&self) -> usize {
        self.block_size * self.num_kv_heads * self.head_dim
    }

    /// Turbo8 scale section bytes — **BF16 per-group scales** (2 bytes
    /// each, vs the 1-byte FP8 scales NVFP4/Turbo3/Turbo4 use). The
    /// BF16 upgrade is what makes Turbo8 viable across many-layer models
    /// like MiniMax M2.7 (58 Turbo8 layers under auto HP=2). Returns
    /// `(num_groups) * 2` bytes total.
    pub fn turbo8_scale_bytes(&self) -> usize {
        // num_groups = elems / GROUP_SIZE; each scale is 2 bytes (BF16).
        self.nvfp4_scale_bytes() * 2
    }
}

/// Per-layer KV cache pool.
struct LayerPool {
    k_pool: DevicePtr,
    v_pool: DevicePtr,
    /// Stride between K blocks in bytes (may differ from V for asym dtypes).
    k_block_stride: usize,
    /// Stride between V blocks in bytes (may differ from K for asym dtypes).
    v_block_stride: usize,
    /// Effective dtype for this layer.
    dtype: KvCacheDtype,
}

/// Paged KV cache across all attention layers.
pub struct PagedKvCache {
    layers: Vec<LayerPool>,
    num_blocks: usize,
    free_blocks: Vec<u32>,
    /// Per-block reference count. Enables shared blocks (prefix caching).
    /// Default: 1 on alloc, freed when decremented to 0.
    block_ref_counts: Vec<u32>,
    config: KvCacheConfig,
    /// Per-block refcount event history (`ATLAS_KV_TRACE=1`; inert otherwise).
    trace: block_trace::BlockTrace,
}

mod block_trace;
mod catalog;
mod paged_impl;
/// Release both pools of every layer.
///
/// Each layer allocates its K and V pools separately, so freeing per layer is
/// correct. The block bookkeeping (`free_blocks`, `block_ref_counts`) is host
/// state indexing into those pools — cleared with them so a released cache
/// cannot hand out a block into freed memory.
impl atlas_core::scope::ModelResource<dyn crate::gpu::GpuBackend> for PagedKvCache {
    fn label(&self) -> &'static str {
        "kv cache"
    }

    fn release(&mut self, gpu: &dyn crate::gpu::GpuBackend) -> anyhow::Result<()> {
        let mut first_error = None;
        for layer in self.layers.drain(..) {
            for ptr in [layer.k_pool, layer.v_pool] {
                if let Err(e) = gpu.free(ptr)
                    && first_error.is_none()
                {
                    first_error = Some(e);
                }
            }
        }
        self.free_blocks.clear();
        self.block_ref_counts.clear();
        self.num_blocks = 0;
        match first_error {
            Some(e) => Err(e),
            None => Ok(()),
        }
    }
}

#[cfg(test)]
mod tests;

#[cfg(test)]
mod tests_tq_plus;