memra_kv/

lib.rs

//! memra-kv — the dual KV/recurrent cache, extracted (Phase D, ARCHITECTURE-H100.md §5).
//!
//! Moved VERBATIM from memra-engine/src/cache.rs behind the `KvDev` seam: the cache only
//! ever needed 7 device ops (alloc/copy/set), so the trait is that surface and nothing
//! more. The append/dequant KERNELS stay in the engine fatbins — this crate owns the
//! structure, sizing math, and the KV format policy (env-selected, shared by the engine's
//! fatbin router and every cache consumer). memra-engine re-exports this as `cache` so
//! call sites are unchanged.


// ---------------- KV format policy (env-selected; moved from memra-engine) ----------------

/// Env-selected KV cache formats (MEMRA_KV_K / MEMRA_KV_V). The engine's flash-fatbin router
/// and the cache sizing below MUST agree — both read this one function.
pub fn kv_cache_formats() -> (&'static str, &'static str) {
    static F: std::sync::OnceLock<(&'static str, &'static str)> = std::sync::OnceLock::new();
    *F.get_or_init(|| {
        let k = match std::env::var("MEMRA_KV_K").as_deref() {
            Ok("fp8") => "fp8",
            Ok("q8_0") | Ok("") | Err(_) => "q8_0",
            Ok(o) => panic!("MEMRA_KV_K={o} unsupported (q8_0 | fp8)"),
        };
        let v = match std::env::var("MEMRA_KV_V").as_deref() {
            Ok("q4_0") => "q4_0",
            Ok("fp8") => "fp8",
            Ok("q5_1") | Ok("") | Err(_) => "q5_1",
            Ok(o) => panic!("MEMRA_KV_V={o} unsupported (q5_1 | q4_0 | fp8)"),
        };
        if (k, v) != ("q8_0", "q5_1") {
            eprintln!("[memra] KV cache format: K={k} V={v} (non-default — new numeric config)");
        }
        (k, v)
    })
}

/// Per-32-element block bytes for the selected (K, V) formats.
pub fn kv_blk_bytes() -> (usize, usize) {
    let (k, v) = kv_cache_formats();
    let kb = match k { "fp8" => 32, _ => 34 };
    let vb = match v { "q4_0" => 18, "fp8" => 32, _ => 24 };
    (kb, vb)
}

/// FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep
/// their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.
pub fn gkv_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_GKV").map(|v| v != "0").unwrap_or(true))
}

/// FP8-WINDOWED switch (MEMRA_GEMMA_WKV; serving-mode default): SPEC serving (MEMRA_DRAFT
/// set) -> OFF, plain -> ON — the acceptance-vs-depth record lives on the engine-side
/// history of `Engine::wkv_on` (git). Explicit env always wins.
pub fn wkv_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("MEMRA_GEMMA_WKV").map(|v| v != "0")
        .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err()))
}

/// Per-model FP8-KV door (-1 = unset → env/default off; 0 = off; 1 = on). Set at qwen
/// model load: the 2026-07-12 arc closed per-model — 9B +0.7-4% scaling with depth,
/// 27B flat (weight-bound), 35B −2% (fp8 format-gates its v3 dp4a lane off). Explicit
/// MEMRA_KV_FP8 wins. 9B adoption attempt REVERTED by measurement 2026-07-29 (−1% at 12k
/// on the then-current build) — loaders currently store 0.
pub static KV_FP8_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);

/// QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set
/// at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8
/// module.
pub fn kv_fp8_on() -> bool {
    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
    if let Some(v) = *ENV.get_or_init(|| std::env::var("MEMRA_KV_FP8").ok()
        .map(|v| v == "1")) { return v; }
    matches!(KV_FP8_FORCE.load(std::sync::atomic::Ordering::Relaxed), 1)
}

// ---------------- the device seam ----------------

