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)
}

/// Step35 SWA-ring experiment (default OFF). The first cut is deliberately architecture-scoped:
/// Gemma4's row-0-addressed window kernels cannot consume a rebased ring view.
pub fn swa_ring_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("MEMRA_SWA_RING").as_deref() == Ok("1"))
}

/// With the SWA-ring door open, `prime_chunk_tokens` caps every legal chunk at this bound. The
/// ring carries one whole maximum-size prime chunk in addition to the reader's window.
pub const PRIME_CHUNK_MAX_TOKENS: usize = 4096;
const SWA_VIEW_ALIGNMENT_ROWS: usize = 32;

/// Physical rows required by the Step35 SWA reader contract. Prime starts at
/// `(base_len - (window - 1)) & !31`, so at most 31 masked rows precede the live window.
pub fn swa_ring_rows(window: usize, max_ctx: usize) -> usize {
    max_ctx.min(window + PRIME_CHUNK_MAX_TOKENS + (SWA_VIEW_ALIGNMENT_ROWS - 1))
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct KvRing {
    rows: usize,
    window: usize,
    base: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KvRingAppend {
    Contiguous {
        write_row: usize,
    },
    Rebase {
        src_row: usize,
        keep_rows: usize,
        new_base: usize,
        write_row: usize,
    },
}

impl KvRing {
    pub fn new(rows: usize, window: usize) -> Self {
        assert!(window > 0 && rows > 0, "invalid SWA ring geometry");
        Self {
            rows,
            window,
            base: 0,
        }
    }

    pub fn rows(&self) -> usize {
        self.rows
    }
    pub fn base(&self) -> usize {
        self.base
    }
    pub fn window(&self) -> usize {
        self.window
    }

    /// Plan a contiguous physical append. When the tail would wrap, retain the caller's exact
    /// aligned read prefix at row zero; the following read remains one contiguous CUDA view.
    pub fn append_plan(
        &self,
        len: usize,
        retain_from: usize,
        append_rows: usize,
    ) -> Result<KvRingAppend, String> {
        if len < self.base || retain_from < self.base || retain_from > len {
            return Err(format!(
                "SWA ring lapped required rows (base {}, retain {retain_from}, len {len})",
                self.base
            ));
        }
        let used = len - self.base;
        if used > self.rows {
            return Err(format!(
                "SWA ring state exceeds capacity ({used} > {})",
                self.rows
            ));
        }
        if used.saturating_add(append_rows) <= self.rows {
            return Ok(KvRingAppend::Contiguous {
                write_row: used % self.rows,
            });
        }

        let keep_rows = len - retain_from;
        if keep_rows.saturating_add(append_rows) > self.rows {
            return Err(format!(
                "SWA ring append does not fit (keep {keep_rows} + append {append_rows} > {})",
                self.rows
            ));
        }
        Ok(KvRingAppend::Rebase {
            src_row: retain_from - self.base,
            keep_rows,
            new_base: retain_from,
            write_row: keep_rows,
        })
    }

    pub fn apply_rebase(&mut self, new_base: usize) {
        debug_assert!(new_base >= self.base);
        self.base = new_base;
    }

    pub fn physical_range(
        &self,
        start: usize,
        end: usize,
    ) -> Result<std::ops::Range<usize>, String> {
        if start < self.base || end < start || end - self.base > self.rows {
            return Err(format!(
                "SWA ring view [{start},{end}) is outside resident [{},{})",
                self.base,
                self.base + self.rows
            ));
        }
        let start_row = (start - self.base) % self.rows;
        let len = end - start;
        debug_assert!(
            start_row + len <= self.rows,
            "ring view must be contiguous after rebase"
        );
        Ok(start_row..start_row + len)
    }

    /// A rewind is usable only when the next aligned Step35 window view is still resident.
    pub fn can_rewind_to(&self, len: usize) -> bool {
        let raw = len.saturating_sub(self.window - 1);
        let view_start = raw & !(SWA_VIEW_ALIGNMENT_ROWS - 1);
        view_start >= self.base
    }
}

// ---------------- 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 cudarc::driver::CudaSlice;
use memra_gguf::config::{LayerKind, ModelConfig};

/// 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,
    /// Step35 SWA physical-row state. `len` remains absolute; `None` keeps the original flat
    /// `[0, max_ctx)` addressing contract.
    pub ring: Option<KvRing>,
    /// 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>,
}

impl KvLayer {
    pub fn physical_rows(
        &self,
        start: usize,
        end: usize,
    ) -> Result<std::ops::Range<usize>, String> {
        match &self.ring {
            Some(ring) => ring.physical_range(start, end),
            None => Ok(start..end),
        }
    }
}

/// 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)
}

