metrale_model_layers/speculative/
draft_stop.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: Per-position confidence stop for MTP draft chains
4//! (`--draft-confidence-stop <tau>`): the drafter keeps extending a chain only
5//! while its last draft's top-1 probability is at least `tau`. The serve
6//! publishes `tau` once; the drafter (single and batched propose) and the
7//! scheduler's depth planner read it here and apply the one rule
8//! [`chain_depth`].
9//!
10//! Owner: model-layers (speculative).
11//! Invariants:
12//! - A published `tau` is in `(0, 1)`; a second publication of a different
13//!   value is refused.
14//! - [`chain_depth`] never returns more than the drafts it is given, and
15//!   returns at least 1 for a non-empty chain: the first draft is always
16//!   verified, as without the stop.
17
18static PUBLISHED_TAU: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
19
20/// 2026-10-02: Publish the serve's stop threshold, before the model builds and
21/// the scheduler starts.
22pub fn set_draft_confidence_stop(tau: f32) -> anyhow::Result<()> {
23    anyhow::ensure!(
24        tau > 0.0 && tau < 1.0,
25        "--draft-confidence-stop must be in (0, 1), got {tau}"
26    );
27    let got = *PUBLISHED_TAU.get_or_init(|| tau);
28    anyhow::ensure!(
29        got.to_bits() == tau.to_bits(),
30        "draft confidence stop already published as {got}, refusing {tau}"
31    );
32    Ok(())
33}
34
35/// 2026-10-02: The published threshold as a log-probability, `ln tau`; `None`
36/// when the serve did not ask for the stop (and in unit tests and tools).
37pub fn draft_stop_logprob() -> Option<f32> {
38    PUBLISHED_TAU.get().map(|t| t.ln())
39}
40
41/// 2026-10-02: How many leading drafts of a chain to verify, from each draft's
42/// top-1 log-probability `lps[j]` (in draft order): drafting continues past
43/// draft `j` only while `lps[j] >= ln_tau`, so the chain ends with the first
44/// draft below the threshold, which is still verified. An empty chain is 0.
45pub fn chain_depth(lps: &[f32], ln_tau: f32) -> usize {
46    lps.iter()
47        .position(|&lp| lp < ln_tau)
48        .map_or(lps.len(), |j| j + 1)
49}
50
51/// 2026-10-02: Whether a chain whose last draft scored `lp` continues.
52pub fn chain_continues(lp: f32, ln_tau: f32) -> bool {
53    lp >= ln_tau
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    const LN_HALF: f32 = -std::f32::consts::LN_2;
61
62    #[test]
63    fn a_confident_chain_keeps_every_draft() {
64        assert_eq!(chain_depth(&[-0.1, -0.2, -0.3], LN_HALF), 3);
65    }
66
67    #[test]
68    fn the_chain_ends_with_its_first_unconfident_draft() {
69        assert_eq!(chain_depth(&[-0.1, -2.0, -0.1], LN_HALF), 2);
70        assert_eq!(chain_depth(&[-2.0, -0.1, -0.1], LN_HALF), 1);
71    }
72
73    #[test]
74    fn the_threshold_itself_continues() {
75        assert!(chain_continues(LN_HALF, LN_HALF));
76        assert!(!chain_continues(LN_HALF - 1e-6, LN_HALF));
77        assert_eq!(chain_depth(&[LN_HALF, LN_HALF], LN_HALF), 2);
78    }
79
80    #[test]
81    fn an_empty_chain_is_zero_and_a_lone_draft_is_one() {
82        assert_eq!(chain_depth(&[], LN_HALF), 0);
83        assert_eq!(chain_depth(&[-9.0], LN_HALF), 1);
84    }
85
86    #[test]
87    fn depth_agrees_with_continuation_at_every_position() {
88        // 2026-10-02: The single-sequence drafter stops with `chain_continues`,
89        // the batched planner truncates with `chain_depth`: the same rule.
90        let lps = [-0.05, -0.4, -0.3, -1.5, -0.01];
91        let ln_tau = (0.6f32).ln();
92        let mut drafted = 0;
93        for &lp in &lps {
94            drafted += 1;
95            if !chain_continues(lp, ln_tau) {
96                break;
97            }
98        }
99        assert_eq!(drafted, chain_depth(&lps, ln_tau));
100    }
101
102    #[test]
103    fn publication_refuses_out_of_range_values() {
104        for bad in [0.0, 1.0, -0.5, 1.5, f32::NAN] {
105            assert!(set_draft_confidence_stop(bad).is_err(), "{bad} accepted");
106        }
107    }
108}
metrale_model_layers/speculative/draft_stop.rs

metrale_model_layers/speculative/
draft_stop.rs