/// The 7 device ops the cache needs — nothing more. Implemented by the engine (and by
/// any future backend); all ops are stream-ordered on the implementor's worker stream.
pub trait KvDev {
    fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>;
    fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>;
    fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>>;
    fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>>;
    fn clone_dtod(&self, src: &CudaSlice<f32>) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>;
    fn copy_into(&self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaSlice<f32>, len: usize)
                 -> Result<(), Box<dyn std::error::Error>>;
    fn set_i32_one(&self, d: &mut CudaSlice<i32>, v: i32) -> Result<(), Box<dyn std::error::Error>>;
}

use memra_gguf::config::{LayerKind, ModelConfig};
use cudarc::driver::CudaSlice;

/// Per-full-attn-layer growing KV cache, resident on GPU. QUANTIZED (KVQUANT-PLAN §B):
/// K stored q8_0 (34 B/32 elem), V stored q5_1 (24 B/32 elem). Per-token byte layout keeps the
/// [token, kv_head, dim] element order so a 32-block never straddles a head (assert head_dim%32==0).
/// Element-within-token index = kv_head*head_dim + d; block = idx/32; lane = idx%32.
pub struct KvLayer {
    pub k: CudaSlice<u8>,   // q8_0 packed, capacity max_ctx*k_tok_bytes
    pub v: CudaSlice<u8>,   // q5_1 packed, capacity max_ctx*v_tok_bytes
    pub kv_dim_k: usize,    // head_dim_k * n_head_kv  (K elements per token)
    pub kv_dim_v: usize,    // head_dim_v * n_head_kv  (V elements per token)
    pub k_tok_bytes: usize, // (kv_dim_k/32)*34
    pub v_tok_bytes: usize, // (kv_dim_v/32)*24
    pub len: usize,
    /// Device-resident mirror of `len` (CUDA-GRAPH-PLAN Phase 2). Holds the KV write SLOT for the
    /// append-dc kernel (old len, before this step's append); after `inc_seqlen` it holds the new
    /// len == t_kv for fa_decode_dc. Kept in lock-step with the host `len`. i32[1].
    pub len_d: CudaSlice<i32>,
}

/// Per-linear-attn-layer fixed recurrent state.
/// conv_state and ssm_state are BOTH kept RESIDENT on GPU — the conv ring assemble + roll runs
/// on-device (conv_assemble_and_roll), so there is no per-step dtoh/htod for either.
pub struct RecurLayer {
    pub conv_state: CudaSlice<f32>, // GPU [conv_dim, d_conv-1] (channel c, tap j at c*pad + j)
    pub ssm_state: CudaSlice<f32>,  // GPU [d_state, d_state, num_v] transposed M[col][i]
    /// PERSISTENT second SSM-state buffer for the gdn-scan double buffer (DECODE DETERMINISM FIX).
    /// gdn_scan needs DISTINCT in/out state buffers. The old eager path allocated a fresh
    /// `state_scratch` via `e.uninit` every step and swapped its pointer into `ssm_state`; that
    /// per-step alloc/free churned the stream-ordered async pool, and the freed prior `ssm_state`
    /// block was recycled by the next step's scratch while a kernel referencing the swapped-in state
    /// was still in flight — a use-after-reuse that produced RUN-TO-RUN nondeterministic decode
    /// (two identical prompt primes diverged). We instead PING-PONG between two STABLE resident
    /// buffers (no per-step alloc/free, no pool churn): step writes into the spare, then swaps the
    /// two owned buffers in place. Stable pointers, identical math. Sized like `ssm_state`.
    pub ssm_state_alt: CudaSlice<f32>,
}

