metrale_model_engine/traits/model/
verify.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ModelVerify`, 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 super::VerifyBatchedOpts;
10use crate::traits::SequenceState;
11use anyhow::{Result, bail};
12
13/// 2026-09-26: Speculative verify passes: the fixed-width, γ-wide, DFlash and batched verifies, and the
14/// GDN fold after a write-on-accept batch.
15pub trait ModelVerify {
16    /// 2026-09-25: Run `tokens` through the model as one verify pass and return the argmax id at
17    /// each position; every token advances the KV and SSM state.
18    fn decode_verify(
19        &self,
20        tokens: &[u32],
21        seq: &mut SequenceState,
22        stream: u64,
23    ) -> Result<Vec<u32>>;
24
25    /// 2026-09-25: K=2 verify (`[last, draft]`); returns the argmax id at each position.
26    fn decode_verify_graphed(
27        &self,
28        tokens: &[u32; 2],
29        seq: &mut SequenceState,
30        stream: u64,
31    ) -> Result<[u32; 2]>;
32
33    /// 2026-09-25: K=3 verify (1 verified token and 2 drafts); returns 3 argmax ids.
34    fn decode_verify_graphed_k3(
35        &self,
36        tokens: &[u32; 3],
37        seq: &mut SequenceState,
38        stream: u64,
39    ) -> Result<[u32; 3]>;
40
41    /// 2026-09-25: K=4 verify (1 verified token and 3 drafts); returns 4 argmax ids.
42    fn decode_verify_graphed_k4(
43        &self,
44        tokens: &[u32; 4],
45        seq: &mut SequenceState,
46        stream: u64,
47    ) -> Result<[u32; 4]>;
48
49    /// 2026-09-25: Whether [`Self::decode_verify_batched`] can run `ks.len()` sequences with
50    /// `ks[i]` verify rows each (the sequence's draft count plus one; the rows may differ per
51    /// sequence). Default `false`, and there is no default batched implementation: a loop of
52    /// per-sequence verifies would leave only the last sequence's rows in the shared logits
53    /// buffer.
54    fn can_batch_verify(&self, _ks: &[usize]) -> bool {
55        false
56    }
57
58    /// 2026-09-25: Verify `ks.len()` sequences in one forward. Rows are flat and sequence-major
59    /// (`tokens.len() == Σ ks`): sequence `i` holds `[last_verified, d0, ..]` in rows
60    /// `off_i..off_i + ks[i]`, with `off_i = Σ_{t<i} ks[t]`. Returns the `Σ ks` argmax ids in the
61    /// same order. On success each sequence advances by its own `ks[i]`, and the caller rewinds
62    /// by its verdict.
63    ///
64    /// Call only when [`Self::can_batch_verify`] is true. With `opts.write_on_accept` the caller
65    /// must run [`Self::gdn_fold_accepted`] with every verdict before any
66    /// `commit_accepted_prefix`. Default: an error.
67    fn decode_verify_batched(
68        &self,
69        tokens: &[u32],
70        ks: &[usize],
71        seqs: &mut [&mut SequenceState],
72        stream: u64,
73        opts: VerifyBatchedOpts,
74    ) -> Result<Vec<u32>> {
75        let _ = (tokens, ks, seqs, stream, opts);
76        bail!("decode_verify_batched: unsupported by this model")
77    }
78
79    /// 2026-09-25: DFlash verify of `tokens` (1 verified token and γ drafts). Default:
80    /// [`Self::decode_verify`].
81    fn decode_verify_graphed_kgamma(
82        &self,
83        tokens: &[u32],
84        seq: &mut SequenceState,
85        stream: u64,
86    ) -> Result<Vec<u32>> {
87        self.decode_verify(tokens, seq, stream)
88    }
89
90    /// 2026-09-25: DFlash verify. Default: [`Self::decode_verify_graphed_kgamma`].
91    fn decode_verify_dflash(
92        &self,
93        tokens: &[u32],
94        seq: &mut SequenceState,
95        stream: u64,
96    ) -> Result<Vec<u32>> {
97        self.decode_verify_graphed_kgamma(tokens, seq, stream)
98    }
99
100    /// 2026-09-25: DFlash decode and verify in one forward: `tokens[0]` is the accepted token and
101    /// `tokens[1..]` the draft block. Default: [`Self::decode_verify_graphed_kgamma`].
102    fn decode_and_verify_fused(
103        &self,
104        tokens: &[u32],
105        seq: &mut SequenceState,
106        stream: u64,
107    ) -> Result<Vec<u32>> {
108        self.decode_verify_graphed_kgamma(tokens, seq, stream)
109    }
110
111    /// 2026-09-25: After a batched verify run with `write_on_accept`, commit the GDN `h` state of
112    /// each slot `slots[i]` at `accepted_rows[i]` verify rows (anchor included), in the verify's
113    /// batch order. `Ok(true)` means the `h` states were committed here, and
114    /// `commit_accepted_prefix` then restores only conv state for those slots. `Ok(false)`
115    /// means nothing was folded. Default: `Ok(false)`.
116    fn gdn_fold_accepted(
117        &self,
118        _slots: &[usize],
119        _accepted_rows: &[u32],
120        _k_rows: usize,
121    ) -> Result<bool> {
122        Ok(false)
123    }
124}
metrale_model_engine/traits/model/verify.rs

metrale_model_engine/traits/model/
verify.rs
