metrale_model_layers/ssm_reserve/
copy_tier.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The prompt-lookup copy tier of the SSM verify pools
4//! (`--prompt-lookup-decoding`): the first `slots` verify slots hold at least
5//! `drafts` h intermediates, so a copy longer than the MTP chain can be
6//! verified there. The serve publishes it once, before the preflight reserve
7//! and the pool allocation, which both read it through [`pool_counts`]
8//! (`super::pool_counts`); neither can then size the pool differently.
9//!
10//! Owner: model-layers (SSM reserve).
11//! Invariants:
12//! - A published tier has `slots >= 1` and `drafts >= 1`; a second publication
13//!   of a different tier is refused.
14//! - [`tier_h`] never lowers a slot's count, and raises only slots below
15//!   `slots`.
16
17/// 2026-10-02: Copy-tier geometry.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub struct CopyTier {
20    /// 2026-10-02: Verify slots `0..slots` get the copy depth.
21    pub slots: usize,
22    /// 2026-10-02: Drafts those slots can verify (h intermediates).
23    pub drafts: usize,
24}
25
26static PUBLISHED: std::sync::OnceLock<CopyTier> = std::sync::OnceLock::new();
27
28/// 2026-10-02: Publish the copy tier, before the preflight reserve.
29pub fn set_copy_tier(tier: CopyTier) -> anyhow::Result<()> {
30    anyhow::ensure!(
31        tier.slots >= 1 && tier.drafts >= 1,
32        "copy tier needs at least one slot and one draft, got {tier:?}"
33    );
34    let got = *PUBLISHED.get_or_init(|| tier);
35    anyhow::ensure!(
36        got == tier,
37        "copy tier already published as {got:?}, refusing {tier:?}"
38    );
39    Ok(())
40}
41
42/// 2026-10-02: The published tier; `None` without prompt lookup.
43pub fn copy_tier() -> Option<CopyTier> {
44    PUBLISHED.get().copied()
45}
46
47/// 2026-10-02: Slot `s`'s h intermediates under `tier`: `base` raised to the
48/// copy depth on a copy slot.
49pub fn tier_h(base: usize, s: usize, tier: Option<CopyTier>) -> usize {
50    match tier {
51        Some(t) if s < t.slots => base.max(t.drafts),
52        _ => base,
53    }
54}
55
56/// 2026-10-02: Rows of the widest verify the pools must hold: the MTP `k`
57/// (drafts + 1), widened to a copy's.
58pub fn tier_rows(k: usize, tier: Option<CopyTier>) -> usize {
59    tier.map_or(k, |t| k.max(t.drafts + 1))
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    const T: CopyTier = CopyTier {
67        slots: 8,
68        drafts: 8,
69    };
70
71    #[test]
72    fn only_copy_slots_are_raised_and_never_lowered() {
73        assert_eq!(tier_h(3, 0, Some(T)), 8);
74        assert_eq!(tier_h(3, 7, Some(T)), 8);
75        assert_eq!(tier_h(3, 8, Some(T)), 3, "slot 8 is outside the tier");
76        assert_eq!(
77            tier_h(12, 0, Some(T)),
78            12,
79            "a deeper ladder slot keeps its depth"
80        );
81        assert_eq!(tier_h(3, 0, None), 3);
82    }
83
84    #[test]
85    fn rows_widen_to_the_copy() {
86        assert_eq!(tier_rows(4, Some(T)), 9);
87        assert_eq!(tier_rows(17, Some(T)), 17);
88        assert_eq!(tier_rows(2, None), 2);
89    }
90
91    #[test]
92    fn an_empty_tier_is_refused() {
93        assert!(
94            set_copy_tier(CopyTier {
95                slots: 0,
96                drafts: 4
97            })
98            .is_err()
99        );
100        assert!(
101            set_copy_tier(CopyTier {
102                slots: 4,
103                drafts: 0
104            })
105            .is_err()
106        );
107    }
108}