metrale_model_layers/speculative/
knobs.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: MTP drafter knobs read from the environment: shadow top-k, multi-sequence mode, the catch-up and
4//! refeed switches, debug taps, and the EP propose command.
5//!
6//! Owner: model-layers (speculative).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11/// 2026-09-25: `METRALE_MTP_SHADOW_TOPK=k`: log the drafter's top-k
12/// candidates at each draft position, and the verify steps' target tokens,
13/// without changing token selection. 0 (unset or unparseable) is off; k is
14/// capped at 8. `ModelLevers::shadow_topk` and `SchedLevers::shadow_topk`
15/// both read it here, once per run.
16pub fn shadow_topk() -> usize {
17    std::env::var("METRALE_MTP_SHADOW_TOPK")
18        .ok()
19        .and_then(|v| v.parse::<usize>().ok())
20        .unwrap_or(0)
21        .min(8)
22}
23
24/// 2026-09-25: True when the MTP dispatch cap is above one sequence. The
25/// catch-up ring, the refeed labels and the carry slot each hold one
26/// sequence's data, so [`mtp_catchup_enabled`], [`mtp_refeed_accepted_enabled`]
27/// and `mtp_carry::carry_armed` are false in this mode.
28pub fn mtp_multi_seq_mode() -> bool {
29    mtp_max_seqs() > 1
30}
31
32/// 2026-09-25: `METRALE_MTP_ACCEPT_DEBUG`, on when set to any value. It turns
33/// on MTP acceptance logs, among them the periodic per-batch-width lines of
34/// the scheduler's `mtp_accept_debug` module, verify-graph outcome counts and
35/// a drafter coverage line. Read once per process.
36pub fn mtp_accept_debug() -> bool {
37    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
38    *ON.get_or_init(|| std::env::var("METRALE_MTP_ACCEPT_DEBUG").is_ok())
39}
40
41/// 2026-09-25: `METRALE_MTP_CATCHUP=1`: the model keeps a ring of target
42/// hiddens, and before a propose it feeds the drafter the pair keys missing
43/// since its newest row. Off otherwise, and off in [`mtp_multi_seq_mode`].
44pub fn mtp_catchup_enabled() -> bool {
45    std::env::var("METRALE_MTP_CATCHUP").ok().as_deref() == Some("1") && !mtp_multi_seq_mode()
46}
47
48/// 2026-09-25: `METRALE_MTP_REFEED_ACCEPTED=1`: after a verify, ring the
49/// target's hiddens for the accepted positions, so the catch-up feed can
50/// rewrite accepted drafter rows with the target's hidden instead of the
51/// drafter's own. It writes into the catch-up ring, which the model allocates
52/// only when [`mtp_catchup_enabled`]. Off otherwise, and off in
53/// [`mtp_multi_seq_mode`]. Measured 2026-07-21 on dgx2 at 2 drafts, about 10k
54/// verify steps per arm: p2_uncond 0.4182 off, 0.4452 on.
55pub fn mtp_refeed_accepted_enabled() -> bool {
56    std::env::var("METRALE_MTP_REFEED_ACCEPTED").ok().as_deref() == Some("1")
57        && !mtp_multi_seq_mode()
58}
59
60/// 2026-09-25: `METRALE_MTP_REFEED_SHIFT`: an offset added to every re-fed
61/// ring label, clamped to `[-4, 4]`; 0 (unset or unparseable) is the derived
62/// mapping. It exists to test that mapping: a run at +1 or -1 is compared
63/// with 0 by acceptance.
64pub fn mtp_refeed_shift() -> isize {
65    std::env::var("METRALE_MTP_REFEED_SHIFT")
66        .ok()
67        .and_then(|v| v.parse::<isize>().ok())
68        .unwrap_or(0)
69        .clamp(-4, 4)
70}
71
72/// 2026-09-25: `METRALE_MTP_REFEED_DEBUG=1`: log a [`hidden_fingerprint`] of
73/// each hidden written to or fed from the catch-up ring. Each one copies a
74/// row to the host.
75pub fn mtp_refeed_debug() -> bool {
76    std::env::var("METRALE_MTP_REFEED_DEBUG").ok().as_deref() == Some("1")
77}
78
79/// 2026-09-25: FNV-1a over the `h` BF16 values at `p`; 0 when the copy to the
80/// host fails.
81pub fn hidden_fingerprint(gpu: &dyn GpuBackend, p: DevicePtr, h: usize) -> u64 {
82    let mut b = vec![0u8; h * 2];
83    if gpu.copy_d2h(p, &mut b).is_err() {
84        return 0;
85    }
86    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
87    for byte in &b {
88        hash ^= *byte as u64;
89        hash = hash.wrapping_mul(0x1000_0000_01b3);
90    }
91    hash
92}
93
94/// 2026-09-25: EP worker command: run one MTP propose alongside rank 0. The
95/// code is followed by `last_token`, `position` and `num_drafts`, each a u32.
96pub const EP_CMD_MTP_PROPOSE: u32 = 0xFFFF_FFF5;
97
98/// 2026-09-25: Run the sharded drafter on every rank with the communicator.
99/// On unless `METRALE_NO_MTP_EP_PROPOSE=1`; read once per process. It is the
100/// GLM-5.3 MTP head's [`DraftProposer::needs_comm`], and a proposer that needs
101/// the communicator also has the worker run each propose
102/// ([`EP_CMD_MTP_PROPOSE`]).
103pub fn mtp_ep_propose_enabled() -> bool {
104    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
105    *ON.get_or_init(|| std::env::var("METRALE_NO_MTP_EP_PROPOSE").ok().as_deref() != Some("1"))
106}
metrale_model_layers/speculative/knobs.rs

metrale_model_layers/speculative/
knobs.rs
