metrale_model_layers/
drafter_context.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Resolution of the MTP drafter context from the environment: whether
4//! the drafter KV is batch-prefilled over the prompt (`prefill`) and whether it is
5//! carried across the turns of a session (`carry`).
6//!
7//! Carry depends on prefill. In the model's `trait_impl/speculative.rs` a carried
8//! drafter KV is adopted (`try_carry_drafter`) only inside
9//! `if !self.mtp_prefill_hidden.is_null()`, and the model allocates that capture
10//! only when `drafter.prefill` is set (`layers::mtp_drafter_prefill_enabled`), so
11//! the two are resolved together.
12//! A configured carry is armed only in single-sequence MTP
13//! (`mtp_carry::carry_armed`).
14//!
15//! | variable (exactly `1`) | result |
16//! |---|---|
17//! | neither | `DrafterContext::BOTH` |
18//! | `METRALE_NO_MTP_DRAFTER_CONTEXT` | `DrafterContext::OFF` |
19//! | `METRALE_MTP_DRAFTER_CONTEXT_PREFILL_ONLY_UNSAFE` | prefill on, carry off; warned at startup |
20//!
21//! Any other value, `0` included, changes nothing.
22//!
23//! Owner: model-layers (MTP drafter).
24//! Invariants:
25//! - [`resolve`] never returns `carry` without `prefill`.
26//! - The kill switch wins over the prefill-only variable.
27
28/// 2026-09-25: Kill switch for both halves; only the exact value `1` counts.
29pub const DISABLE_ENV: &str = "METRALE_NO_MTP_DRAFTER_CONTEXT";
30
31/// 2026-09-25: Prefill without carry; only the exact value `1` counts, and
32/// [`DISABLE_ENV`] wins over it.
33pub const PREFILL_ONLY_ENV: &str = "METRALE_MTP_DRAFTER_CONTEXT_PREFILL_ONLY_UNSAFE";
34
35/// 2026-09-25: Which halves of the drafter context are configured. [`resolve`]
36/// never returns `carry` without `prefill`.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct DrafterContext {
39    /// 2026-09-25: Capture prompt hiddens and batch-prefill the drafter KV.
40    pub prefill: bool,
41    /// 2026-09-25: Carry the drafter KV across turns of a session; it runs only
42    /// when `mtp_carry::carry_armed` also holds.
43    pub carry: bool,
44}
45
46impl Default for DrafterContext {
47    /// 2026-09-25: [`BOTH`](DrafterContext::BOTH). Written out because a derived
48    /// `Default` would be [`OFF`](DrafterContext::OFF).
49    fn default() -> Self {
50        Self::BOTH
51    }
52}
53
54impl DrafterContext {
55    /// 2026-09-25: Both halves on: what [`resolve`] returns when neither variable is `1`.
56    pub const BOTH: Self = Self {
57        prefill: true,
58        carry: true,
59    };
60    /// 2026-09-25: Both halves off: what [`resolve`] returns for the kill switch.
61    pub const OFF: Self = Self {
62        prefill: false,
63        carry: false,
64    };
65}
66
67/// 2026-09-25: Resolve the two variables' values with no environment access, so
68/// the rule is unit-tested (SBIO). `disable` wins over `prefill_only`.
69pub fn resolve(disable: Option<&str>, prefill_only: Option<&str>) -> DrafterContext {
70    if disable == Some("1") {
71        return DrafterContext::OFF;
72    }
73    if prefill_only == Some("1") {
74        // 2026-09-25: The only branch that returns prefill without carry.
75        return DrafterContext {
76            prefill: true,
77            carry: false,
78        };
79    }
80    DrafterContext::BOTH
81}
82
83/// 2026-09-25: Resolve the policy from the environment. `ModelLevers::from_env`
84/// calls it, and the result is `ModelLevers::drafter`. The outcome is logged once
85/// per process (`REPORTED`).
86pub fn resolve_from_env() -> DrafterContext {
87    {
88        let disable = std::env::var(DISABLE_ENV).ok();
89        let prefill_only = std::env::var(PREFILL_ONLY_ENV).ok();
90        let cfg = resolve(disable.as_deref(), prefill_only.as_deref());
91
92        REPORTED.call_once(|| report(cfg));
93        cfg
94    }
95}
96
97/// 2026-09-25: Makes `report` run once per process, on the first [`resolve_from_env`].
98static REPORTED: std::sync::Once = std::sync::Once::new();
99
100fn report(cfg: DrafterContext) {
101    // 2026-09-25: Report the armed carry (`mtp_carry::carry_armed`, the predicate
102    // the runtime gates on), not the configured one: the MTP dispatch cap can
103    // force the carry off under this config.
104    let carry_armed = crate::mtp_carry::carry_armed(cfg);
105    tracing::info!(
106        "MTP drafter context: prefill={} carry={} ({}). Disable both with {}=1.",
107        on_off(cfg.prefill),
108        on_off(carry_armed),
109        if cfg == DrafterContext::BOTH {
110            "default"
111        } else {
112            "overridden by environment"
113        },
114        DISABLE_ENV,
115    );
116    if cfg.carry && !carry_armed {
117        tracing::warn!(
118            "MTP cross-turn carry is CONFIGURED ON but INERT: the MTP \
119             dispatch cap is {} (>1), and the carry slot is single-sequence \
120             by design, so it is force-disabled. Pass --mtp-max-seqs 1 \
121             to arm it; leave it unset to keep multi-sequence MTP.",
122            crate::speculative::mtp_max_seqs(),
123        );
124    }
125    if cfg.prefill && !cfg.carry {
126        tracing::warn!(
127            "{PREFILL_ONLY_ENV}=1: drafter prefill is ON with cross-turn carry \
128             OFF. This is a MEASUREMENT ARM, not a deployment — a warm turn \
129             rebuilds the drafter for a measured 1136 ms against ~211 ms/turn \
130             of decode saving (net -927 ms/turn, spent on TTFT).",
131        );
132    }
133}
134
135fn on_off(b: bool) -> &'static str {
136    if b { "ON" } else { "OFF" }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn bare_environment_enables_both() {
145        assert_eq!(resolve(None, None), DrafterContext::BOTH);
146    }
147
148    #[test]
149    fn kill_switch_disables_both() {
150        assert_eq!(resolve(Some("1"), None), DrafterContext::OFF);
151    }
152
153    #[test]
154    fn prefill_only_arm_disables_carry_only() {
155        assert_eq!(
156            resolve(None, Some("1")),
157            DrafterContext {
158                prefill: true,
159                carry: false
160            }
161        );
162    }
163
164    #[test]
165    fn kill_switch_beats_the_research_arm() {
166        assert_eq!(resolve(Some("1"), Some("1")), DrafterContext::OFF);
167    }
168
169    #[test]
170    fn only_exactly_one_switches_anything() {
171        for v in ["0", "", "true", "yes", "2", "1 "] {
172            assert_eq!(
173                resolve(Some(v), None),
174                DrafterContext::BOTH,
175                "{DISABLE_ENV}={v:?} must not disable",
176            );
177            assert_eq!(
178                resolve(None, Some(v)),
179                DrafterContext::BOTH,
180                "{PREFILL_ONLY_ENV}={v:?} must not change anything",
181            );
182        }
183    }
184}
metrale_model_layers/drafter_context.rs

metrale_model_layers/
drafter_context.rs
