metrale_model_engine/traits/model/
ssm_state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ModelSsmState`, one of the supertraits `Model` is made of. Its methods, default
4//! bodies and docs are the ones `Model` declared before the split.
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9use crate::traits::SequenceState;
10use anyhow::Result;
11
12/// 2026-09-26: Recurrent (SSM) state: checkpoint and rollback around a verify, decode-rollback
13/// snapshots, and the commit of an accepted prefix on the secondary stream.
14pub trait ModelSsmState {
15    fn checkpoint_ssm_states(&self, seq: &mut SequenceState) -> Result<()>;
16
17    /// 2026-09-25: Restore the SSM state to where `num_accepted` tokens of the last verify were
18    /// accepted.
19    fn rollback_ssm_states(&self, seq: &mut SequenceState, num_accepted: usize) -> Result<()>;
20
21    /// 2026-09-25: Whether the model has recurrent SSM layers, whose state is advanced in place
22    /// each token. Lowering `seq_len` rewinds paged KV but not that state, so a decode rollback
23    /// must also restore it from a snapshot. Default `false`.
24    fn has_ssm_layers(&self) -> bool {
25        false
26    }
27
28    /// 2026-09-26: Whether some layer keeps per-sequence state that lowering the KV cursor does
29    /// not rewind (`LayerCapabilities::decode_rollback_unsupported`); the scheduler then declines
30    /// a decode rollback (`rollback_to_boundary`). Default `false`.
31    fn decode_rollback_unsupported(&self) -> bool {
32        false
33    }
34
35    /// 2026-09-25: Decode-rollback SSM snapshot slots per active sequence; the scheduler sizes
36    /// each sequence's snapshot ring from it. `0` (the default) means no decode-rollback
37    /// snapshots. `TransformerModel` returns the depth `ssm_reserve::decode_rollback_ring_slots`
38    /// chose.
39    fn decode_rollback_ring_slots(&self) -> usize {
40        0
41    }
42
43    /// 2026-09-25: Save `seq`'s SSM `h_state` and `conv_state` (every SSM layer) into its
44    /// decode-rollback ring slot `ring_slot`, in `0..decode_rollback_ring_slots()`. Default:
45    /// `Ok(())`.
46    fn save_decode_ssm_snapshot(&self, _seq: &SequenceState, _ring_slot: usize) -> Result<()> {
47        Ok(())
48    }
49
50    /// 2026-09-25: Restore `seq`'s SSM state from ring slot `ring_slot`, written earlier by
51    /// [`Self::save_decode_ssm_snapshot`]. Default: `Ok(())`.
52    fn restore_decode_ssm_snapshot(&self, _seq: &SequenceState, _ring_slot: usize) -> Result<()> {
53        Ok(())
54    }
55
56    /// 2026-09-25: During decode, save a block-aligned SSM snapshot at a checkpoint boundary, so
57    /// the next turn's prefix-cache hit can restore from decode-produced state. The scheduler
58    /// calls it after a step's SSM state is final. Default: no-op.
59    fn decode_marconi_checkpoint(&self, _seq: &mut SequenceState) {}
60
61    /// 2026-09-25: Start the SSM checkpoint copies on the secondary stream; call
62    /// [`Self::sync_secondary`] before the next verify. Default: the synchronous
63    /// [`Self::checkpoint_ssm_states`].
64    fn start_checkpoint_async(&self, seq: &mut SequenceState) -> Result<()> {
65        self.checkpoint_ssm_states(seq)
66    }
67
68    /// 2026-09-25: Roll the SSM state back to `num_accepted` and checkpoint it again, on the
69    /// secondary stream. Default: the synchronous rollback, then checkpoint.
70    fn start_rollback_and_checkpoint_async(
71        &self,
72        seq: &mut SequenceState,
73        num_accepted: usize,
74    ) -> Result<()> {
75        self.rollback_ssm_states(seq, num_accepted)?;
76        self.checkpoint_ssm_states(seq)
77    }
78
79    /// 2026-09-25: Wait for the secondary stream's work. Default: `Ok(())`.
80    fn sync_secondary(&self) -> Result<()> {
81        Ok(())
82    }
83
84    /// 2026-09-25: Commit the accepted prefix of a verify onto the live SSM `h_state` and
85    /// `conv_state`, on the secondary stream (pair with [`Self::sync_secondary`]). Default:
86    /// `Ok(())`.
87    fn commit_accepted_prefix(
88        &self,
89        _seq: &mut SequenceState,
90        _num_accepted: usize,
91        _k: usize,
92    ) -> Result<()> {
93        Ok(())
94    }
95
96    /// 2026-09-25: SSM snapshot-pool occupancy `(used, total)`, or `None` without a snapshot
97    /// pool (the default; `TransformerModel` also returns `None` when the pool has no slots).
98    fn ssm_snapshot_occupancy(&self) -> Option<(u32, u32)> {
99        None
100    }
101}
metrale_model_engine/traits/model/ssm_state.rs

metrale_model_engine/traits/model/
ssm_state.rs