pub struct Cache {
    pub kv: Vec<Option<KvLayer>>,
    pub recur: Vec<Option<RecurLayer>>,
    pub pos: usize,
    pub max_ctx: usize,
    /// BATCHED-TICK increment 2 component 3 (lean logits, 2026-08-01): device-side park of
    /// this session's LAST logits row. Device-sampled rows in the batched serving tick skip
    /// the [n_vocab] logits D2H entirely; the tick instead dtod-copies the row here (device
    /// bandwidth, ~µs) so the ONE consumer that truly needs the final row — the KV-reuse
    /// pool's park-at-retire (an empty-suffix resume samples from parked last_logits) —
    /// can D2H it once at retire. Lazily allocated on the first lean tick; None on every
    /// non-lean path (zero cost). Travels with the Cache into the reuse pool.
    pub last_logits_dev: Option<CudaSlice<f32>>,
    /// DFlash tap sink (dflash lane, 2026-07-13): when armed, the gemma4 verify/prime
    /// trunks copy the residual stream AFTER each tapped layer into `buf` rows
    /// ([t, n_taps*hidden] row-major — the drafter fc input layout). None on every
    /// non-dflash path (zero cost).
    pub dflash_taps: Option<DflashTapSink>,
}

/// The context-linear K/V layout for one full-attention layer. This is the single sizing source
/// used by both `Cache::new_inner` and `cache_bytes_per_token`: admission must never reimplement
/// Gemma's per-layer geometry or the active KV-format doors independently from the allocator.
fn full_attention_kv_layout(cfg: &ModelConfig, il: u32) -> (usize, usize, usize, usize) {
    debug_assert_eq!(cfg.layer_kind(il), LayerKind::FullAttention);
    let n_head_kv = cfg.n_head_kv as usize;
    let (kv_dim_k, kv_dim_v) = match &cfg.gemma4 {
        Some(g) => {
            let hd = if g.swa_pattern[il as usize] {
                g.key_length_swa
            } else {
                g.key_length_global
            } as usize;
            // E4B ships a SCALAR head_count_kv (per-layer vec empty; scalar = 2 in
            // the gguf, landing in cfg.n_head_kv): kv_dim = hd * 2 for BOTH kinds —
            // swa 2x256 = 512, global 2x512 = 1024. The old fallback used
            // key_length_global (512) for both, which HALVED the global layers' K/V
            // (the attn writes wk.out_features = 1024 rows): every E4B global layer
            // stored/attended half its K/V and the batched append read row strides
            // wrong — THE cross-mode maxdiff-30 root (2026-07-12 bisect, il=5 slot-1
            // byte forensics). 26B/31B keep the per-layer vec.
            let d = match g.head_count_kv.get(il as usize) {
                Some(n) => hd * *n as usize,
                None => hd * n_head_kv,
            };
            (d, d)
        }
        None => (
            cfg.head_dim_k as usize * n_head_kv,
            cfg.head_dim_v as usize * n_head_kv,
        ),
    };
    assert!(
        kv_dim_k % 32 == 0 && kv_dim_v % 32 == 0,
        "KVQUANT requires per-layer kv_dim_k%32==0 && kv_dim_v%32==0 \
         (layer {il}: k={kv_dim_k} v={kv_dim_v})"
    );
    let (kbb, vbb) = kv_blk_bytes();
    let g4_global_fp8 = gkv_on()
        && cfg
            .gemma4
            .as_ref()
            .is_some_and(|g| !g.swa_pattern[il as usize]);
    let g4_windowed_fp8 = wkv_on()
        && cfg
            .gemma4
            .as_ref()
            .is_some_and(|g| g.swa_pattern[il as usize]);
    let qwen_fp8 = kv_fp8_on() && cfg.gemma4.is_none();
    let (kbb_l, vbb_l) = if g4_global_fp8 || g4_windowed_fp8 || qwen_fp8 {
        (32, 32)
    } else {
        (kbb, vbb)
    };
    (kv_dim_k, kv_dim_v, kbb_l, vbb_l)
}

