metrale_speculative/
spec_capacity.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Draft-count limits for speculative dispatch: the verify-slot
4//! capacity of every active sequence, and the serial path's verify arm.
5//!
6//! Owner: speculative.
7//! Invariants:
8//! - `clamp_drafts_to_slot_capacity` returns no more than `drafts` and no
9//!   more than any capacity it is given.
10//! - `serial_verify_plan` keeps at most `step_drafts` of an MTP sequence's
11//!   pending drafts.
12//!
13//! The verify pools size each slot for the deepest draft count the static
14//! ladder can give a sequence in it (`ssm_reserve::verify_slot_drafts`).
15//! With the default ladder `4:3,8:3,16:1,32:1` and `--num-drafts 3`, slots
16//! 0..8 hold 3 drafts and slots from 8 on hold 1. A step can ask for more:
17//! a sequence can sit on a high slot while `n_active` is small, and
18//! `adaptive_rung` can raise n in 9..=16 to 2 drafts. A step's draft count is
19//! therefore clamped to the smallest capacity among its active sequences,
20//! read from the model's pools (`Model::mtp_slot_draft_capacity`).
21//! `METRALE_MTP_POOL_FULL_WIDTH` gives every slot `num_drafts`, which makes
22//! the clamp a no-op. With the tiered default the adaptive lift never
23//! survives the clamp: nine or more active sequences hold nine distinct
24//! slots, one of them at index 8 or above, whose capacity is 1.
25
26/// 2026-09-25: Clamps a step's draft count to the smallest verify-slot
27/// capacity among the active sequences. `usize::MAX` (the `Model` default,
28/// for a model without SSM verify pools) does not clamp; an empty iterator
29/// leaves `drafts` unchanged.
30pub fn clamp_drafts_to_slot_capacity(
31    drafts: usize,
32    slot_capacities: impl IntoIterator<Item = usize>,
33) -> usize {
34    slot_capacities
35        .into_iter()
36        .fold(drafts, |acc, cap| acc.min(cap))
37}
38
39/// 2026-09-25: The verify arm the serial (per-sequence) path runs for one
40/// sequence.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum SerialArm {
43    DFlash,
44    /// 2026-10-02: Four or more drafts on an MTP serve (a long prompt-lookup
45    /// copy): one K-row verify, then the K4 verdict.
46    KN,
47    K4,
48    K3,
49    K2,
50}
51
52/// 2026-09-25: How many of a sequence's `pending` drafts the serial path
53/// keeps, and which arm verifies them. An MTP sequence can carry more
54/// pending drafts than this step's width `step_drafts` (the ladder count
55/// after [`clamp_drafts_to_slot_capacity`]); the surplus is dropped so the
56/// verify never exceeds the slot. With `dflash` every pending draft is kept.
57/// Four or more kept drafts run the DFlash arm on a DFlash serve and the KN arm
58/// otherwise; fewer run K4, K3 or K2 by the kept count, each also bounded by
59/// `step_drafts`.
60pub fn serial_verify_plan(pending: usize, step_drafts: usize, dflash: bool) -> (usize, SerialArm) {
61    let keep = if dflash {
62        pending
63    } else {
64        pending.min(step_drafts)
65    };
66    let arm = if keep >= 4 && dflash {
67        SerialArm::DFlash
68    } else if keep >= 4 {
69        SerialArm::KN
70    } else if step_drafts >= 3 && keep >= 3 {
71        SerialArm::K4
72    } else if step_drafts >= 2 && keep >= 2 {
73        SerialArm::K3
74    } else {
75        SerialArm::K2
76    };
77    (keep, arm)
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn a_low_capacity_slot_bounds_the_whole_step() {
86        // 2026-09-25: One capacity-1 slot limits the whole step to one draft.
87        assert_eq!(clamp_drafts_to_slot_capacity(3, [3, 3, 1]), 1);
88        // 2026-09-25: One sequence on a high slot at small n.
89        assert_eq!(clamp_drafts_to_slot_capacity(3, [1]), 1);
90        // 2026-09-25: The adaptive-rung lift (2 drafts at n in 9..=16) is
91        // clamped by a capacity-1 slot in the batch.
92        assert_eq!(clamp_drafts_to_slot_capacity(2, [3, 1, 3]), 1);
93    }
94
95    #[test]
96    fn full_capacity_slots_do_not_clamp() {
97        assert_eq!(clamp_drafts_to_slot_capacity(3, [3, 3, 3]), 3);
98        // 2026-09-25: `usize::MAX` (no SSM verify pools) does not clamp.
99        assert_eq!(clamp_drafts_to_slot_capacity(3, [usize::MAX; 4]), 3);
100        // 2026-09-25: An empty iterator leaves the count unchanged.
101        assert_eq!(clamp_drafts_to_slot_capacity(2, std::iter::empty()), 2);
102    }
103
104    #[test]
105    fn a_ladder_below_num_drafts_never_reaches_the_k4_verify_at_n1() {
106        // 2026-09-25: A one-draft ladder (`METRALE_MTP_K_LADDER=1:1`) under
107        // `--num-drafts 3`: slot 0 holds one draft, so the step width is 1 and
108        // a sequence carrying three or two pending drafts keeps one and runs
109        // K2.
110        let step = clamp_drafts_to_slot_capacity(1, [1]);
111        assert_eq!(serial_verify_plan(3, step, false), (1, SerialArm::K2));
112        assert_eq!(serial_verify_plan(2, step, false), (1, SerialArm::K2));
113        // 2026-09-25: A two-draft step keeps two of three and runs K3.
114        assert_eq!(serial_verify_plan(3, 2, false), (2, SerialArm::K3));
115    }
116
117    #[test]
118    fn a_full_width_step_dispatches_as_before() {
119        assert_eq!(serial_verify_plan(3, 3, false), (3, SerialArm::K4));
120        assert_eq!(serial_verify_plan(2, 3, false), (2, SerialArm::K3));
121        assert_eq!(serial_verify_plan(1, 3, false), (1, SerialArm::K2));
122        assert_eq!(serial_verify_plan(0, 3, false), (0, SerialArm::K2));
123        // 2026-09-25: DFlash γ-blocks are never truncated to an MTP width.
124        assert_eq!(serial_verify_plan(16, 16, true), (16, SerialArm::DFlash));
125        assert_eq!(serial_verify_plan(3, 3, true), (3, SerialArm::K4));
126    }
127
128    #[test]
129    fn a_long_mtp_copy_takes_the_kn_arm_never_dflash() {
130        // 2026-10-02: A prompt-lookup copy of 8 on an MTP serve: the step width
131        // is the copy's, and 4+ drafts take the K-row MTP verify, not DFlash's
132        // (whose drafter trims and telemetry belong to a block drafter).
133        assert_eq!(serial_verify_plan(8, 8, false), (8, SerialArm::KN));
134        assert_eq!(serial_verify_plan(4, 4, false), (4, SerialArm::KN));
135        assert_eq!(serial_verify_plan(8, 3, false), (3, SerialArm::K4));
136    }
137}
metrale_speculative/spec_capacity.rs

metrale_speculative/
spec_capacity.rs
