spark_model/layers/ops/
model_levers.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Model-side kernel-path levers, resolved once and then carried.
4//!
5//! The second of the two lever categories on [`crate::layer::ForwardContext`]:
6//!
7//! * [`super::GemmDispatch`] — which GEMM implementation each projection takes.
8//! * [`ModelLevers`] — everything else the model's kernel paths branch on:
9//!   the SSM/GDN recurrence variant, FFN routing, MoE quantization, LoRA
10//!   application mode, diagnostics.
11//!
12//! Both were `OnceLock<bool>` statics reading `ATLAS_*` at first touch. Two
13//! problems with that, and only the first is about hot-swap:
14//!
15//! 1. A static outlives the model whose flags it encodes. Load a second model
16//!    whose recipe sets different levers and the process keeps taking the
17//!    previous model's branches — silently, because a cached `bool` cannot
18//!    report that it is stale.
19//! 2. It hides the dependency. A function that reads the environment through a
20//!    static declares nothing in its signature, cannot be exercised with a
21//!    different configuration without mutating the process, and gives the
22//!    compiler nothing to check.
23//!
24//! Carrying it fixes both, and a site that forgets the field fails to build.
25
26/// Kernel-path levers for one loaded model.
27///
28/// Plain `Copy` data resolved from the environment at model construction. Group
29/// membership follows the subsystem the lever steers, so a reader can see at a
30/// glance which part of the forward pass a flag reaches.
31#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
32pub struct ModelLevers {
33    // ── SSM / GDN recurrence ──
34    /// Keep GDN recurrent state in registers across the prefill chunk loop.
35    /// Default ON (the fold that shipped in PR #369, −7.25 % wall); the env var
36    /// is an opt-OUT, which is why the field is stored positively and the
37    /// resolution inverts it.
38    pub gdn_regresident: bool,
39    /// Batched FLA path for multi-sequence GDN decode.
40    pub gdn_batched_fla: bool,
41    /// WY17 GDN recurrence variant. Ships ON; `ATLAS_GDN_WY17=0` opts out.
42    pub gdn_wy17: bool,
43    /// WY-N GDN recurrence variant. Ships ON; `ATLAS_GDN_WYN=0` opts out.
44    pub gdn_wyn: bool,
45
46    // ── FFN / MoE ──
47    /// Lossless single-warp decode GEMV (`w4a16_gemv_sw`, `w4a16_gemv_dual_sw`).
48    /// Ships ON; `ATLAS_NO_GEMV_SW=1` restores the 64-thread kernels.
49    pub gemv_sw: bool,
50    /// Route decode FFN through the tile GEMM rather than the scalar GEMV.
51    pub decode_ffn_via_gemm: bool,
52    /// Small-M FFN GEMM tile shape. Ships ON; `ATLAS_FFN_SMALLM=0` opts out.
53    pub ffn_small_m: bool,
54    /// FP4 holo layout for the MoE down projection.
55    pub holo_moe_down_fp4: bool,
56    /// FP4 holo layout for the MoE gate/up projections.
57    pub holo_moe_gateup_fp4: bool,
58    /// Collect per-layer MoE expert-union statistics. Diagnostic.
59    pub moe_union_stats: bool,
60
61    // ── Attention ──
62    /// Contiguous-attention path for the DFlash head.
63    pub dflash_contig_attn: bool,
64
65    // ── LoRA ──
66    /// Apply LoRA eagerly at load instead of at each forward.
67    pub lora_eager: bool,
68    /// Allow hot rotation of LoRA adapters.
69    pub lora_rotate: bool,
70
71    // ── Diagnostics ──
72    /// K=4 chain-widening diagnostics.
73    pub k4_diag: bool,
74    /// Per-layer hidden-state norm dumps on the Gemma-4 decode path. Heavy —
75    /// one device-to-host copy per layer.
76    pub gemma4_diag: bool,
77
78    // ── Attention (cont.) ──
79    /// BF16 tensor-core attention projections: dequant FP4 to BF16 and use a
80    /// BF16 MMA instead of the default path, which crushes activations to FP8
81    /// E4M3. Removes the FP8 prefill perturbation on those projections.
82    pub bf16_tc_proj: bool,
83    /// Configured max decode batch (`--max-batch-size`), the reference count
84    /// the split-K attention split count is pinned to. Not from the
85    /// environment: `TransformerModel::new` writes it from the serve arg.
86    ///
87    /// It pins DETERMINISM — the online-softmax split-merge is
88    /// non-associative, so a sequence decoded alone must see the same
89    /// reduction tree as one co-batched with fifteen others. Held in a
90    /// `OnceLock` it was also idempotent, so a second model with a different
91    /// max batch would silently keep the first model's split count.
92    pub max_decode_seqs: u32,
93    /// `ATLAS_MTP_SHADOW_TOPK=k` (0 = off, clamped to 8): the drafter D2Hs
94    /// its logits and logs the top-k candidates. Observational only.
95    pub shadow_topk: usize,
96    /// `ATLAS_KV_POISON=1` — fill a fresh KV block with NaN instead of zero,
97    /// the discriminator for the "unwritten fresh tail block read"
98    /// hypothesis. A diagnostic that changes what the kernels READ, so it
99    /// must not leak across a swap.
100    pub kv_poison: bool,
101    /// MTP drafter context policy (`ATLAS_NO_DRAFTER_CONTEXT` /
102    /// `ATLAS_DRAFTER_PREFILL_ONLY`), resolved and logged once per model.
103    /// The two halves are coupled — prefill without carry is a measured
104    /// −927 ms/turn loss — so they travel as one value.
105    pub drafter: crate::model::drafter_context::DrafterContext,
106}
107
108fn from_values(
109    mut value: impl FnMut(&str) -> Option<String>,
110    mut present: impl FnMut(&str) -> bool,
111    shadow_topk: usize,
112    drafter: crate::model::drafter_context::DrafterContext,
113) -> ModelLevers {
114    fn opt_in(value: Option<&str>) -> bool {
115        value == Some("1")
116    }
117    fn opt_out(value: Option<&str>) -> bool {
118        value != Some("0")
119    }
120    fn opt_in_truthy(value: Option<&str>) -> bool {
121        value.is_some_and(|value| value == "1" || value.eq_ignore_ascii_case("true"))
122    }
123
124    ModelLevers {
125        max_decode_seqs: 1,
126        shadow_topk,
127        kv_poison: opt_in(value("ATLAS_KV_POISON").as_deref()),
128        drafter,
129        gdn_regresident: value("ATLAS_NO_GDN_REGRESIDENT").as_deref() != Some("1"),
130        gdn_batched_fla: opt_in(value("ATLAS_GDN_BATCHED_FLA").as_deref()),
131        gdn_wy17: opt_out(value("ATLAS_GDN_WY17").as_deref()),
132        gdn_wyn: opt_out(value("ATLAS_GDN_WYN").as_deref()),
133        ffn_small_m: opt_out(value("ATLAS_FFN_SMALLM").as_deref()),
134        gemv_sw: super::gemv_sw::gemv_sw_from(value("ATLAS_NO_GEMV_SW").as_deref()),
135        decode_ffn_via_gemm: opt_in(value("ATLAS_DECODE_FFN_VIA_GEMM").as_deref()),
136        holo_moe_down_fp4: opt_in_truthy(value("ATLAS_HOLO_MOE_DOWN_FP4").as_deref()),
137        holo_moe_gateup_fp4: opt_in_truthy(value("ATLAS_HOLO_MOE_GATEUP_FP4").as_deref()),
138        moe_union_stats: opt_in(value("ATLAS_MOE_UNION_STATS").as_deref()),
139        dflash_contig_attn: opt_in(value("ATLAS_DFLASH_CONTIG_ATTN").as_deref()),
140        lora_eager: opt_in_truthy(value("ATLAS_LORA_EAGER").as_deref()),
141        lora_rotate: opt_in_truthy(value("ATLAS_LORA_ROTATE").as_deref()),
142        k4_diag: opt_in(value("ATLAS_K4_DIAG").as_deref()),
143        gemma4_diag: opt_in_truthy(value("ATLAS_DIAG_GEMMA4").as_deref()),
144        bf16_tc_proj: present("ATLAS_BF16_TC_PROJ"),
145    }
146}
147
148impl ModelLevers {
149    /// Resolve from the environment. Called once, when the model is built.
150    pub fn from_env() -> Self {
151        from_values(
152            |var| std::env::var(var).ok(),
153            |var| std::env::var_os(var).is_some(),
154            crate::speculative::shadow_topk(),
155            crate::model::drafter_context::resolve_from_env(),
156        )
157    }
158
159    /// What a build resolves to with no `ATLAS_*` set — every opt-in off, the
160    /// one opt-out lever on. Tests construct a context with this instead of
161    /// mutating the process environment.
162    pub fn defaults() -> Self {
163        Self {
164            max_decode_seqs: 1,
165            shadow_topk: 0,
166            kv_poison: false,
167            drafter: crate::model::drafter_context::DrafterContext::BOTH,
168            gdn_regresident: true,
169            gdn_wy17: true,
170            gdn_wyn: true,
171            ffn_small_m: true,
172            gemv_sw: true,
173            ..Self::default()
174        }
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181    use std::collections::HashMap;
182
183    fn resolve(values: &[(&str, &str)]) -> ModelLevers {
184        let values: HashMap<_, _> = values.iter().copied().collect();
185        from_values(
186            |name| values.get(name).map(|value| (*value).to_owned()),
187            |name| values.contains_key(name),
188            0,
189            crate::model::drafter_context::DrafterContext::BOTH,
190        )
191    }
192
193    #[test]
194    fn the_opt_out_lever_is_on_by_default_and_every_opt_in_is_off() {
195        let d = ModelLevers::defaults();
196        assert_eq!(
197            resolve(&[]),
198            d,
199            "absent environment uses the public default"
200        );
201        assert_eq!(
202            d,
203            ModelLevers {
204                gdn_regresident: true,
205                gdn_wy17: true,
206                gdn_wyn: true,
207                gemv_sw: true,
208                ffn_small_m: true,
209                max_decode_seqs: 1,
210                drafter: crate::model::drafter_context::DrafterContext::BOTH,
211                ..ModelLevers::default()
212            }
213        );
214    }
215
216    #[test]
217    fn exact_one_opt_ins_map_to_their_own_fields() {
218        let cases = [
219            ("ATLAS_KV_POISON", [true, false, false, false, false, false]),
220            (
221                "ATLAS_GDN_BATCHED_FLA",
222                [false, true, false, false, false, false],
223            ),
224            (
225                "ATLAS_DECODE_FFN_VIA_GEMM",
226                [false, false, true, false, false, false],
227            ),
228            (
229                "ATLAS_MOE_UNION_STATS",
230                [false, false, false, true, false, false],
231            ),
232            (
233                "ATLAS_DFLASH_CONTIG_ATTN",
234                [false, false, false, false, true, false],
235            ),
236            ("ATLAS_K4_DIAG", [false, false, false, false, false, true]),
237        ];
238        for (name, expected) in cases {
239            let d = resolve(&[(name, "1")]);
240            assert_eq!(
241                [
242                    d.kv_poison,
243                    d.gdn_batched_fla,
244                    d.decode_ffn_via_gemm,
245                    d.moe_union_stats,
246                    d.dflash_contig_attn,
247                    d.k4_diag
248                ],
249                expected,
250                "{name}"
251            );
252        }
253        assert!(!resolve(&[("ATLAS_K4_DIAG", "true")]).k4_diag);
254    }
255
256    #[test]
257    fn truthy_opt_ins_map_independently_and_presence_is_distinct() {
258        let cases = [
259            (
260                "ATLAS_HOLO_MOE_DOWN_FP4",
261                [true, false, false, false, false],
262            ),
263            (
264                "ATLAS_HOLO_MOE_GATEUP_FP4",
265                [false, true, false, false, false],
266            ),
267            ("ATLAS_LORA_EAGER", [false, false, true, false, false]),
268            ("ATLAS_LORA_ROTATE", [false, false, false, true, false]),
269            ("ATLAS_DIAG_GEMMA4", [false, false, false, false, true]),
270        ];
271        for (name, expected) in cases {
272            let d = resolve(&[(name, "TrUe")]);
273            assert_eq!(
274                [
275                    d.holo_moe_down_fp4,
276                    d.holo_moe_gateup_fp4,
277                    d.lora_eager,
278                    d.lora_rotate,
279                    d.gemma4_diag
280                ],
281                expected,
282                "{name}"
283            );
284        }
285        assert!(resolve(&[("ATLAS_BF16_TC_PROJ", "0")]).bf16_tc_proj);
286    }
287
288    #[test]
289    fn kill_switches_and_zero_opt_outs_keep_their_distinct_polarities() {
290        let d = resolve(&[
291            ("ATLAS_NO_GDN_REGRESIDENT", "1"),
292            ("ATLAS_NO_GEMV_SW", "1"),
293            ("ATLAS_GDN_WY17", "0"),
294            ("ATLAS_GDN_WYN", "0"),
295            ("ATLAS_FFN_SMALLM", "0"),
296        ]);
297        assert!(!d.gdn_regresident);
298        assert!(!d.gemv_sw);
299        assert!(!d.gdn_wy17);
300        assert!(!d.gdn_wyn);
301        assert!(!d.ffn_small_m);
302        assert!(resolve(&[("ATLAS_NO_GDN_REGRESIDENT", "0")]).gdn_regresident);
303        assert!(resolve(&[("ATLAS_GDN_WY17", "1")]).gdn_wy17);
304    }
305
306    #[test]
307    fn externally_resolved_shadow_and_drafter_values_are_carried() {
308        let d = from_values(
309            |_| None,
310            |_| false,
311            7,
312            crate::model::drafter_context::DrafterContext::OFF,
313        );
314        assert_eq!(d.shadow_topk, 7);
315        assert_eq!(
316            d.drafter,
317            crate::model::drafter_context::DrafterContext::OFF
318        );
319    }
320}
spark_model/layers/ops/model_levers.rs

spark_model/layers/ops/
model_levers.rs
