metrale_model_layers/ssm_reserve/
decode_ring.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Depth of the per-sequence SSM decode-rollback ring: the depth
4//! decision, the published depth, and the fit to free memory.
5//!
6//! Owner: model-layers (SSM reserve).
7//! Invariants:
8//! - The published depth is set at most once per process; later writes leave
9//!   the first value in force.
10//! - `parse_decode_ring_slots` accepts only `auto` and
11//!   `0..=DECODE_ROLLBACK_RING_SLOTS`.
12//!
13//! The server's preflight reserve (`preflight/decode_ring.rs`, unless
14//! `METRALE_SSM_RESERVE_RING_FULL` is present) and `TransformerModel::new`
15//! (`impl_a1.rs`) both take the depth from [`decode_rollback_ring_slots`].
16//! When preflight fits a smaller depth to free memory, it publishes that depth
17//! through [`set_decode_ring_slots`] so the allocation follows it. The ring is
18//! written by the scheduler's `rollback::snapshot_boundary_if_ssm` and read by
19//! `rollback::rollback_to_boundary`; on a model with SSM layers and a depth of
20//! 0, a rollback declines with `RollbackFallback::NoSsmSnapshot`.
21//!
22//! Environment, read in [`decode_rollback_ring_slots`]:
23//! - `METRALE_SSM_DECODE_RING`: `1` forces the full depth and `0` forces 0,
24//!   when no depth is published.
25//! - `METRALE_DISABLE_WATCHDOGS`: `1` or `true` skips the ring when no depth
26//!   is published and `METRALE_SSM_DECODE_RING` is not `1`.
27
28/// 2026-09-25: Outcome of the ring-depth decision. `skip_reason` is `Some` only for
29/// the implicit skips (speculative decode, watchdogs disabled), so the
30/// allocating call site can log the saving.
31pub struct DecodeRingDecision {
32    pub slots: usize,
33    pub skip_reason: Option<&'static str>,
34}
35
36/// 2026-09-25: Depths the preflight fit may choose, largest first. The first equals
37/// [`metrale_kernels::DECODE_ROLLBACK_RING_SLOTS`] (8); each next depth halves
38/// the ring's share of the reserve.
39pub const DECODE_RING_FIT_LADDER: [usize; 5] = [8, 4, 2, 1, 0];
40
41/// 2026-09-25: The published depth: an explicit `--ssm-decode-ring-slots N`
42/// (`serve_flags.rs`, before preflight) or the depth preflight fitted. `auto`
43/// publishes nothing, so `METRALE_SSM_DECODE_RING` applies unless preflight
44/// publishes a fitted depth.
45static DECODE_RING_SLOTS: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
46
47/// 2026-09-25: Publishes the ring depth and returns the depth in force; the first
48/// write wins. Preflight skips its fit when a depth is already published
49/// ([`published_decode_ring_slots`]), so an explicit depth is not shrunk.
50pub fn set_decode_ring_slots(slots: usize) -> usize {
51    let _ = DECODE_RING_SLOTS.set(slots);
52    *DECODE_RING_SLOTS.get().expect("just set")
53}
54
55/// 2026-09-25: The published depth, or `None`. Reads the cell without setting it.
56pub fn published_decode_ring_slots() -> Option<usize> {
57    DECODE_RING_SLOTS.get().copied()
58}
59
60/// 2026-09-25: Parses `--ssm-decode-ring-slots`: `auto` gives `None` (fit at
61/// preflight), a number up to `DECODE_ROLLBACK_RING_SLOTS` gives `Some`, and
62/// anything else an error. `validate_serve_args` and `publish_kernel_flags`
63/// both parse through this.
64pub fn parse_decode_ring_slots(s: &str) -> Result<Option<usize>, String> {
65    if s == "auto" {
66        return Ok(None);
67    }
68    let n: usize = s
69        .parse()
70        .map_err(|_| format!("unknown ssm-decode-ring-slots '{s}' (valid: auto, 0..=8)"))?;
71    if n > metrale_kernels::DECODE_ROLLBACK_RING_SLOTS {
72        return Err(format!(
73            "ssm-decode-ring-slots {n} exceeds the {} the ring is sized for",
74            metrale_kernels::DECODE_ROLLBACK_RING_SLOTS
75        ));
76    }
77    Ok(Some(n))
78}
79
80/// 2026-09-25: The per-sequence ring depth, from the published depth and the
81/// environment ([`decode_rollback_ring_slots_with`]). Preflight passes
82/// `args.speculative || args.dflash` as `use_speculative`; the allocation must
83/// receive the same flag, or the reserve and the ring differ.
84pub fn decode_rollback_ring_slots(
85    num_ssm_layers: usize,
86    use_speculative: bool,
87) -> DecodeRingDecision {
88    let watchdogs_value = std::env::var("METRALE_DISABLE_WATCHDOGS").ok();
89    let watchdogs_disabled = watchdogs_disabled_from_value(watchdogs_value.as_deref());
90    let ring_override = std::env::var("METRALE_SSM_DECODE_RING").ok();
91    decode_rollback_ring_slots_with(
92        num_ssm_layers,
93        use_speculative,
94        published_decode_ring_slots(),
95        ring_override.as_deref(),
96        watchdogs_disabled,
97    )
98}
99
100/// 2026-09-25: `METRALE_DISABLE_WATCHDOGS` truthiness: `1` or `true`, trimmed and
101/// case-insensitive, as the server's `parse_disable_watchdogs` reads it.
102pub fn watchdogs_disabled_from_value(value: Option<&str>) -> bool {
103    value
104        .map(|v| {
105            let v = v.trim().to_ascii_lowercase();
106            v == "1" || v == "true"
107        })
108        .unwrap_or(false)
109}
110
111/// 2026-09-25: Env-free core of [`decode_rollback_ring_slots`]. Precedence, highest
112/// first:
113///
114/// 1. no SSM layers: 0;
115/// 2. `published`: that depth, also under speculative decode or with watchdogs
116///    disabled;
117/// 3. `ring_override` `"1"` or `"0"`: `DECODE_ROLLBACK_RING_SLOTS` or 0;
118/// 4. speculative decode or watchdogs disabled: 0, with a `skip_reason`;
119/// 5. otherwise `DECODE_ROLLBACK_RING_SLOTS`.
120pub fn decode_rollback_ring_slots_with(
121    num_ssm_layers: usize,
122    use_speculative: bool,
123    published: Option<usize>,
124    ring_override: Option<&str>,
125    watchdogs_disabled: bool,
126) -> DecodeRingDecision {
127    if num_ssm_layers == 0 {
128        return DecodeRingDecision {
129            slots: 0,
130            skip_reason: None,
131        };
132    }
133    if let Some(slots) = published {
134        return DecodeRingDecision {
135            slots,
136            skip_reason: None,
137        };
138    }
139    match ring_override {
140        Some("1") => DecodeRingDecision {
141            slots: metrale_kernels::DECODE_ROLLBACK_RING_SLOTS,
142            skip_reason: None,
143        },
144        Some("0") => DecodeRingDecision {
145            slots: 0,
146            skip_reason: None,
147        },
148        _ if use_speculative || watchdogs_disabled => DecodeRingDecision {
149            slots: 0,
150            skip_reason: Some(if use_speculative {
151                "speculative decode active"
152            } else {
153                "watchdogs disabled"
154            }),
155        },
156        _ => DecodeRingDecision {
157            slots: metrale_kernels::DECODE_ROLLBACK_RING_SLOTS,
158            skip_reason: None,
159        },
160    }
161}
162
163/// 2026-09-25: Largest [`DECODE_RING_FIT_LADDER`] depth, at most `start_slots`, with
164/// `reserve_without_ring + depth * bytes_per_ring_slot <= free_mem` (saturating
165/// arithmetic). `bytes_per_ring_slot` is one depth unit of the ring. Returns 0
166/// also when no depth fits; the caller tells the two apart by comparing the
167/// total against its limit (`preflight/decode_ring.rs`).
168pub fn fit_decode_ring_slots(
169    start_slots: usize,
170    reserve_without_ring: usize,
171    bytes_per_ring_slot: usize,
172    free_mem: usize,
173) -> usize {
174    DECODE_RING_FIT_LADDER
175        .iter()
176        .copied()
177        .filter(|&slots| slots <= start_slots)
178        .find(|&slots| {
179            reserve_without_ring.saturating_add(slots.saturating_mul(bytes_per_ring_slot))
180                <= free_mem
181        })
182        .unwrap_or(0)
183}
184
185#[cfg(test)]
186#[path = "decode_ring_tests.rs"]
187mod decode_ring_tests;
metrale_model_layers/ssm_reserve/decode_ring.rs

metrale_model_layers/ssm_reserve/
decode_ring.rs
