Struct Cache 

pub struct Cache {
    pub kv: Vec<Option<KvLayer>>,
    pub recur: Vec<Option<RecurLayer>>,
    pub pos: usize,
    pub max_ctx: usize,
    pub last_logits_dev: Option<CudaSlice<f32>>,
    pub dflash_taps: Option<DflashTapSink>,
}

Fields

kv: Vec<Option<KvLayer>>

recur: Vec<Option<RecurLayer>>

pos: usize

max_ctx: usize

last_logits_dev: Option<CudaSlice<f32>>

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.

dflash_taps: Option<DflashTapSink>

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

Implementations

impl Cache

pub fn new( e: &impl KvDev, cfg: &ModelConfig, max_ctx: usize, ) -> Result<Cache, Box<dyn Error>>

Allocate GPU-resident caches sized by arch + max context.

pub fn new_pp2( dev0: &dyn KvDev, dev1: &dyn KvDev, split: usize, cfg: &ModelConfig, max_ctx: usize, ) -> Result<Cache, Box<dyn Error>>

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_ppn<'a>( devs: &[&'a dyn KvDev], fence: &[usize], cfg: &ModelConfig, max_ctx: usize, ) -> Result<Cache, Box<dyn Error>>

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 has_swa_ring(&self) -> bool

pub fn can_rollback(&self, snap: &CacheSnapshot, accept_len: usize) -> bool

pub fn snapshot(&self, e: &impl KvDev) -> Result<CacheSnapshot, Box<dyn Error>>

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_into( &self, e: &impl KvDev, snap: &mut CacheSnapshot, ) -> Result<(), Box<dyn Error>>

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 rollback( &mut self, e: &impl KvDev, snap: &CacheSnapshot, accept_len: usize, ) -> Result<(), Box<dyn Error>>

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.