fn kv_plane_allocation_bytes(rows: usize, token_bytes: usize) -> usize {
    rows * token_bytes + 8
}

/// 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 {
    cache_bytes_per_token_for_layers(cfg, 0, cfg.n_layer as usize)
}

/// Context-linear cache bytes per token owned by layers in `[lo, hi)`. PP admission uses the
/// same layer ranges as `Cache::new_ppn`, so each device is charged for exactly the cache planes
/// it allocates rather than for the aggregate model geometry.
pub fn cache_bytes_per_token_for_layers(cfg: &ModelConfig, lo: usize, hi: usize) -> usize {
    assert!(
        lo <= hi && hi <= cfg.n_layer as usize,
        "cache layer range out of bounds"
    );
    let shared = cfg.gemma4.as_ref().map(|g| g.shared_kv_layers).unwrap_or(0);
    (lo as u32..hi as u32)
        .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()
}

/// Portion of [`cache_bytes_per_token`] whose physical row count is capped by the Step35 SWA
/// ring. Zero with the flag off and for every non-Step35 architecture.
pub fn cache_ring_bytes_per_token(cfg: &ModelConfig) -> usize {
    cache_ring_bytes_per_token_for_layers(cfg, 0, cfg.n_layer as usize)
}

/// Ring-capped portion of [`cache_bytes_per_token_for_layers`] for `[lo, hi)`.
pub fn cache_ring_bytes_per_token_for_layers(cfg: &ModelConfig, lo: usize, hi: usize) -> usize {
    assert!(
        lo <= hi && hi <= cfg.n_layer as usize,
        "cache layer range out of bounds"
    );
    if !swa_ring_on() || !cfg.arch.is_step35() {
        return 0;
    }
    let shared = cfg.gemma4.as_ref().map(|g| g.shared_kv_layers).unwrap_or(0);
    (lo as u32..hi as u32)
        .filter(|&il| cfg.layer_kind(il) == LayerKind::FullAttention)
        .filter(|&il| shared == 0 || il < cfg.n_layer - shared)
        .filter(|&il| {
            cfg.layer_geometry(il)
                .is_some_and(|geometry| geometry.window.is_some())
        })
        .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()
}

/// Physical row cap shared by the Step35 SWA trunk and MTP scratch; zero when no ring is active.
pub fn cache_ring_row_cap(cfg: &ModelConfig) -> usize {
    if !swa_ring_on() || !cfg.arch.is_step35() {
        return 0;
    }
    cfg.geometry
        .as_ref()
        .and_then(|table| table.classes().iter().find_map(|geometry| geometry.window))
        .map(|window| swa_ring_rows(window as usize, usize::MAX))
        .unwrap_or(0)
}

