metrale_model_layers/speculative/
verify_key.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Batched-verify ordering and graph key: the dispatch order and
4//! depth assignment of one verify batch, and the CUDA-graph cache key built
5//! from them.
6//!
7//! Owner: model-layers (speculative).
8//! Invariants:
9//! - `verify_batch_permutation` returns a permutation of `0..n`, and panics
10//!   when `slots` and `ks` differ in length.
11//! - With `canonical = true`, `verify_batch_order` places slots in
12//!   non-decreasing order and depths in non-increasing order, and returns the
13//!   multiset of depths it was given.
14//! - With `canonical = false`, each member keeps its own depth, ordered
15//!   deepest first, then by slot.
16//!
17//! The scheduler orders each batch with these functions (`mtp_dcut::plan`,
18//! `mtp_step`) and the model keys its batched-verify graphs with
19//! [`verify_graph_key`] (`verify_e2::verify_batched_graph_key`), so the
20//! dispatch and the key see one order. Under the canonical assignment the key
21//! depends on the slot set and the depth multiset, not on which sequence D-Cut
22//! gave which depth, which reduces the number of distinct keys. Ascending
23//! slots also keep a depth run on consecutive slots when the pool slots are
24//! consecutive; the batched GDN path checks the actual pointers and declines
25//! otherwise (`trait_decode_batched_conv_gdn_multi.rs`).
26//!
27//! [`canonical_assignment`] chooses the arm from the batch width.
28//! `METRALE_NO_CANONICAL_VERIFY_KEY`, present with any value (`0` included),
29//! turns the canonical arm off at every width.
30
31/// 2026-09-25: True unless `METRALE_NO_CANONICAL_VERIFY_KEY` is present. Read
32/// once per process.
33pub fn canonical_verify_key_enabled() -> bool {
34    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
35    *ON.get_or_init(|| std::env::var_os("METRALE_NO_CANONICAL_VERIFY_KEY").is_none())
36}
37
38/// 2026-09-25: Batch width (sequences) at or above which [`canonical_assignment`]
39/// selects the canonical arm, unless `METRALE_CANONICAL_KEY_MIN_WIDTH` overrides
40/// it. Above the scheduler's D-Cut width cap (`SchedLevers::dcut_width_cap`),
41/// `mtp_dcut::plan` leaves every depth equal, and then both arms give the same
42/// order (`uniform_depths_are_identical_under_both_arms`).
43pub const CANONICAL_KEY_MIN_WIDTH: usize = 8;
44
45/// 2026-09-25: The width threshold: `METRALE_CANONICAL_KEY_MIN_WIDTH` when it
46/// parses as a `usize` (0 selects the canonical arm at every width), else
47/// [`CANONICAL_KEY_MIN_WIDTH`]. Read once per process.
48pub fn canonical_key_min_width() -> usize {
49    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
50    *N.get_or_init(|| min_width_from_env(std::env::var_os(ENV_MIN_WIDTH)))
51}
52
53const ENV_MIN_WIDTH: &str = "METRALE_CANONICAL_KEY_MIN_WIDTH";
54
55/// 2026-09-25: Parses a raw [`ENV_MIN_WIDTH`] value; [`canonical_key_min_width`]
56/// does the environment read, so tests call this directly.
57fn min_width_from_env(raw: Option<std::ffi::OsString>) -> usize {
58    raw.and_then(|v| v.into_string().ok())
59        .and_then(|v| v.trim().parse::<usize>().ok())
60        .unwrap_or(CANONICAL_KEY_MIN_WIDTH)
61}
62
63/// 2026-09-25: Whether a batch of `n` sequences gets the canonical depth-to-slot
64/// assignment. `mtp_dcut::plan` asks with the batch width and `mtp_step` asks
65/// with the same width for each chunk it dispatches, so the order `plan`
66/// assigned and the permutation `mtp_step` applies use the same arm.
67pub fn canonical_assignment(n: usize) -> bool {
68    canonical_assignment_at(n, canonical_key_min_width(), canonical_verify_key_enabled())
69}
70
71/// 2026-09-25: The policy of [`canonical_assignment`] with its two
72/// environment inputs passed in: the threshold, and whether the kill switch
73/// is clear. With the kill switch set, no width or threshold selects the
74/// canonical arm.
75fn canonical_assignment_at(n: usize, min_width: usize, kill_switch_clear: bool) -> bool {
76    kill_switch_clear && n >= min_width
77}
78
79/// 2026-09-25: Dispatch order for one verify batch, as a permutation only.
80///
81/// `slots[i]` and `ks[i]` describe member `i` in the caller's order (`ks[i]`
82/// is its row count, drafts + 1). `order[p]` is the input index dispatched at
83/// position `p`.
84///
85/// * `canonical = true`: ascending ssm slot; `ks` is unread.
86/// * `canonical = false`: deepest first, then ascending ssm slot.
87///
88/// Ties break on input index, so the order is a deterministic function of
89/// the inputs. Under `canonical = true` a second application to the ordered
90/// batch returns the identity (`canonical_order_is_idempotent`).
91///
92/// Panics when `slots` and `ks` differ in length.
93pub fn verify_batch_permutation(slots: &[usize], ks: &[usize], canonical: bool) -> Vec<usize> {
94    assert_eq!(
95        slots.len(),
96        ks.len(),
97        "verify_batch_permutation: slots/ks mismatch"
98    );
99    let n = slots.len().min(ks.len());
100    let mut order: Vec<usize> = (0..n).collect();
101    if canonical {
102        order.sort_by_key(|&i| (slots[i], i));
103    } else {
104        order.sort_by_key(|&i| (std::cmp::Reverse(ks[i]), slots[i], i));
105    }
106    order
107}
108
109/// 2026-09-25: Orders one verify batch and assigns its depths; the entry point
110/// of `mtp_dcut::plan`, which truncates each sequence's drafts to the depth
111/// assigned here.
112///
113/// Returns `(order, depths)`: `order` is [`verify_batch_permutation`] and
114/// `depths[p]` is the row count position `p` verifies.
115///
116/// * `canonical = true`: the depth multiset is sorted descending onto the
117///   ordered batch, so `depths[p]` need not be `ks[order[p]]`. Slots are
118///   non-decreasing in `p` and depths non-increasing.
119/// * `canonical = false`: each member keeps its own depth.
120///
121/// Stages that only need the dispatch order use [`verify_batch_permutation`],
122/// which leaves every member's depth attached to it.
123pub fn verify_batch_order(
124    slots: &[usize],
125    ks: &[usize],
126    canonical: bool,
127) -> (Vec<usize>, Vec<usize>) {
128    let order = verify_batch_permutation(slots, ks, canonical);
129    let depths: Vec<usize> = if canonical {
130        let mut d: Vec<usize> = ks[..order.len()].to_vec();
131        d.sort_unstable_by(|a, b| b.cmp(a));
132        d
133    } else {
134        order.iter().map(|&i| ks[i]).collect()
135    };
136    (order, depths)
137}
138
139/// 2026-09-25: The batched-verify CUDA-graph cache key: the `(ssm slot, row
140/// count)` pairs in dispatch order, then one sentinel word,
141/// `u32::MAX - wy_tables_null - 2 * write_on_accept`, which differs for each
142/// of the four flag combinations. `write_on_accept` is in the key because it
143/// can select a different WY launch (the write-on-accept K=4 twin,
144/// `trait_decode_batched_conv_gdn_multi.rs`).
145pub fn verify_graph_key(
146    pairs: &[(u32, u32)],
147    wy_tables_null: bool,
148    write_on_accept: bool,
149) -> Vec<u32> {
150    let mut key: Vec<u32> = Vec::with_capacity(2 * pairs.len() + 1);
151    for &(slot, k) in pairs {
152        key.push(slot);
153        key.push(k);
154    }
155    key.push(u32::MAX - u32::from(wy_tables_null) - 2 * u32::from(write_on_accept));
156    key
157}
158
159#[cfg(test)]
160#[path = "verify_key_tests.rs"]
161mod tests;
metrale_model_layers/speculative/verify_key.rs

metrale_model_layers/speculative/
verify_key.rs