/// Context-linear bytes allocated by one trunk cache token.
///
/// Fixed allocations (the 8-byte plane tail pads, `len_d`, recurrent state, and optional lazy
/// buffers) are deliberately excluded. Admission adds their measured high-water residual as a
/// request-independent activation term; multiplying this coefficient by the request's own
/// `ctx_cap` exactly mirrors the context-scaled allocations in `Cache::new_inner`.
pub fn cache_bytes_per_token(cfg: &ModelConfig) -> usize {
    let shared = cfg.gemma4.as_ref().map(|g| g.shared_kv_layers).unwrap_or(0);
    (0..cfg.n_layer)
        .filter(|&il| cfg.layer_kind(il) == LayerKind::FullAttention)
        .filter(|&il| shared == 0 || il < cfg.n_layer - shared)
        .map(|il| {
            let (kv_dim_k, kv_dim_v, kbb, vbb) = full_attention_kv_layout(cfg, il);
            (kv_dim_k / 32) * kbb + (kv_dim_v / 32) * vbb
        })
        .sum()
}

/// See [`Cache::dflash_taps`]. Armed per forward by the dflash round (t = that forward's
/// row count); the trunk writes tap slot s of row r at buf[r*n_taps*hidden + s*hidden ..].
pub struct DflashTapSink {
    pub layer_ids: Vec<usize>,
    pub buf: CudaSlice<f32>,
    pub hidden: usize,
    pub t: usize,
}

/// Snapshot of the dual cache taken BEFORE a spec-decode draft+verify round (MTP-PLAN §C/§D.4).
/// - Full-attn KV: only the per-layer `len` is recorded; rollback truncates (append-only,
///   position-addressed — no copy). C.1.
/// - Linear-attn conv/ssm: real device-to-device COPIES of the recurrent state, because those
///   buffers are mutated IN PLACE by the verify pass and have no position index to truncate. C.2.
///   (CudaSlice::clone is an Arc refcount, NOT a buffer copy — so we alloc fresh + memcpy_dtod.)
pub struct CacheSnapshot {
    pub kv_len: Vec<Option<usize>>, // per layer (Some for full-attn layers)
    pub conv: Vec<Option<CudaSlice<f32>>>, // per layer (Some for linear-attn layers, D2D copy)
    pub ssm: Vec<Option<CudaSlice<f32>>>,
    pub pos: usize,
}

impl Cache {
    /// Allocate GPU-resident caches sized by arch + max context.
    pub fn new(
        e: &impl KvDev,
        cfg: &ModelConfig,
        max_ctx: usize,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        Self::new_inner(&|_| e, cfg, max_ctx)
    }

    /// M1-PP2 increment 2 (stage-owned KV): layers [0, split) allocate through `dev0`,
    /// layers [split, n) through `dev1` — each pipeline stage's cache lives on the
    /// device that runs the stage. With dev0 == dev1 this is byte-for-byte `new`
    /// (the single-device plumbing gate). Sizing math is IDENTICAL either way.
    pub fn new_pp2(
        dev0: &dyn KvDev,
        dev1: &dyn KvDev,
        split: usize,
        cfg: &ModelConfig,
        max_ctx: usize,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        Self::new_inner(&|il| if il < split { dev0 } else { dev1 }, cfg, max_ctx)
    }

    /// M2 N-stage twin of `new_pp2`: `fence` is the stage map from `memra_engine::pp::
    /// pp_cuts` ([0, c1, .., n_trunk]); layer il allocates through the engine of the
    /// stage that runs it. Layers at/beyond the fence end (MTP/NextN blocks) allocate
    /// through the LAST stage. Sizing math is IDENTICAL to `new` — only the allocating
    /// device varies.
    pub fn new_ppn<'a>(
        devs: &[&'a dyn KvDev],
        fence: &[usize],
        cfg: &ModelConfig,
        max_ctx: usize,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        assert_eq!(devs.len() + 1, fence.len(), "ppn cache: devs vs fence mismatch");
        let pick = |il: usize| -> &dyn KvDev {
            let s = match fence[1..fence.len() - 1].binary_search(&il) {
                Ok(k) => k + 1,
                Err(k) => k,
            };
            devs[s.min(devs.len() - 1)]
        };
        Self::new_inner(&pick, cfg, max_ctx)
    }

