metrale_model_engine/model/trait_impl/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `Model` and its supertraits for `TransformerModel`: one `impl_<supertrait>.rs`
4//! file per supertrait, mostly delegating to `<method>_dispatch` helpers in the sibling modules.
5//!
6//! Owner: model-engine.
7//! Invariants:
8//! - `prefill`, `prefill_chunk`, `prefill_twophase` and `mixed_forward` call
9//!   `try_eager_drafter_prefill` only after their forward returns `Ok`.
10//! - When the forward returns `Err`, `decode_batch`, `mixed_forward` and every
11//!   verify method that goes through `release_verify_capture_on_err` end any
12//!   capture still open on the default stream.
13
14#![allow(unused_imports, dead_code, clippy::too_many_arguments)]
15
16use anyhow::Result;
17use metrale_cache::kv_cache::PagedKvCache;
18use metrale_gpu_runtime::gpu::DevicePtr;
19
20use super::types::{PinnedMetaStaging, TransformerModel};
21use crate::traits::{
22    ChunkedPrefillPageMetadata, FeedSource, Model, PrefillSlice, RowMask, SequenceState,
23};
24use metrale_model_layers::layer::{AttnMetadataDev, LayerState};
25use metrale_model_layers::speculative::DraftProposer;
26use metrale_model_layers::weight_map::{DenseWeight, MtpWeights};
27
28mod async_chkpt;
29mod borrow_streak;
30mod decode_a;
31mod decode_a2;
32mod decode_a3;
33mod decode_a_diag;
34mod decode_b;
35mod decode_b2;
36pub(in crate::model) mod decode_checkpoint;
37mod decode_graph_key;
38mod slot_graph_key;
39pub(crate) use slot_graph_key::SlotGraphKey;
40mod decode_multi_seq_gate;
41mod drafter_prefill;
42mod ep_misc;
43mod feed;
44mod gdn_woa;
45mod graph_borrow;
46mod impl_adapters;
47mod impl_circuit;
48mod impl_circuit_run;
49mod impl_circuit_verify_batch;
50mod impl_device_feed;
51mod impl_draft;
52mod impl_ep;
53mod impl_forward;
54mod impl_lifecycle;
55mod impl_logits;
56mod impl_ssm_state;
57mod impl_streams;
58mod impl_verify;
59mod impl_vision;
60mod lm_head_batched;
61mod mamba2_state_guard;
62mod meta;
63mod meta_argmax;
64mod prefill_a;
65mod prefill_b;
66mod prefill_c;
67mod prefill_d;
68mod prefix_reuse;
69mod sequence;
70mod sequence_compact;
71mod speculative;
72mod speculative_mtp;
73pub(in crate::model) mod ssm_fault_in;
74mod verify_a;
75mod verify_a_ssm;
76mod verify_b;
77mod verify_c;
78mod verify_c2;
79mod verify_d;
80mod verify_e;
81pub(in crate::model) mod verify_e2;
82mod verify_fused;
83
84impl Model for TransformerModel {}
85
86impl TransformerModel {
87    /// 2026-09-25: On `Err`, end and discard any capture still open on the
88    /// default stream, as `decode_batch` does, so a verify that fails during
89    /// capture does not leave the stream capturing. A no-op when the stream is
90    /// not capturing. The verify modules that capture (verify_b, verify_c,
91    /// verify_c2, verify_d, verify_e, verify_fused) capture on the default
92    /// stream.
93    fn release_verify_capture_on_err<T>(&self, r: Result<T>) -> Result<T> {
94        if r.is_err() {
95            self.gpu.abort_capture_if_active(self.gpu.default_stream());
96        }
97        r
98    }
99
100    /// 2026-09-26: Collect each layer's aux sequence state
101    /// (`LayerAuxState::snapshot_aux`) as `(layer index, blob)` pairs to
102    /// store with an SSM snapshot. Empty when no layer returns a blob.
103    pub(in crate::model) fn collect_aux_states(
104        &self,
105        seq: &SequenceState,
106        stream: u64,
107    ) -> Result<Vec<(u32, Vec<u8>)>> {
108        let mut out = Vec::new();
109        for (i, l) in self.layers.iter().enumerate() {
110            if let Some(blob) =
111                l.snapshot_aux(seq.layer_states[i].as_ref(), self.gpu.as_ref(), stream)?
112            {
113                out.push((i as u32, blob));
114            }
115        }
116        Ok(out)
117    }
118
119    /// 2026-09-25: True when some layer reports `has_aux_state`; restore sites
120    /// then decline a snapshot that has no aux blobs.
121    pub(in crate::model) fn requires_aux_state(&self) -> bool {
122        self.layers.iter().any(|l| l.has_aux_state())
123    }
124
125    /// 2026-09-25: Apply a snapshot's aux blobs to the layers they came from.
126    /// The first layer error is returned.
127    pub(in crate::model) fn apply_aux_states(
128        &self,
129        seq: &mut SequenceState,
130        blobs: &[(u32, Vec<u8>)],
131        stream: u64,
132    ) -> Result<()> {
133        for (i, blob) in blobs {
134            self.layers[*i as usize].restore_aux(
135                seq.layer_states[*i as usize].as_mut(),
136                blob,
137                self.gpu.as_ref(),
138                stream,
139            )?;
140        }
141        Ok(())
142    }
143}
144
145#[cfg(test)]
146mod gdn_woa_folded_slots_tests {
147    /// 2026-09-25: The folded-slot record must hold only the slots folded by
148    /// the most recent fold.
149    ///
150    /// The logic is modelled directly because the real path needs a GPU.
151    /// `record_new` is the recording in `gdn_fold_accepted_dispatch` (clear,
152    /// then extend only when a layer folded). `record_old` clears only when a
153    /// layer folded; it is the negative control.
154    fn record_old(list: &mut Vec<usize>, any: bool, slots: &[usize]) {
155        if any {
156            list.clear();
157            list.extend_from_slice(slots);
158        }
159    }
160    fn record_new(list: &mut Vec<usize>, any: bool, slots: &[usize]) {
161        list.clear();
162        if any {
163            list.extend_from_slice(slots);
164        }
165    }
166    /// 2026-09-25: The reader in async_chkpt.rs: removes `slot` and returns
167    /// true when it was listed, which skips that slot's `h` restore.
168    fn consume(list: &mut Vec<usize>, slot: usize) -> bool {
169        match list.iter().position(|&s| s == slot) {
170            Some(p) => {
171                list.swap_remove(p);
172                true
173            }
174            None => false,
175        }
176    }
177
178    /// 2026-09-25: A fold that folds nothing must not leave a previous batch's
179    /// slots listed. The reader is keyed by slot index, so a stale entry makes
180    /// the slot's next occupant skip its `h` restore.
181    #[test]
182    fn a_no_op_fold_does_not_leave_a_stale_claim() {
183        let mut new = vec![];
184        record_new(&mut new, true, &[1, 2, 3, 4]);
185        for s in [1, 2, 3, 4] {
186            assert!(consume(&mut new, s), "batch A's own slots are folded");
187        }
188        record_new(&mut new, false, &[7]);
189        assert!(
190            !consume(&mut new, 7),
191            "slot 7 was NOT folded and must restore h"
192        );
193
194        // 2026-09-25: Control: `record_old` fails this. The no-op fold leaves
195        // batch A's entry for slot 3 in the list.
196        let mut old = vec![];
197        record_old(&mut old, true, &[1, 2, 3, 4]);
198        record_old(&mut old, false, &[7]);
199        assert!(
200            consume(&mut old, 3),
201            "the shipped behaviour leaves slot 3 claimed — this assertion \
202             documents the bug, and its inverse is the fix above"
203        );
204    }
205}
metrale_model_engine/model/trait_impl/mod.rs

metrale_model_engine/model/trait_impl/
mod.rs
