metrale_model_engine/model/trait_impl/
impl_ssm_state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `impl ModelSsmState for TransformerModel`, mostly delegating to `<method>_dispatch`
4//! helpers in the sibling modules.
5//!
6//! Owner: model-engine.
7//! Invariants: the ones in `trait_impl/mod.rs`.
8
9use anyhow::Result;
10
11use crate::model::types::TransformerModel;
12use crate::traits::{ModelSsmState, SequenceState};
13
14impl ModelSsmState for TransformerModel {
15    fn checkpoint_ssm_states(&self, seq: &mut SequenceState) -> Result<()> {
16        self.gdn_carry_flush_pending()?;
17        self.checkpoint_ssm_states_dispatch(seq)
18    }
19
20    fn rollback_ssm_states(&self, seq: &mut SequenceState, num_accepted: usize) -> Result<()> {
21        self.gdn_carry_flush_pending()?;
22        self.rollback_ssm_states_dispatch(seq, num_accepted)
23    }
24
25    fn has_ssm_layers(&self) -> bool {
26        self.ssm_pool.num_ssm_layers > 0
27    }
28
29    fn decode_rollback_unsupported(&self) -> bool {
30        self.layers.iter().any(|l| l.decode_rollback_unsupported())
31    }
32
33    fn decode_rollback_ring_slots(&self) -> usize {
34        if self.ssm_snapshots.decode_rollback_enabled() {
35            self.ssm_snapshots.decode_ring_slots
36        } else {
37            0
38        }
39    }
40
41    fn save_decode_ssm_snapshot(&self, seq: &SequenceState, ring_slot: usize) -> Result<()> {
42        self.gdn_carry_flush_pending()?;
43        self.save_decode_ssm_snapshot_dispatch(seq, ring_slot)
44    }
45
46    fn restore_decode_ssm_snapshot(&self, seq: &SequenceState, ring_slot: usize) -> Result<()> {
47        self.gdn_carry_flush_pending()?;
48        self.restore_decode_ssm_snapshot_dispatch(seq, ring_slot)
49    }
50
51    fn decode_marconi_checkpoint(&self, seq: &mut SequenceState) {
52        if let Err(e) = self.gdn_carry_flush_pending() {
53            tracing::error!(
54                "gdn carry flush before {}: {e:#}",
55                "decode_marconi_checkpoint"
56            );
57        }
58        self.decode_marconi_checkpoint_dispatch(seq)
59    }
60
61    fn ssm_snapshot_occupancy(&self) -> Option<(u32, u32)> {
62        let (used, total) = self.ssm_snapshots.occupancy();
63        (total > 0).then_some((used as u32, total as u32))
64    }
65
66    fn start_checkpoint_async(&self, seq: &mut SequenceState) -> Result<()> {
67        self.gdn_carry_flush_pending()?;
68        self.start_checkpoint_async_dispatch(seq)
69    }
70
71    fn start_rollback_and_checkpoint_async(
72        &self,
73        seq: &mut SequenceState,
74        num_accepted: usize,
75    ) -> Result<()> {
76        self.gdn_carry_flush_pending()?;
77        self.start_rollback_and_checkpoint_async_dispatch(seq, num_accepted)
78    }
79
80    fn sync_secondary(&self) -> Result<()> {
81        self.sync_secondary_dispatch()
82    }
83
84    fn commit_accepted_prefix(
85        &self,
86        seq: &mut SequenceState,
87        num_accepted: usize,
88        k: usize,
89    ) -> Result<()> {
90        self.commit_accepted_prefix_dispatch(seq, num_accepted, k)
91    }
92}
metrale_model_engine/model/trait_impl/impl_ssm_state.rs

metrale_model_engine/model/trait_impl/
impl_ssm_state.rs