    /// Shared allocation walk: `pick(il)` supplies the device that OWNS layer il's
    /// cache state (always the same device outside the pp2 door).
    fn new_inner<'a>(
        pick: &dyn Fn(usize) -> &'a dyn KvDev,
        cfg: &ModelConfig,
        max_ctx: usize,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let n = cfg.n_layer as usize;
        let mut kv = Vec::with_capacity(n);
        let mut recur = Vec::with_capacity(n);
        let head_dim_k = cfg.head_dim_k as usize;
        let head_dim_v = cfg.head_dim_v as usize;
        assert!(head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
                "KVQUANT requires head_dim_k%32==0 && head_dim_v%32==0 (got k={head_dim_k} v={head_dim_v})");
        let (conv_dim, d_state, num_v, d_conv) = if let Some(s) = &cfg.ssm {
            let num_k = s.group_count as usize;
            let num_v = s.time_step_rank as usize;
            let ds = s.state_size as usize;
            (
                ds * num_k * 2 + ds * num_v,
                ds,
                num_v,
                s.conv_kernel as usize,
            )
        } else {
            (0, 0, 0, 0)
        };
        for il in 0..cfg.n_layer {
            // stage-owned allocation (pp2): the device that runs this layer allocates it.
            let e = pick(il as usize);
            // E4B KV-SHARING: the trailing shared_kv_layers have no k/v of their own — they
            // attend an earlier layer's cache (hybrid_forward resolves the target). No KvLayer
            // here: any accidental use is a loud unwrap at bring-up, and rewind/len loops
            // (iter_mut().flatten()) skip None naturally.
            let g4_shared = cfg.gemma4.as_ref().map(|g| g.shared_kv_layers).unwrap_or(0);
            if g4_shared > 0 && il >= cfg.n_layer - g4_shared {
                kv.push(None);
                recur.push(None);
                continue;
            }
            match cfg.layer_kind(il) {
                LayerKind::FullAttention => {
                    // Gemma per-layer geometry and every KV-format door are resolved by the same
                    // helper admission uses for its analytic byte coefficient.
                    let (kv_dim_k, kv_dim_v, kbb_l, vbb_l) =
                        full_attention_kv_layout(cfg, il);
                    let k_tok_bytes = (kv_dim_k / 32) * kbb_l;
                    let v_tok_bytes = (kv_dim_v / 32) * vbb_l;
                    kv.push(Some(KvLayer {
                        // +8B tail pad: the v4 stage's aligned funnelshift window reads up to
                        // 4B past the final block (PR #3's finding, adopted pad-style — the
                        // expert-dot precedent; zero hot-loop branches, values discarded).
                        k: e.alloc_u8(max_ctx * k_tok_bytes + 8)?,
                        v: e.alloc_u8(max_ctx * v_tok_bytes + 8)?,
                        kv_dim_k,
                        kv_dim_v,
                        k_tok_bytes,
                        v_tok_bytes,
                        len: 0,
                        len_d: e.htod_i32(&[0])?,
                    }));
                    recur.push(None);
                }
                LayerKind::LinearAttention => {
                    kv.push(None);
                    recur.push(Some(RecurLayer {
                        conv_state: e.zeros(conv_dim * (d_conv - 1))?,
                        ssm_state: e.zeros(d_state * d_state * num_v)?,
                        ssm_state_alt: e.zeros(d_state * d_state * num_v)?,
                    }));
                }
            }
        }
        Ok(Cache { kv, recur, pos: 0, max_ctx, dflash_taps: None, last_logits_dev: None })
    }

    /// Snapshot the dual cache before a spec-decode draft+verify round (MTP-PLAN §C/§D.4).
    /// Records each full-attn `len` (cheap) and makes a REAL device copy of each linear-attn
    /// conv_state/ssm_state (a fresh alloc + memcpy_dtod — NOT an Arc clone).
    pub fn snapshot(&self, e: &impl KvDev) -> Result<CacheSnapshot, Box<dyn std::error::Error>> {
        let n = self.kv.len();
        let mut kv_len = Vec::with_capacity(n);
        let mut conv = Vec::with_capacity(n);
        let mut ssm = Vec::with_capacity(n);
        for il in 0..n {
            match &self.kv[il] {
                Some(kvl) => kv_len.push(Some(kvl.len)),
                None => kv_len.push(None),
            }
            match &self.recur[il] {
                Some(rl) => {
                    conv.push(Some(e.clone_dtod(&rl.conv_state)?));
                    ssm.push(Some(e.clone_dtod(&rl.ssm_state)?));
                }
                None => {
                    conv.push(None);
                    ssm.push(None);
                }
            }
        }
        Ok(CacheSnapshot {
            kv_len,
            conv,
            ssm,
            pos: self.pos,
        })
    }

    /// PERSISTENT-BUFFER snapshot (spec-decode hot loop): refresh `snap` IN PLACE — same values as
    /// `snapshot()` but the conv/ssm device buffers are reused across rounds (D2D copy-into, ZERO
    /// allocations vs 2 fresh clones per linear layer per round). `snap` must come from a prior
    /// `snapshot()` of THIS cache (same layer shapes).
    pub fn snapshot_into(
        &self,
        e: &impl KvDev,
        snap: &mut CacheSnapshot,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let n = self.kv.len();
        for il in 0..n {
            snap.kv_len[il] = self.kv[il].as_ref().map(|kvl| kvl.len);
            if let Some(rl) = &self.recur[il] {
                let dc = snap.conv[il]
                    .as_mut()
                    .expect("snapshot_into: shape mismatch (conv)");
                let ds = snap.ssm[il]
                    .as_mut()
                    .expect("snapshot_into: shape mismatch (ssm)");
                let (cn, sn) = (rl.conv_state.len(), rl.ssm_state.len());
                e.copy_into(dc, 0, &rl.conv_state, cn)?;
                e.copy_into(ds, 0, &rl.ssm_state, sn)?;
            }
        }
        snap.pos = self.pos;
        Ok(())
    }

    /// Roll the cache back to exactly `snap.pos + accept_len` committed tokens (MTP-PLAN §C).
    /// - Full-attn KV (C.1): set len = snapshot_len + accept_len (truncate, no copy).
    /// - Linear-attn (C.2): RESTORE the snapshot conv/ssm (real D2D copy back into the resident
    ///   buffers). The caller must then REPLAY the `accept_len` committed tokens through the full
    ///   T=1 decode path to rebuild the recurrent state for those positions. We restore (not
    ///   replay here) because replay needs the model; this only resets state to the pre-round value.
    /// `cache.pos` is set to `snap.pos` so the caller's replay advances it back to the commit point.
    pub fn rollback(
        &mut self,
        e: &impl KvDev,
        snap: &CacheSnapshot,
        accept_len: usize,
    ) -> Result<(), Box<dyn std::error::Error>> {
        for il in 0..self.kv.len() {
            if let (Some(kvl), Some(saved)) = (self.kv[il].as_mut(), snap.kv_len[il]) {
                kvl.len = saved + accept_len;
                // keep the device mirror in lock-step (CUDA-GRAPH-PLAN Phase 2). Set IN PLACE
                // (stable pointer): a fresh htod_i32 would reallocate len_d, but its old pointer is
                // baked into the captured decode graph's append/inc/fa_decode kernels — replacing it
                // strands the graph on a freed buffer (stale-pointer hazard). memcpy_htod in place.
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            if let Some(rl) = self.recur[il].as_mut() {
                if let Some(c) = &snap.conv[il] {
                    e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
                }
                if let Some(s) = &snap.ssm[il] {
                    e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
                }
            }
        }
        self.pos = snap.pos;
        Ok(())
    }
}