metrale_model_layers/ssm_reserve/
rollback.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The process-wide SSM verify-rollback mode and the replay-ring
4//! sizing it implies.
5//!
6//! Owner: model-layers (SSM reserve).
7//! Invariants:
8//! - The published mode is set at most once per process; later
9//!   `set_ssm_rollback_mode` calls return the mode already in force.
10
11/// 2026-09-25: SSM verify-rollback mode (`--ssm-rollback-mode`, experimental).
12///
13/// * `Snapshot` (the CLI default): verify writes per-token h/conv state
14///   intermediates, and a partial accept restores from them. The only mode
15///   with a device path.
16/// * `Replay`: keeps only the pre-verify checkpoint per verify slot plus a
17///   ring sized for the verify window's per-token GDN inputs. Capture and
18///   replay are not implemented: a serve in this mode boots, and every
19///   speculative verify entry refuses
20///   (`SsmStatePool::require_verify_rollback_supported`).
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub enum SsmRollbackMode {
23    Snapshot,
24    Replay,
25}
26
27impl std::str::FromStr for SsmRollbackMode {
28    type Err = String;
29    /// 2026-09-25: The parse for the `--ssm-rollback-mode` value; CLI validation
30    /// and the serve's publication both call it.
31    fn from_str(s: &str) -> Result<Self, Self::Err> {
32        match s {
33            "snapshot" => Ok(Self::Snapshot),
34            "replay" => Ok(Self::Replay),
35            other => Err(format!(
36                "unknown ssm-rollback-mode '{other}' (valid: snapshot, replay)"
37            )),
38        }
39    }
40}
41
42/// 2026-09-25: The published rollback mode, written from the serve command line
43/// and read by pool construction and preflight. The first write wins, as in
44/// `gdn_flags::set_from_cli`.
45static ROLLBACK_MODE: std::sync::OnceLock<SsmRollbackMode> = std::sync::OnceLock::new();
46
47/// 2026-09-25: Publish the command line's mode. Returns the mode in force (first
48/// write wins).
49pub fn set_ssm_rollback_mode(mode: SsmRollbackMode) -> SsmRollbackMode {
50    let _ = ROLLBACK_MODE.set(mode);
51    *ROLLBACK_MODE.get().expect("just set")
52}
53
54/// 2026-09-25: The mode in force. `Snapshot` when nothing was published (tests,
55/// examples); reading it that way also fixes it for the process.
56pub fn ssm_rollback_mode() -> SsmRollbackMode {
57    *ROLLBACK_MODE.get_or_init(|| SsmRollbackMode::Snapshot)
58}
59
60/// 2026-09-25: Bytes of one cached verify row of GDN inputs, per SSM layer: the
61/// deinterleaved qkvz row (`qkvz_elems` BF16, the rows of
62/// `ConvGdnArgs::deinterleaved`) plus the gate/beta row (`nv * 2` FP32, the
63/// rows of `ConvGdnArgs::gates_buf`).
64pub fn ssm_replay_row_bytes(qkvz_elems: usize, nv: usize) -> usize {
65    qkvz_elems * 2 + nv * 2 * 4
66}
67
68/// 2026-09-25: Replay-mode verify-window input ring:
69/// `mtp_state_slots * (k_ceiling - 1) * num_ssm_layers * row_bytes`, since a
70/// partial accept replays at most K-1 tokens. Preflight reserves and
71/// `SsmStatePool::new` allocates through this function, with different
72/// arguments: preflight passes `num_drafts + 1` and `mtp_state_slots`, the
73/// pool its `num_intermediates` and `mtp_slots + 1` (the dummy slot).
74pub fn ssm_replay_ring_bytes(
75    num_ssm_layers: usize,
76    row_bytes: usize,
77    k_ceiling: usize,
78    mtp_state_slots: usize,
79) -> usize {
80    mtp_state_slots * k_ceiling.saturating_sub(1) * num_ssm_layers * row_bytes
81}
metrale_model_layers/ssm_reserve/rollback.rs

metrale_model_layers/ssm_reserve/
rollback.rs
