prefix cache: caching on = caching off at temperature 0 on SSM-hybrid models (fixes #158)

crates/spark-model/src/layer.rs

@@ -380,6 +380,8 @@ pub struct ForwardContext<'a> {
     /// cached K/V, and the replay range [snap_tok, matched) is rewritten
     /// into SHARED prefix-cache blocks — non-exact recompute poisons them
     /// and the drift ratchets across turns (2026-06-10 warm-hit stutter).
+    /// Exact path: register-resident token-sequential kernel (default), WY4
+    /// fallback. Off under the prefill grid (see `grid_replay_is_cold`).
     pub gdn_exact_replay: bool,
     /// Device `[num_tokens]` u32 token IDs for the tokens being processed this
     /// pass, in the SAME order the per-token MoE loop visits them. Required by

crates/spark-model/src/model/trait_impl/decode_checkpoint.rs

@@ -37,7 +37,10 @@ impl TransformerModel {
     /// matched is within `interval` blocks → tiny replay tail. Live SSM state
     /// must be canonical at call time (post-commit on the MTP path).
     pub(super) fn decode_marconi_checkpoint_dispatch(&self, seq: &mut SequenceState) {
+        // R11 fix 1: decode-produced state is not on the prefill pass grid;
+        // under the grid no decode checkpoint is registered.
         if !self.ssm_snapshots.is_enabled()
+            || self.prefill_grid() > 0
             || !self.prefix_cache.is_active()
             || crate::model::mtp_carry::marconi_prefill_only()
             || self.config.num_ssm_layers() == 0

crates/spark-model/src/model/trait_impl/drafter_anchor.rs

@@ -0,0 +1,209 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+//! R11 fix 7: MTP drafter context on a warm (prefix-restored) turn, made
+//! identical to a cold turn.
+//!
+//! A cold turn captures every prompt row's final-layer hidden into
+//! `mtp_prefill_hidden` and, at the end of prefill, builds the drafter KV over
+//! the whole prompt in one `prefill_drafter` call. A warm turn used to adopt
+//! the previous turn's drafter KV (carry), whose rows were made by decode, so
+//! the drafts — and with them the accepted tokens — differed from a cold run.
+//!
+//! Here each grid SSM checkpoint at token `a` also stores the hidden rows
+//! `[0, a)` (host RAM, `a * hidden * 2` bytes). A warm hit restores them into
+//! `mtp_prefill_hidden`, claims the capture generation, and the replay passes
+//! append rows from `a`, so the end-of-prefill consume runs the SAME
+//! whole-prompt `prefill_drafter` a cold run does. A hit whose rows are not
+//! stored is declined (full recompute): correct, slower.
+//!
+//! Memory: one entry per checkpoint slot, in host RAM, bounded by
+//! `ATLAS_MTP_ANCHOR_MAX_MB` (default 512; oldest evicted first). On
+//! unified-memory boxes (GB10, Strix Halo) host RAM is GPU memory, so the
+//! default is kept small: a 30k-token anchor is ~120 MB at hidden 2048, and
+//! an evicted anchor only costs a declined hit (full recompute), never a
+//! wrong answer.
+//! `ATLAS_MTP_ANCHOR=0` restores the old carry behaviour.
+
+use std::collections::{HashMap, VecDeque};
+use std::sync::atomic::Ordering;
+
+use parking_lot::Mutex;
+use spark_runtime::gpu::DevicePtr;
+
+use super::super::types::TransformerModel;
+use crate::traits::SequenceState;
+
+struct Entry {
+    hash: u64,
+    rows: usize,
+    bytes: Vec<u8>,
+}
+
+#[derive(Default)]
+struct Store {
+    map: HashMap<usize, Entry>,
+    order: VecDeque<usize>,
+    total: usize,
+}
+
+static STORE: std::sync::OnceLock<Mutex<Store>> = std::sync::OnceLock::new();
+
+fn store() -> &'static Mutex<Store> {
+    STORE.get_or_init(|| Mutex::new(Store::default()))
+}
+
+/// Default anchor budget in MiB (see the module doc for why it is small).
+const DEFAULT_CAP_MB: usize = 512;
+
+fn cap_bytes() -> usize {
+    static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
+    *C.get_or_init(|| {
+        std::env::var("ATLAS_MTP_ANCHOR_MAX_MB")
+            .ok()
+            .and_then(|v| v.parse::<usize>().ok())
+            .unwrap_or(DEFAULT_CAP_MB)
+            * 1024
+            * 1024
+    })
+}
+
+fn prefix_hash(tokens: &[u32]) -> u64 {
+    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
+    for t in tokens {
+        for b in t.to_le_bytes() {
+            h ^= b as u64;
+            h = h.wrapping_mul(0x0000_0100_0000_01b3);
+        }
+    }
+    h
+}
+
+impl Store {
+    fn remove(&mut self, slot: usize) {
+        if let Some(e) = self.map.remove(&slot) {
+            self.total -= e.bytes.len();
+        }
+        self.order.retain(|&s| s != slot);
+    }
+}
+
+impl TransformerModel {
+    /// Whether warm turns must reproduce the cold drafter context exactly.
+    pub(in crate::model) fn mtp_anchor_on(&self) -> bool {
+        static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+        !self.mtp_prefill_hidden.is_null()
+            && self.proposer.is_some()
+            && self.prefill_grid() > 0
+            // Off under the multi-rank protocol: the rows live on the head
+            // only, so a head-side decline after `pc_agree_restore` would
+            // restore on one rank and recompute on the other.
+            && !self.multi_rank_protocol_active()
+            && !*OFF.get_or_init(|| std::env::var("ATLAS_MTP_ANCHOR").as_deref() == Ok("0"))
+    }
+
+    /// Save hidden rows `[0, rows)` for checkpoint slot `snap_slot`. Only when
+    /// this sequence still owns a capture that covers them.
+    pub(in crate::model) fn mtp_anchor_save(
+        &self,
+        seq: &SequenceState,
+        tokens: &[u32],
+        snap_slot: usize,
+        rows: usize,
+        stream: u64,
+    ) {
+        let mut st = store().lock();
+        st.remove(snap_slot);
+        if !self.mtp_anchor_on() {
+            return;
+        }
+        let captured = self.mtp_prefill_capture_len.load(Ordering::Relaxed);
+        let owns = seq.mtp_capture_gen != 0
+            && seq.mtp_capture_gen == self.mtp_prefill_capture_gen.load(Ordering::Relaxed);
+        if !owns || captured < rows || rows > tokens.len() {
+            return;
+        }
+        let n = rows * self.config.hidden_size * 2;
+        if n > cap_bytes() {
+            return;
+        }
+        while st.total + n > cap_bytes() {
+            let Some(old) = st.order.pop_front() else {
+                break;
+            };
+            if let Some(e) = st.map.remove(&old) {
+                st.total -= e.bytes.len();
+            }
+        }
+        let mut bytes = vec![0u8; n];
+        if let Err(e) = self
+            .gpu
+            .copy_d2h_on_stream(self.mtp_prefill_hidden, &mut bytes, stream)
+        {
+            tracing::warn!("MTP anchor save failed (anchor stays drafter-less): {e:#}");
+            return;
+        }
+        st.total += n;
+        st.order.push_back(snap_slot);
+        st.map.insert(
+            snap_slot,
+            Entry {
+                hash: prefix_hash(&tokens[..rows]),
+                rows,
+                bytes,
+            },
+        );
+    }
+
+    /// True when a warm restore from `snap_slot` at `rows` can rebuild the
+    /// cold drafter context (or when exactness is off).
+    pub(in crate::model) fn mtp_anchor_available(
+        &self,
+        snap_slot: usize,
+        tokens: &[u32],
+        rows: usize,
+    ) -> bool {
+        if !self.mtp_anchor_on() {
+            return true;
+        }
+        let st = store().lock();
+        let ok = st.map.get(&snap_slot).is_some_and(|e| {
+            e.rows == rows && rows <= tokens.len() && e.hash == prefix_hash(&tokens[..rows])
+        });
+        if !ok {
+            tracing::info!(
+                "MTP anchor: no drafter rows stored for snapshot {snap_slot} @ {rows} — \
+                 declining the warm restore (full recompute keeps MTP output equal to cold)"
+            );
+        }
+        ok
+    }
+
+    /// Restore the rows into the capture buffer and claim the capture for
+    /// `seq`, so the replay passes append from `rows` and the end-of-prefill
+    /// consume runs the cold whole-prompt `prefill_drafter`.
+    pub(in crate::model) fn mtp_anchor_restore(
+        &self,
+        seq: &mut SequenceState,
+        snap_slot: usize,
+        stream: u64,
+    ) -> anyhow::Result<()> {
+        if !self.mtp_anchor_on() {
+            return Ok(());
+        }
+        let st = store().lock();
+        let Some(e) = st.map.get(&snap_slot) else {
+            anyhow::bail!("MTP anchor vanished between check and restore");
+        };
+        self.gpu
+            .copy_h2d_async(&e.bytes, self.mtp_prefill_hidden, stream)?;
+        // The H2D reads host memory owned by the store; finish it before the
+        // lock is released so a concurrent eviction cannot free it mid-copy.
+        self.gpu.synchronize(stream)?;
+        let generation = self.mtp_prefill_capture_gen.fetch_add(1, Ordering::Relaxed) + 1;
+        seq.mtp_capture_gen = generation;
+        self.mtp_prefill_capture_len
+            .store(e.rows, Ordering::Relaxed);
+        let _ = DevicePtr::NULL;
+        Ok(())
+    }
+}

crates/spark-model/src/model/trait_impl/meta.rs

@@ -28,6 +28,9 @@ use crate::speculative::DraftProposer;
 use crate::traits::{ChunkedPrefillPageMetadata, Model, SequenceState};
 use crate::weight_map::{DenseWeight, MtpWeights, QuantizedWeight};
 
+mod grid;
+mod kv_cap;
+
 impl TransformerModel {
     pub(super) fn vocab_size_dispatch(&self) -> usize {
         self.config.vocab_size

crates/spark-model/src/model/trait_impl/meta/grid.rs

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+//! R11 prefill pass-grid queries (split from `meta.rs` for the 500-LoC cap).
+
+use super::super::super::types::TransformerModel;
+
+impl TransformerModel {
+    /// R11 fix 1: active prefill pass grid in tokens for this model, 0 = off.
+    /// Only hybrid-SSM models (whose warm path replays from an SSM anchor)
+    /// use it. Off unless `--exact-prefix-cache` or `ATLAS_PREFILL_GRID`.
+    pub(in super::super) fn prefill_grid(&self) -> usize {
+        if self.config.num_ssm_layers() == 0 {
+            return 0;
+        }
+        let g = spark_runtime::prefill_grid::prefill_grid();
+        if g > 0 && !spark_runtime::prefill_grid::subblock_disabled() {
+            spark_runtime::prefill_grid::disable_subblock();
+            tracing::info!("exact prefix cache: prefill pass grid {g} tokens");
+        }
+        g
+    }
+
+    /// R11 fix 6: under the prefill grid a warm replay runs exactly the cold
+    /// passes, so it must use the same GDN kernels (FLA) rather than the forced
+    /// "exact replay" path (register-resident token-sequential by default, WY4
+    /// only with `ATLAS_NO_GDN_REGRESIDENT=1`); both differ from the cold FLA
+    /// pass. `ATLAS_GRID_EXACT_REPLAY=1` restores the old forcing.
+    pub(in super::super) fn grid_replay_is_cold(&self) -> bool {
+        static FORCE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+        self.prefill_grid() > 0
+            && !*FORCE
+                .get_or_init(|| std::env::var("ATLAS_GRID_EXACT_REPLAY").as_deref() == Ok("1"))
+    }
+
+    /// R11 fix 1 (`cache_sequence`): under the pass grid insert only the KV of
+    /// the prompt's full grid passes. Returns false (nothing done) when the
+    /// grid is off.
+    pub(in super::super) fn cache_grid_prefix(
+        &self,
+        seq: &crate::traits::SequenceState,
+        bs: usize,
+    ) -> bool {
+        if self.prefill_grid() == 0 {
+            return false;
+        }
+        let n = spark_runtime::prefill_grid::cuts()
+            .floor(seq.prompt_len)
+            .min(seq.tokens.len());
+        let nb = (n / bs).min(seq.block_table.len());
+        if nb > 0 {
+            let acquired = self.prefix_cache.insert(
+                &seq.tokens[..nb * bs],
+                &seq.block_table[..nb],
+                &seq.disk_block_ids[..nb.min(seq.disk_block_ids.len())],
+                bs,
+                nb * bs,
+                seq.adapter_id,
+            );
+            super::super::super::block_mgmt::cache_acquires_refs(
+                &acquired,
+                &mut self.kv_cache.lock(),
+            );
+        }
+        true
+    }
+
+    /// R11 fix 1 (`prefill_chunk_dispatch_with`): position-fixed pass grid
+    /// (`spark_runtime::prefill_grid`). Every pass ends at a grid cut or the
+    /// prompt end, with prefix caching on or off, so a warm replay from a grid
+    /// anchor runs the same passes as a cold run. A chunk that crosses a cut
+    /// runs as two passes; `None` when the grid is off, the prompt has vision
+    /// pads, or no cut falls inside the chunk.
+    pub(in super::super) fn grid_cut(
+        &self,
+        tokens: &[u32],
+        seq: &mut crate::traits::SequenceState,
+        chunk_start: usize,
+        chunk_len: usize,
+        is_last_chunk: bool,
+        stream: u64,
+    ) -> Option<anyhow::Result<spark_runtime::gpu::DevicePtr>> {
+        if self.prefill_grid() == 0 || self.tokens_have_vision_pad(tokens) {
+            return None;
+        }
+        let end = chunk_start + chunk_len;
+        let cut = spark_runtime::prefill_grid::cuts().cut(chunk_start, end)?;
+        Some(
+            self.prefill_chunk_dispatch(tokens, seq, chunk_start, cut - chunk_start, false, stream)
+                .and_then(|_| {
+                    self.prefill_chunk_dispatch(tokens, seq, cut, end - cut, is_last_chunk, stream)
+                }),
+        )
+    }
+
+    /// The tail-split checkpoint (`prefill_chunk_dispatch_with`) runs only
+    /// with snapshots on and the pass grid off (the grid replaces it).
+    pub(in super::super) fn tail_split_snapshots(&self) -> bool {
+        self.prefill_grid() == 0 && self.ssm_snapshots.is_enabled()
+    }
+
+    /// R11 fix 1 (`finalize_last`): under the pass grid only KV written by
+    /// FULL grid passes is shared (the last pass has a prompt-dependent row
+    /// count) and no full-length leaf snapshot is saved (a warm replay must
+    /// start on a grid point); routing through the stale-V cap branch inserts
+    /// the grid-floored prefix without a snapshot. Returns
+    /// `(cache_blocks, cap_applied, grid_cap)`.
+    pub(in super::super) fn grid_cap_blocks(
+        &self,
+        cache_blocks: usize,
+        full_blocks: usize,
+        total: usize,
+        bs: usize,
+    ) -> (usize, bool, bool) {
+        if self.prefill_grid() == 0 {
+            return (cache_blocks, cache_blocks < full_blocks, false);
+        }
+        let grid_blocks = spark_runtime::prefill_grid::cuts().floor(total) / bs;
+        (cache_blocks.min(grid_blocks), true, true)
+    }
+}