spark_model/layers/ops/
model_levers_resolve.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! How [`ModelLevers`] is READ — the resolution table and the three
4//! constructors, split from what the levers ARE.
5//!
6//! Split out of `model_levers.rs` at 483 lines to stay under the repository's
7//! 500-LoC cap. The seam is deliberate rather than arbitrary: the parent file
8//! is now the DECLARATION — one documented field per lever, which is what a
9//! reader looking for "what does this flag do" wants — and this file is the
10//! single place that touches the environment, which is what a reader looking
11//! for "how is it spelled" wants. `from_values` stays pure over two closures
12//! so tests drive the production resolution instead of a copy of it.
13
14use super::ModelLevers;
15use crate::layers::ops::gemv_sw;
16
17pub(super) fn from_values(
18    mut value: impl FnMut(&str) -> Option<String>,
19    mut present: impl FnMut(&str) -> bool,
20    shadow_topk: usize,
21    drafter: crate::model::drafter_context::DrafterContext,
22    // ★ Passed IN, like `shadow_topk` and `drafter`, NOT read here. Calling
23    // `speculative::draft_conf_tau()` from inside this function broke its
24    // purity: it reads the real environment whatever the closures say, so a
25    // sibling test that set the variable made `resolve(&[])` return 0.99 and
26    // `the_opt_out_lever_is_on_by_default_and_every_opt_in_is_off` failed
27    // under parallel test execution. `from_values` is pure over its inputs;
28    // that is the property the whole test suite rests on.
29    draft_conf_tau: f32,
30    // The compiled target's `[defaults] decode_split_silu`. Passed IN for the
31    // same purity reason as `draft_conf_tau` above: `target_defaults::resolved`
32    // reads the real environment whatever the closures say.
33    default_split_silu: bool,
34) -> ModelLevers {
35    fn opt_in(value: Option<&str>) -> bool {
36        value == Some("1")
37    }
38    fn opt_out(value: Option<&str>) -> bool {
39        value != Some("0")
40    }
41    fn opt_in_truthy(value: Option<&str>) -> bool {
42        value.is_some_and(|value| value == "1" || value.eq_ignore_ascii_case("true"))
43    }
44    /// `"1"` or `"true"` EXACTLY — case-SENSITIVE, unlike [`opt_in_truthy`].
45    ///
46    /// The two decode-graph levers spelled it `is_ok_and(|v| v == "1" || v ==
47    /// "true")`, and widening them to accept `TRUE` would arm an experimental
48    /// CUDA-graph capture on a spelling that previously did nothing — the
49    /// direction that turns capture ON unexpectedly. Preserved rather than
50    /// unified; the divergence between the two helpers is the point.
51    fn opt_in_truthy_exact(value: Option<&str>) -> bool {
52        matches!(value, Some("1") | Some("true"))
53    }
54
55    ModelLevers {
56        max_decode_seqs: 1,
57        shadow_topk,
58        kv_poison: opt_in(value("ATLAS_KV_POISON").as_deref()),
59        drafter,
60        gdn_regresident: value("ATLAS_NO_GDN_REGRESIDENT").as_deref() != Some("1"),
61        gdn_batched_fla: opt_in(value("ATLAS_GDN_BATCHED_FLA").as_deref()),
62        gdn_wy17: opt_out(value("ATLAS_GDN_WY17").as_deref()),
63        gdn_wyn: opt_out(value("ATLAS_GDN_WYN").as_deref()),
64        ffn_small_m: opt_out(value("ATLAS_FFN_SMALLM").as_deref()),
65        gemv_sw: gemv_sw::gemv_sw_from(value("ATLAS_NO_GEMV_SW").as_deref()),
66        decode_ffn_via_gemm: opt_in(value("ATLAS_DECODE_FFN_VIA_GEMM").as_deref()),
67        holo_moe_down_fp4: opt_in_truthy(value("ATLAS_HOLO_MOE_DOWN_FP4").as_deref()),
68        holo_moe_gateup_fp4: opt_in_truthy(value("ATLAS_HOLO_MOE_GATEUP_FP4").as_deref()),
69        moe_union_stats: opt_in(value("ATLAS_MOE_UNION_STATS").as_deref()),
70        fp32_routing: opt_in(value("ATLAS_FP32_ROUTING").as_deref()),
71        fp32_gate: opt_in(value("ATLAS_FP32_GATE").as_deref()),
72        frankenstein_decode_via_prefill: opt_in(
73            value("ATLAS_FRANKENSTEIN_DECODE_VIA_PREFILL").as_deref(),
74        ),
75        k2_diag: opt_in(value("ATLAS_K2_DIAG").as_deref()),
76        dflash_debug_dump_full: opt_in(value("ATLAS_DFLASH_DEBUG_DUMP_FULL").as_deref()),
77        mtp_debug_norms: opt_in(value("ATLAS_MTP_DEBUG_NORMS").as_deref()),
78        draft_conf_tau,
79        // The compiled target declares this (`kernels/<hw>/HARDWARE.toml`
80        // `[defaults] decode_split_silu`); every current target declares it
81        // ON, which is the shipped default. `ATLAS_NO_DECODE_SPLIT_SILU` stays
82        // the PRESENCE kill switch, unchanged, and still wins. Resolved
83        // through `present` rather than `target_defaults::resolved()` so
84        // `from_values` stays pure over its closures — the property the whole
85        // test suite rests on. The declaration reaches it as the
86        // `default_split_silu` argument.
87        decode_split_silu: crate::layers::ops::target_defaults::resolve_toggle(
88            default_split_silu,
89            None,
90            present("ATLAS_NO_DECODE_SPLIT_SILU"),
91        )
92        .value,
93        bf16_tc_prefill: present("ATLAS_BF16_TC_PREFILL"),
94        fp8_m64_prefill: present("ATLAS_FP8_M64_PREFILL"),
95        int8_prefill: present("ATLAS_INT8_PREFILL"),
96        int8_faith5: present("ATLAS_INT8_FAITH5"),
97        ffn_nvfp4_mmq: !present("ATLAS_NO_FFN_NVFP4_MMQ"),
98        ffn_nvfp4_mmq_down: !present("ATLAS_NO_FFN_NVFP4_MMQ_DOWN"),
99        ffn_mmq: present("ATLAS_FFN_MMQ"),
100        ffn_mmq_down_q4k: present("ATLAS_FFN_MMQ_DOWN_Q4K"),
101        fp4_prefill: present("ATLAS_FP4_PREFILL"),
102        prefill_v2: !present("ATLAS_DISABLE_PREFILL_V2"),
103        moe_grouped_cutlass: opt_in(value("ATLAS_HOLO_MOE_GROUPED_CUTLASS").as_deref()),
104        moe_grouped_down: opt_in(value("ATLAS_HOLO_MOE_GROUPED_DOWN").as_deref()),
105        moe_prefill_exact_tiles: match value("ATLAS_MOE_PREFILL_EXACT_TILES").as_deref() {
106            Some("0") => Some(false),
107            Some("1") => Some(true),
108            _ => None,
109        },
110        moe_prefill_max_load_factor: value("ATLAS_MOE_PREFILL_MAX_LOAD_FACTOR")
111            .as_deref()
112            .and_then(|v| v.parse::<usize>().ok())
113            .filter(|&factor| factor > 0),
114        moe_prefill_zero: opt_in(value("ATLAS_MOE_PREFILL_ZERO").as_deref()),
115        moe_prefill_fp8_down: opt_in(value("ATLAS_MOE_PREFILL_FP8_DOWN").as_deref()),
116        ssm_w4a4: !present("ATLAS_NO_SSM_W4A4"),
117        ssd: !present("ATLAS_NO_SSD"),
118        ssm_persistent: !present("ATLAS_NO_SSM_PERSISTENT"),
119        moe_zero_intermediates: !present("ATLAS_MOE_NO_ZERO_INTERMEDIATES"),
120        moe_max_m_tiles_estimate: present("ATLAS_MOE_MAX_M_TILES_ESTIMATE"),
121        moe_w4a4: present("ATLAS_MOE_W4A4"),
122        shared_w4a4: !present("ATLAS_NO_SHARED_W4A4"),
123        shared_w4a4_down: present("ATLAS_SHARED_W4A4_DOWN"),
124        dflash_contig_attn: opt_in(value("ATLAS_DFLASH_CONTIG_ATTN").as_deref()),
125        lora_eager: opt_in_truthy(value("ATLAS_LORA_EAGER").as_deref()),
126        lora_rotate: opt_in_truthy(value("ATLAS_LORA_ROTATE").as_deref()),
127        k4_diag: opt_in(value("ATLAS_K4_DIAG").as_deref()),
128        gemma4_diag: opt_in_truthy(value("ATLAS_DIAG_GEMMA4").as_deref()),
129        mla_perseq_fallback: opt_in_truthy_exact(value("ATLAS_MLA_PERSEQ_FALLBACK").as_deref()),
130        hc_perseq_decode: opt_in(value("ATLAS_HC_PERSEQ_DECODE").as_deref()),
131        decode_batch_log: opt_in(value("ATLAS_DECODE_BATCH_LOG").as_deref()),
132        ms_profile: opt_in(value("ATLAS_MS_PROFILE").as_deref()),
133        conc_hsd: opt_in_truthy_exact(value("ATLAS_CONC_HSD").as_deref()),
134        ssm_save_dump: present("ATLAS_SSM_SAVE_DUMP"),
135        ep_graphs: opt_in_truthy_exact(value("ATLAS_EP_GRAPHS").as_deref()),
136        gdn_decode_graph: opt_in_truthy_exact(value("ATLAS_GDN_DECODE_GRAPH").as_deref()),
137        bf16_tc_proj: present("ATLAS_BF16_TC_PROJ"),
138        weight_pre_rotated: opt_in_truthy(value("TQ_PLUS_WEIGHT_ROTATION").as_deref()),
139        ssm_ms_profile: opt_in(value("ATLAS_SSM_MS_PROFILE").as_deref()),
140        ssm_detail_profile: opt_in(value("ATLAS_SSM_DETAIL_PROFILE").as_deref()),
141        ssm_gemv_batch4: opt_out(value("ATLAS_SSM_GEMV_BATCH4").as_deref()),
142        gdn_fused_conv: opt_in(value("ATLAS_GDN_FUSED_CONV").as_deref()),
143        moe_legacy_pertoken_decode: opt_in(value("ATLAS_MOE_LEGACY_PERTOKEN_DECODE").as_deref()),
144    }
145}
146
147impl ModelLevers {
148    /// The process-wide levers, resolved from the environment EXACTLY ONCE.
149    ///
150    /// ★ USE THIS, NOT [`Self::from_env`]. Every field here is a pure function
151    /// of `ATLAS_*` environment variables, which cannot change after start —
152    /// the runtime `set_var` that could have changed them was deliberately
153    /// removed. So this is a process constant and must be computed once.
154    ///
155    /// It was not. `from_env` reads ~30 environment variables, each allocating
156    /// a `String`, and three call sites invoked it from hot paths.
157    /// MEASURED: 32,513 resolutions in a single `concurrency-sweep` — which
158    /// matches 48 layers x ~680 prefills, i.e. once per layer per prefill from
159    /// `qwen3_attention::prefill_weights`. Each of those also re-ran
160    /// `drafter_context::resolve_from_env` and its logging.
161    ///
162    /// Returns a reference so callers cannot accidentally keep re-resolving;
163    /// `ModelLevers` is `Copy`, so `*ModelLevers::get()` is free when an owned
164    /// value is wanted.
165    pub fn get() -> &'static Self {
166        static LEVERS: std::sync::OnceLock<ModelLevers> = std::sync::OnceLock::new();
167        LEVERS.get_or_init(Self::from_env)
168    }
169
170    /// Resolve from the environment, unconditionally.
171    ///
172    /// Prefer [`Self::get`]. This exists for the one caller that needs an OWNED,
173    /// MUTABLE copy — the model build overwrites `max_decode_seqs` with the
174    /// batch size — and for tests that want a fresh read. Calling it in a hot
175    /// path re-reads every `ATLAS_*` variable.
176    pub fn from_env() -> Self {
177        from_values(
178            |var| std::env::var(var).ok(),
179            |var| std::env::var_os(var).is_some(),
180            crate::speculative::shadow_topk(),
181            crate::model::drafter_context::resolve_from_env(),
182            crate::speculative::draft_conf_tau(),
183            crate::layers::ops::target_defaults::declared().decode_split_silu,
184        )
185    }
186
187    /// What a build resolves to with no `ATLAS_*` set — every opt-in off, the
188    /// one opt-out lever on. Tests construct a context with this instead of
189    /// mutating the process environment.
190    pub fn defaults() -> Self {
191        Self {
192            max_decode_seqs: 1,
193            shadow_topk: 0,
194            kv_poison: false,
195            drafter: crate::model::drafter_context::DrafterContext::BOTH,
196            gdn_regresident: true,
197            gdn_wy17: true,
198            gdn_wyn: true,
199            ffn_small_m: true,
200            gemv_sw: true,
201            // Opt-out: ships ON, `ATLAS_SSM_GEMV_BATCH4=0` disables. Every
202            // opt-out lever must appear here or
203            // `the_opt_out_lever_is_on_by_default_and_every_opt_in_is_off`
204            // fails — which is exactly how this line came to be written.
205            ssm_gemv_batch4: true,
206            // The dense-FFN opt-outs. Each ships ON and is disabled by the
207            // PRESENCE of its variable, at any value — `=0` does not
208            // re-enable them, which is why they are listed here explicitly
209            // rather than left to `Default`.
210            decode_split_silu: true,
211            ffn_nvfp4_mmq: true,
212            ffn_nvfp4_mmq_down: true,
213            prefill_v2: true,
214            // The Nemotron prefill opt-outs, presence-gated like the four
215            // above: `=0` does NOT re-enable them.
216            ssm_w4a4: true,
217            ssd: true,
218            ssm_persistent: true,
219            moe_zero_intermediates: true,
220            shared_w4a4: true,
221            ..Self::default()
222        }
223    }
224}
spark_model/layers/ops/model_levers_resolve.rs

spark_model/layers/ops/
model_levers_resolve.rs
