metrale_model_layers/speculative/
ladder.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Per-step MTP draft count as a function of the number of active
4//! sequences (the K-vs-batch ladder), and the multi-sequence MTP width cap.
5//!
6//! Owner: model-layers (speculative).
7//! Invariants:
8//! - `mtp_ladder_drafts` returns 0 when `num_drafts` is 0, and a value in
9//!   `1..=num_drafts` otherwise.
10//! - The ladder, its kill switch and the cap are read from the environment
11//!   once per process.
12//!
13//! The default ladder is `4:3,8:3,16:1,32:1`: three drafts up to 8 active
14//! sequences, one draft above that. A width above the last step takes the last
15//! step's count. `metrale_speculative::adaptive_rung` uses this count as its
16//! floor and may raise widths 9..=16 to two drafts; it keeps the static count
17//! when `METRALE_MTP_STATIC_RUNG` or `METRALE_MTP_K_LADDER` is present.
18//!
19//! Overrides:
20//! - `METRALE_MTP_K_LADDER`, e.g. `"4:3,8:2,16:1"`: comma-separated
21//!   `n_max:drafts` steps, sorted by `n_max`. A value that does not parse as a
22//!   whole gives the default ladder. Draft counts are clamped to
23//!   `1..=num_drafts`.
24//! - `METRALE_NO_MTP_K_LADDER`, present with any value (`0` included): every
25//!   width gets `num_drafts`, and the default of [`mtp_max_seqs`] drops from
26//!   32 to 4.
27
28/// 2026-09-25: Whether `METRALE_NO_MTP_K_LADDER` is present, with any value.
29/// Read once per process.
30pub fn mtp_ladder_disabled() -> bool {
31    static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
32    *OFF.get_or_init(|| std::env::var_os("METRALE_NO_MTP_K_LADDER").is_some())
33}
34
35/// 2026-09-25: Parses `n_max:drafts,...` into steps sorted by `n_max`. `None`
36/// when any step fails to parse; the caller then uses the default ladder.
37fn parse_ladder(value: &str) -> Option<Vec<(usize, usize)>> {
38    let mut steps = Vec::new();
39    for part in value.split(',') {
40        let (n, k) = part.trim().split_once(':')?;
41        steps.push((n.trim().parse().ok()?, k.trim().parse().ok()?));
42    }
43    if steps.is_empty() {
44        return None;
45    }
46    steps.sort_by_key(|&(n, _)| n);
47    Some(steps)
48}
49
50/// 2026-09-25: The default ladder. Measured 2026-07-31 on dgx1 at C=16, one binary,
51/// three reps per arm: 16:1 181.63 tok/s, 16:2 172.70, 16:3 152.97.
52const DEFAULT_LADDER: [(usize, usize); 4] = [(4, 3), (8, 3), (16, 1), (32, 1)];
53
54/// 2026-10-03: The ladder `--mtp-k-ladder` published ([`set_mtp_k_ladder`]).
55static PUBLISHED_LADDER: std::sync::OnceLock<Vec<(usize, usize)>> = std::sync::OnceLock::new();
56
57fn mtp_ladder_steps() -> &'static [(usize, usize)] {
58    if let Some(published) = PUBLISHED_LADDER.get() {
59        return published;
60    }
61    static STEPS: std::sync::OnceLock<Vec<(usize, usize)>> = std::sync::OnceLock::new();
62    STEPS.get_or_init(|| {
63        std::env::var("METRALE_MTP_K_LADDER")
64            .ok()
65            .and_then(|value| parse_ladder(&value))
66            .unwrap_or_else(|| DEFAULT_LADDER.to_vec())
67    })
68}
69
70/// 2026-10-03: Parses a `--mtp-k-ladder` value (`n_max:drafts,...`); an error names
71/// what does not parse, where the environment variable falls back to the default.
72pub fn parse_mtp_k_ladder(value: &str) -> anyhow::Result<Vec<(usize, usize)>> {
73    let steps = parse_ladder(value).ok_or_else(|| {
74        anyhow::anyhow!("--mtp-k-ladder {value:?} is not a list of n_max:drafts steps")
75    })?;
76    anyhow::ensure!(
77        steps.iter().all(|&(n, k)| n >= 1 && k >= 1),
78        "--mtp-k-ladder {value:?}: widths and draft counts start at 1"
79    );
80    Ok(steps)
81}
82
83/// 2026-10-03: Publish `--mtp-k-ladder` once, before the model sizes its verify pools
84/// (`ssm_reserve::verify_slot_drafts` reads the ladder) and the scheduler starts. It wins
85/// over `METRALE_MTP_K_LADDER`. A second publication of a different ladder is refused.
86pub fn set_mtp_k_ladder(steps: Vec<(usize, usize)>) -> anyhow::Result<()> {
87    let got = PUBLISHED_LADDER.get_or_init(|| steps.clone());
88    anyhow::ensure!(
89        *got == steps,
90        "MTP K-ladder already published as {got:?}, refusing {steps:?}"
91    );
92    Ok(())
93}
94
95/// 2026-10-03: Whether the serve runs an explicit ladder (`--mtp-k-ladder`, or
96/// `METRALE_MTP_K_LADDER` present with any value): the adaptive rung then keeps the
97/// static counts, so the operator gets exactly the rungs spelled out.
98pub fn mtp_ladder_pinned() -> bool {
99    PUBLISHED_LADDER.get().is_some() || std::env::var_os("METRALE_MTP_K_LADDER").is_some()
100}
101
102/// 2026-10-03: The draft count of `steps` at `n_active` sequences under the
103/// `num_drafts` ceiling: the first step whose width covers `n_active`, else the last.
104pub fn ladder_drafts_from_steps(
105    steps: &[(usize, usize)],
106    n_active: usize,
107    num_drafts: usize,
108) -> usize {
109    if num_drafts == 0 {
110        return 0;
111    }
112    steps
113        .iter()
114        .find(|&&(n_max, _)| n_active <= n_max)
115        .or(steps.last())
116        .map(|&(_, k)| k.clamp(1, num_drafts))
117        .unwrap_or(num_drafts)
118}
119
120/// 2026-09-25: The per-step draft count for `n_active` concurrent sequences.
121///
122/// `num_drafts` is the configured ceiling. The result is in `1..=num_drafts`,
123/// or 0 when `num_drafts` is 0. With the ladder disabled it is `num_drafts`.
124/// A width above the last step uses the last step's count.
125pub fn mtp_ladder_drafts(n_active: usize, num_drafts: usize) -> usize {
126    if num_drafts == 0 {
127        return 0;
128    }
129    if mtp_ladder_disabled() {
130        return num_drafts;
131    }
132    ladder_drafts_from_steps(mtp_ladder_steps(), n_active, num_drafts)
133}
134
135/// 2026-09-29: Engine default of the multi-sequence MTP dispatch cap, when neither
136/// `--mtp-max-seqs` nor MODEL.toml `[behavior] mtp_max_seqs` gives one. 4 when
137/// `METRALE_NO_MTP_K_LADDER` is present.
138pub const DEFAULT_MTP_MAX_SEQS: usize = 32;
139
140/// 2026-09-29: The widest cap a serve accepts: one batched verify's sequence width
141/// (`layer::VERIFY_WY_TABLE_SEQS`), which is also the serve's widest batch.
142pub const MAX_MTP_MAX_SEQS: usize = crate::layer::VERIFY_WY_TABLE_SEQS;
143
144/// 2026-09-29: Where the serve's dispatch cap came from.
145#[derive(Clone, Copy, Debug, PartialEq, Eq)]
146pub enum MtpMaxSeqsSource {
147    Cli,
148    ModelDefault,
149    EngineDefault,
150}
151
152/// 2026-09-29: Resolve the multi-sequence MTP dispatch cap: `--mtp-max-seqs` when given, else
153/// MODEL.toml `[behavior] mtp_max_seqs` when above 0, else the engine default
154/// ([`DEFAULT_MTP_MAX_SEQS`], 4 under `METRALE_NO_MTP_K_LADDER`). A cap outside
155/// `1..=MAX_MTP_MAX_SEQS` is refused, from either source.
156pub fn resolve_mtp_max_seqs(
157    cli: Option<usize>,
158    model_default: u32,
159    ladder_disabled: bool,
160) -> anyhow::Result<(usize, MtpMaxSeqsSource)> {
161    let (n, source) = match (cli, model_default) {
162        (Some(n), _) => (n, MtpMaxSeqsSource::Cli),
163        (None, m) if m > 0 => (m as usize, MtpMaxSeqsSource::ModelDefault),
164        (None, _) => (
165            if ladder_disabled {
166                4
167            } else {
168                DEFAULT_MTP_MAX_SEQS
169            },
170            MtpMaxSeqsSource::EngineDefault,
171        ),
172    };
173    anyhow::ensure!(
174        (1..=MAX_MTP_MAX_SEQS).contains(&n),
175        "MTP dispatch cap {n} ({source:?}) is outside 1..={MAX_MTP_MAX_SEQS}"
176    );
177    Ok((n, source))
178}
179
180static PUBLISHED_MTP_MAX_SEQS: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
181
182/// 2026-09-29: Publish the serve's dispatch cap once, before the model sizes its verify pools
183/// and the scheduler reads its levers. A second publication of a different value is refused.
184pub fn set_mtp_max_seqs(n: usize) -> anyhow::Result<()> {
185    let got = *PUBLISHED_MTP_MAX_SEQS.get_or_init(|| n);
186    anyhow::ensure!(
187        got == n,
188        "MTP dispatch cap already published as {got}, refusing {n}"
189    );
190    Ok(())
191}
192
193/// 2026-09-29: The multi-sequence MTP dispatch cap the serve published
194/// ([`set_mtp_max_seqs`]); before publication (unit tests, tools) the engine default.
195///
196/// The scheduler runs a speculative step only while the active count is at
197/// most this (`SchedLevers::mtp_max_seqs`, checked in `lane_decode.rs`), and
198/// the model sizes its verify pools and its single-sequence MTP structures
199/// from the same value (`ssm_reserve::mtp_state_slots`,
200/// `speculative::mtp_multi_seq_mode`).
201pub fn mtp_max_seqs() -> usize {
202    PUBLISHED_MTP_MAX_SEQS
203        .get()
204        .copied()
205        .unwrap_or(if mtp_ladder_disabled() {
206            4
207        } else {
208            DEFAULT_MTP_MAX_SEQS
209        })
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    // 2026-09-25: Assumes the test process sets neither `METRALE_MTP_K_LADDER`
217    // nor `METRALE_NO_MTP_K_LADDER`, which are latched once per process.
218    #[test]
219    fn default_ladder_holds_depth_to_the_cap() {
220        assert_eq!(mtp_ladder_drafts(1, 3), 3);
221        assert_eq!(mtp_ladder_drafts(4, 3), 3);
222        assert_eq!(mtp_ladder_drafts(5, 3), 3);
223        assert_eq!(mtp_ladder_drafts(8, 3), 3);
224        assert_eq!(mtp_ladder_drafts(9, 3), 1);
225        assert_eq!(mtp_ladder_drafts(16, 3), 1);
226        assert_eq!(mtp_ladder_drafts(17, 3), 1);
227        assert_eq!(mtp_ladder_drafts(32, 3), 1);
228        assert_eq!(mtp_ladder_drafts(64, 3), 1);
229        assert_eq!(mtp_ladder_drafts(16, 1), 1);
230    }
231
232    /// 2026-10-03: The flag's parser refuses what the environment variable silently
233    /// replaces with the default ladder.
234    #[test]
235    fn the_flag_parser_refuses_a_malformed_or_zero_ladder() {
236        assert_eq!(parse_mtp_k_ladder("4:1, 2:2").unwrap(), [(2, 2), (4, 1)]);
237        for bad in ["", "2", "2:x", "0:2", "2:0", "2:2,,4:1"] {
238            assert!(parse_mtp_k_ladder(bad).is_err(), "{bad:?} accepted");
239        }
240    }
241
242    #[test]
243    fn depth_at_width_env_rungs_parse_shape() {
244        let steps = parse_ladder("32:2, 4:3,8:3, 16:2,24:2").unwrap();
245        assert_eq!(steps, [(4, 3), (8, 3), (16, 2), (24, 2), (32, 2)]);
246        assert_eq!(ladder_drafts_from_steps(&steps, 17, 3), 2);
247        assert_eq!(ladder_drafts_from_steps(&steps, 24, 3), 2);
248        assert_eq!(ladder_drafts_from_steps(&steps, 25, 3), 2);
249        assert_eq!(ladder_drafts_from_steps(&steps, 32, 3), 2);
250    }
251
252    /// 2026-09-29: The flag beats MODEL.toml, MODEL.toml beats the engine default, 0 in
253    /// MODEL.toml means unset, and a cap outside 1..=128 is refused from either source.
254    #[test]
255    fn mtp_max_seqs_resolution_order_and_bounds() {
256        use MtpMaxSeqsSource::*;
257        assert_eq!(resolve_mtp_max_seqs(Some(8), 128, false).unwrap(), (8, Cli));
258        assert_eq!(
259            resolve_mtp_max_seqs(None, 128, false).unwrap(),
260            (128, ModelDefault)
261        );
262        assert_eq!(
263            resolve_mtp_max_seqs(None, 0, false).unwrap(),
264            (32, EngineDefault)
265        );
266        assert_eq!(
267            resolve_mtp_max_seqs(None, 0, true).unwrap(),
268            (4, EngineDefault)
269        );
270        assert_eq!(
271            resolve_mtp_max_seqs(Some(32), 128, false).unwrap(),
272            (32, Cli)
273        );
274        assert!(resolve_mtp_max_seqs(Some(0), 128, false).is_err());
275        assert!(resolve_mtp_max_seqs(Some(129), 0, false).is_err());
276        assert!(resolve_mtp_max_seqs(None, 129, false).is_err());
277        assert_eq!(MAX_MTP_MAX_SEQS, 128);
278    }
279
280    #[test]
281    fn explicit_steps_are_honored() {
282        let steps = parse_ladder("4:3,8:2").unwrap();
283        assert_eq!(ladder_drafts_from_steps(&steps, 4, 3), 3);
284        assert_eq!(ladder_drafts_from_steps(&steps, 5, 3), 2);
285        assert_eq!(ladder_drafts_from_steps(&steps, 8, 3), 2);
286        assert_eq!(ladder_drafts_from_steps(&steps, 9, 3), 2);
287    }
288
289    #[test]
290    fn malformed_ladder_is_rejected_as_a_unit() {
291        assert_eq!(parse_ladder(""), None);
292        assert_eq!(parse_ladder("4:3,broken,8:2"), None);
293        assert_eq!(parse_ladder("4:three"), None);
294    }
295
296    #[test]
297    fn ladder_clamps_to_configured_ceiling() {
298        assert_eq!(mtp_ladder_drafts(2, 1), 1);
299        assert_eq!(mtp_ladder_drafts(2, 0), 0);
300    }
301}
metrale_model_layers/speculative/ladder.rs

metrale_model_layers/speculative/
ladder.rs
