metrale_model_arch/dflash_head/levers.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: DFlash drafter levers, resolved from the environment once per head and carried on it.
4//!
5//! Owner: model-arch (DFlash drafter).
6//! Invariants:
7//! - `DFlashLevers::from_env` is called only while a head is built (`from_weights.rs`),
8//! so a head's levers are fixed for its lifetime.
9//! - Each field reads one `METRALE_*` variable, except `any_diagnostic_armed`, which
10//! reads every name in `GRAPH_SUPPRESSING_DIAGNOSTICS`.
11//!
12//! The levers live on the head rather than in a process-wide `OnceLock`, so a second
13//! model built in the same process resolves the environment again.
14
15/// 2026-09-25: Diagnostic and A/B levers for one loaded DFlash drafter.
16// 2026-09-25: No `Eq`: `conf_tau` is an `f32`.
17#[derive(Clone, Copy, Debug, PartialEq, Default)]
18pub struct DFlashLevers {
19 /// 2026-09-25: A variable in `GRAPH_SUPPRESSING_DIAGNOSTICS` is set, at any
20 /// value. `forward_block` then runs eagerly, without CUDA graphs. This is
21 /// presence, not truth: `METRALE_DFLASH_BLOCK_DUMP=0` suppresses graphs and
22 /// arms no dump.
23 pub any_diagnostic_armed: bool,
24
25 /// 2026-09-25: `METRALE_DFLASH_DEBUG_DUMP=1`: log the first 10 BF16 values
26 /// of key `forward_block` intermediates.
27 pub debug_dump: bool,
28 /// 2026-09-25: `METRALE_DFLASH_DEBUG_DUMP_FULL=1`: one-shot whole-tensor dumps
29 /// to `/tmp`, and the first propose's drafts to the log.
30 pub debug_dump_full: bool,
31 /// 2026-09-25: `METRALE_DFLASH_LOG_DRAFTS=1`: log the first propose's drafts
32 /// (one-shot per model).
33 pub log_drafts: bool,
34 /// 2026-09-25: `METRALE_DFLASH_BLOCK_DUMP=1`: one-shot `/tmp` dumps of the
35 /// block logits and per-layer intermediates, gated by `block_dump_armed_at`.
36 pub block_dump: bool,
37 /// 2026-09-25: `METRALE_DFLASH_BLOCK_DUMP_AT_POS=<n>`: arm the block dump only
38 /// at decode position >= n. Unset or unparseable is 0.
39 pub block_dump_at_pos: usize,
40 /// 2026-09-25: `METRALE_DFLASH_OPTION_B_DIAG=1`: one-shot read-back of the
41 /// K row layer 0 just wrote to the paged drafter cache, compared with its
42 /// source.
43 pub option_b_diag: bool,
44
45 /// 2026-09-25: `METRALE_DFLASH_DEBUG_FORCE_PATTERN=1`: overwrite the first ctx
46 /// slot `forward_block` reads with a deterministic pattern.
47 pub force_pattern: bool,
48 /// 2026-09-25: `METRALE_DFLASH_DEBUG_FORCE_NOISE_PATTERN=1`: the same, for the
49 /// noise rows.
50 pub force_noise_pattern: bool,
51 /// 2026-09-25: `METRALE_DFLASH_DEBUG_CTX_OFF=1`: drop ctx conditioning
52 /// (`eff_ctx = 0`).
53 pub force_no_ctx: bool,
54 /// 2026-09-25: `METRALE_DFLASH_DEBUG_CTX_USED=<n>`: cap `eff_ctx` at n (it
55 /// stays bounded by the ctx rows present and `ctx_window`).
56 pub force_ctx_used: Option<usize>,
57
58 /// 2026-09-25: `METRALE_DFLASH_PRECOMPUTE=1`: also run the ctx K/V precompute
59 /// from `forward_block`; `propose_drafts` runs it either way.
60 pub precompute: bool,
61 /// 2026-09-25: `METRALE_DFLASH_PRECOMPUTE_COMMIT=1`: let that diagnostic run
62 /// write to the paged cache (its `commit` argument).
63 pub precompute_commit: bool,
64
65 /// 2026-09-25: `METRALE_DFLASH_PROPOSE_WARMUP_N=<n>`: eager passes per block
66 /// width before graph capture. Unset or unparseable is 2.
67 pub propose_warmup_n: usize,
68
69 /// 2026-09-25: The paged drafter cache for ctx K/V ("Option B"). On unless
70 /// `METRALE_DFLASH_OPTION_B=0`; resolved by `option_b_from` in `dflash_head.rs`.
71 pub option_b: bool,
72 /// 2026-09-25: `METRALE_DFLASH_OPTION_B_NO_CTX=1`: pass `ctx_count = 0`, so
73 /// paged attention sees only the γ K/V written in-layer.
74 pub option_b_no_ctx: bool,
75 /// 2026-09-25: The DFlash2 conv and selector path. Runs when the checkpoint
76 /// carries its components, unless `METRALE_DFLASH2=0`.
77 pub dflash2: bool,
78 /// 2026-09-25: `METRALE_DFLASH_BATCH_PROPOSE=<width>` caps the cross-sequence
79 /// batch. Unset is `usize::MAX` (as wide as the scratch bands allow); `0` or
80 /// `1` makes `propose_batch_max` return 1.
81 pub batch_propose_width: usize,
82 /// 2026-09-25: `METRALE_DFLASH_DRAFT_CAP=<n>`: return at most n drafts per
83 /// propose. `None` caps at the block width in flight (`block_g`).
84 pub draft_cap: Option<usize>,
85
86 /// 2026-09-25: `METRALE_DFLASH_VERIFY_TRACE=1`: log every propose's drafts
87 /// before the row-0 drop and the cap.
88 pub verify_trace: bool,
89 /// 2026-09-25: `METRALE_DFLASH_PRECOMPUTE_DUMP=1`: one-shot `/tmp` dumps of
90 /// the ctx K/V precompute intermediates.
91 pub precompute_dump: bool,
92 /// 2026-09-25: `METRALE_DFLASH_CTX_PARITY_DUMP=1`: one-shot `/tmp` dump of the
93 /// accumulated ctx hidden rows.
94 pub ctx_parity_dump: bool,
95 /// 2026-09-25: `METRALE_DFLASH_DEBUG_NO_DECODE_APPEND=1`: skip the ctx append
96 /// at the start of `propose_drafts`.
97 pub no_decode_append: bool,
98 /// 2026-09-25: `METRALE_DFLASH_DEBUG_FULL_PRECOMPUTE=1`: the non-batched
99 /// propose recomputes the whole ctx prefix (`committed = 0`) each step.
100 pub full_precompute: bool,
101 /// 2026-09-25: `METRALE_DFLASH_CTXLEN_PROBE=1`: warn when `ctx_positions` is
102 /// not strictly increasing, and log `ctx_len` against the position when the
103 /// position is a multiple of 16.
104 pub ctxlen_probe: bool,
105
106 /// 2026-09-25: The sequential Markov fixup. Runs when the drafter carries
107 /// the head, unless `METRALE_DSPARK_MARKOV=0`, which takes the plain
108 /// per-row argmax loop.
109 pub dspark_markov: bool,
110 /// 2026-09-25: `METRALE_DSPARK_CONF_TAU=<t>`: sigmoid-space acceptance
111 /// threshold for the confidence head. Unset or unparseable is `0.0`, which
112 /// leaves the head off.
113 pub conf_tau: f32,
114 /// 2026-09-25: `METRALE_DSPARK_SHIFT=1|0` forces the shifted-row convention
115 /// on or off; any other value (`None`) defers to the drafter config.
116 pub dspark_shift: Option<bool>,
117 /// 2026-09-25: Row 0 carries the Markov bias. On unless
118 /// `METRALE_DSPARK_ANCHOR_BIAS=0`. Confidence truncation also requires it:
119 /// an unbiased row never writes its confidence slot.
120 pub dspark_anchor_bias: bool,
121 /// 2026-09-25: `METRALE_DSPARK_CONF_TRACE=1`: log the confidence logits and
122 /// sigmoids.
123 pub dspark_conf_trace: bool,
124}
125
126/// 2026-09-25: Variables whose presence, at any value, sets `any_diagnostic_armed`
127/// and so keeps `forward_block` off CUDA graphs.
128const GRAPH_SUPPRESSING_DIAGNOSTICS: [&str; 11] = [
129 "METRALE_DFLASH_PROPOSE_NO_GRAPH",
130 "METRALE_DFLASH_DEBUG_DUMP_FULL",
131 "METRALE_DFLASH_OPTION_B_DIAG",
132 "METRALE_DFLASH_PRECOMPUTE_DUMP",
133 "METRALE_DFLASH_VERIFY_TRACE",
134 "METRALE_DFLASH_LOG_DRAFTS",
135 "METRALE_DFLASH_DEBUG_FORCE_PATTERN",
136 "METRALE_DFLASH_DEBUG_FORCE_NOISE_PATTERN",
137 "METRALE_DFLASH_DEBUG_CTX_OFF",
138 "METRALE_DFLASH_DEBUG_CTX_USED",
139 "METRALE_DFLASH_BLOCK_DUMP",
140];
141
142fn from_values(
143 mut value: impl FnMut(&str) -> Option<String>,
144 mut present: impl FnMut(&str) -> bool,
145) -> DFlashLevers {
146 fn opt_in(value: Option<&str>) -> bool {
147 value == Some("1")
148 }
149 fn parsed<T: std::str::FromStr>(value: Option<&str>) -> Option<T> {
150 value.and_then(|v| v.parse().ok())
151 }
152
153 DFlashLevers {
154 any_diagnostic_armed: GRAPH_SUPPRESSING_DIAGNOSTICS.iter().any(|var| present(var)),
155
156 debug_dump: opt_in(value("METRALE_DFLASH_DEBUG_DUMP").as_deref()),
157 debug_dump_full: opt_in(value("METRALE_DFLASH_DEBUG_DUMP_FULL").as_deref()),
158 log_drafts: opt_in(value("METRALE_DFLASH_LOG_DRAFTS").as_deref()),
159 block_dump: opt_in(value("METRALE_DFLASH_BLOCK_DUMP").as_deref()),
160 block_dump_at_pos: parsed(value("METRALE_DFLASH_BLOCK_DUMP_AT_POS").as_deref())
161 .unwrap_or(0),
162 option_b_diag: opt_in(value("METRALE_DFLASH_OPTION_B_DIAG").as_deref()),
163
164 force_pattern: opt_in(value("METRALE_DFLASH_DEBUG_FORCE_PATTERN").as_deref()),
165 force_noise_pattern: opt_in(value("METRALE_DFLASH_DEBUG_FORCE_NOISE_PATTERN").as_deref()),
166 force_no_ctx: opt_in(value("METRALE_DFLASH_DEBUG_CTX_OFF").as_deref()),
167 force_ctx_used: parsed(value("METRALE_DFLASH_DEBUG_CTX_USED").as_deref()),
168
169 precompute: opt_in(value("METRALE_DFLASH_PRECOMPUTE").as_deref()),
170 precompute_commit: opt_in(value("METRALE_DFLASH_PRECOMPUTE_COMMIT").as_deref()),
171
172 option_b: super::option_b_from(value("METRALE_DFLASH_OPTION_B").as_deref()),
173 option_b_no_ctx: opt_in(value("METRALE_DFLASH_OPTION_B_NO_CTX").as_deref()),
174 dflash2: value("METRALE_DFLASH2").as_deref() != Some("0"),
175 batch_propose_width: parsed(value("METRALE_DFLASH_BATCH_PROPOSE").as_deref())
176 .unwrap_or(usize::MAX),
177 draft_cap: parsed(value("METRALE_DFLASH_DRAFT_CAP").as_deref()),
178
179 verify_trace: opt_in(value("METRALE_DFLASH_VERIFY_TRACE").as_deref()),
180 precompute_dump: opt_in(value("METRALE_DFLASH_PRECOMPUTE_DUMP").as_deref()),
181 ctx_parity_dump: opt_in(value("METRALE_DFLASH_CTX_PARITY_DUMP").as_deref()),
182 no_decode_append: opt_in(value("METRALE_DFLASH_DEBUG_NO_DECODE_APPEND").as_deref()),
183 full_precompute: opt_in(value("METRALE_DFLASH_DEBUG_FULL_PRECOMPUTE").as_deref()),
184 ctxlen_probe: opt_in(value("METRALE_DFLASH_CTXLEN_PROBE").as_deref()),
185
186 dspark_markov: value("METRALE_DSPARK_MARKOV").as_deref() != Some("0"),
187 conf_tau: parsed(value("METRALE_DSPARK_CONF_TAU").as_deref()).unwrap_or(0.0),
188
189 propose_warmup_n: parsed(value("METRALE_DFLASH_PROPOSE_WARMUP_N").as_deref()).unwrap_or(2),
190
191 dspark_shift: match value("METRALE_DSPARK_SHIFT").as_deref() {
192 Some("1") => Some(true),
193 Some("0") => Some(false),
194 _ => None,
195 },
196 dspark_anchor_bias: value("METRALE_DSPARK_ANCHOR_BIAS").as_deref() != Some("0"),
197 dspark_conf_trace: opt_in(value("METRALE_DSPARK_CONF_TRACE").as_deref()),
198 }
199}
200
201impl DFlashLevers {
202 /// 2026-09-25: Resolve from the environment. The head calls it once, in
203 /// `from_weights.rs`; per-step code reads the head's `levers` field.
204 pub fn from_env() -> Self {
205 from_values(metrale_config::levers::var, |var| {
206 metrale_config::levers::var_os(var).is_some()
207 })
208 }
209
210 /// 2026-09-25: What `from_env` resolves to with none of these variables set
211 /// (checked by `nothing_set_resolves_to_defaults`).
212 pub fn defaults() -> Self {
213 Self {
214 // 2026-09-25: Every field whose unset value is not its type's
215 // `Default` is spelled out here.
216 option_b: true,
217 dflash2: true,
218 dspark_markov: true,
219 batch_propose_width: usize::MAX,
220 dspark_anchor_bias: true,
221 propose_warmup_n: 2,
222 ..Self::default()
223 }
224 }
225
226 /// 2026-09-25: The block dump is armed for this decode position.
227 pub fn block_dump_armed_at(&self, position: usize) -> bool {
228 self.block_dump && position >= self.block_dump_at_pos
229 }
230}
231
232#[cfg(test)]
233#[path = "levers_tests.rs"]
234mod tests;