/// 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;
                    let ring = if swa_ring_on() && cfg.arch.is_step35() {
                        cfg.layer_geometry(il)
                            .and_then(|geometry| geometry.window)
                            .map(|window| {
                                let window = window as usize;
                                KvRing::new(swa_ring_rows(window, max_ctx), window)
                            })
                    } else {
                        None
                    };
                    let alloc_rows = ring.as_ref().map(KvRing::rows).unwrap_or(max_ctx);
                    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(kv_plane_allocation_bytes(alloc_rows, k_tok_bytes))?,
                        v: e.alloc_u8(kv_plane_allocation_bytes(alloc_rows, v_tok_bytes))?,
                        kv_dim_k,
                        kv_dim_v,
                        k_tok_bytes,
                        v_tok_bytes,
                        len: 0,
                        ring,
                        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,
        })
    }

    pub fn has_swa_ring(&self) -> bool {
        self.kv.iter().flatten().any(|layer| layer.ring.is_some())
    }

    pub fn can_rollback(&self, snap: &CacheSnapshot, accept_len: usize) -> bool {
        self.kv
            .iter()
            .zip(&snap.kv_len)
            .all(|(layer, saved)| match (layer, saved) {
                (Some(layer), Some(saved)) => layer
                    .ring
                    .as_ref()
                    .is_none_or(|ring| ring.can_rewind_to(saved + accept_len)),
                _ => true,
            })
    }

    /// 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>> {
        if !self.can_rollback(snap, accept_len) {
            return Err(
                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
            );
        }
        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(())
    }
}

#[cfg(test)]
mod swa_ring_tests {
    use super::{KvRing, KvRingAppend, kv_plane_allocation_bytes, swa_ring_rows};

    #[test]
    fn allocation_rows_cover_window_max_prime_and_alignment_slack() {
        assert_eq!(swa_ring_rows(512, 262_144), 512 + 4096 + 31);
        assert_eq!(swa_ring_rows(512, 4096), 4096);
        assert_eq!(
            kv_plane_allocation_bytes(4639, 1088),
            4639 * 1088 + 8,
            "the Step35 session plane allocates ring rows plus the existing tail pad",
        );
    }

    #[test]
    fn ring_matches_flat_bytes_before_wrap() {
        let ring = KvRing::new(swa_ring_rows(512, 262_144), 512);
        let flat: Vec<u32> = (0..1024).collect();
        let mut physical = vec![u32::MAX; ring.rows()];
        let KvRingAppend::Contiguous { write_row } = ring.append_plan(0, 0, flat.len()).unwrap()
        else {
            panic!("first append unexpectedly wrapped")
        };
        physical[write_row..write_row + flat.len()].copy_from_slice(&flat);
        let view = ring.physical_range(0, flat.len()).unwrap();
        assert_eq!(&physical[view], flat.as_slice());
    }

    #[test]
    fn wrap_rebases_the_exact_aligned_prime_view() {
        let mut ring = KvRing::new(swa_ring_rows(512, 262_144), 512);
        let flat: Vec<u32> = (0..8192).collect();
        let mut physical = vec![u32::MAX; ring.rows()];
        let KvRingAppend::Contiguous { write_row } = ring.append_plan(0, 0, 4096).unwrap() else {
            panic!("first prime chunk unexpectedly wrapped")
        };
        physical[write_row..write_row + 4096].copy_from_slice(&flat[..4096]);

        let off = (4096usize - (512 - 1)) & !31usize;
        let KvRingAppend::Rebase {
            src_row,
            keep_rows,
            new_base,
            write_row,
        } = ring.append_plan(4096, off, 4096).unwrap()
        else {
            panic!("second prime chunk did not wrap")
        };
        let retained = physical[src_row..src_row + keep_rows].to_vec();
        physical[..keep_rows].copy_from_slice(&retained);
        ring.apply_rebase(new_base);
        physical[write_row..write_row + 4096].copy_from_slice(&flat[4096..8192]);

        let view = ring.physical_range(off, 8192).unwrap();
        assert_eq!(&physical[view], &flat[off..8192]);
        assert_eq!(ring.base(), off);
    }

    #[test]
    fn rewind_declines_once_the_required_window_was_lapped() {
        let mut ring = KvRing::new(swa_ring_rows(512, 262_144), 512);
        let KvRingAppend::Rebase { new_base, .. } = ring.append_plan(4096, 3584, 4096).unwrap()
        else {
            panic!("expected wrap")
        };
        ring.apply_rebase(new_base);
        assert!(ring.can_rewind_to(4095));
        assert!(!ring.can_rewind_to(4094));
        assert!(!ring.can_rewind_to(0));
    }
}