metrale_model_layers/layers/ops/
model_levers_resolve.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: How [`ModelLevers`] is read: the resolution table (`from_values`) and the `get`, `from_env` and `defaults` constructors.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - `from_values` reads only its two closures and its arguments; it never touches the process
8//!   environment, so the tests drive this production table directly.
9//! - [`ModelLevers::get`] resolves once per process (`OnceLock`).
10//! - `resolve(&[])` in the tests equals [`ModelLevers::defaults`]
11//!   (`the_opt_out_lever_is_on_by_default_and_every_opt_in_is_off`).
12
13use super::ModelLevers;
14use crate::layers::ops::gemv_sw;
15
16pub(super) fn from_values(
17    mut value: impl FnMut(&str) -> Option<String>,
18    mut present: impl FnMut(&str) -> bool,
19    shadow_topk: usize,
20    drafter: crate::drafter_context::DrafterContext,
21    // 2026-09-25: The compiled target's `[defaults] decode_split_silu`
22    // (`target_defaults::declared()` in `from_env`), passed in so a test can choose it.
23    default_split_silu: bool,
24) -> ModelLevers {
25    fn opt_in(value: Option<&str>) -> bool {
26        value == Some("1")
27    }
28    fn opt_out(value: Option<&str>) -> bool {
29        value != Some("0")
30    }
31    fn opt_in_truthy(value: Option<&str>) -> bool {
32        value.is_some_and(|value| value == "1" || value.eq_ignore_ascii_case("true"))
33    }
34    /// 2026-09-25: `"1"` or `"true"` exactly: case-sensitive, unlike [`opt_in_truthy`], so `TRUE`
35    /// does not arm. Used by `ep_graphs`, `gdn_decode_graph`, `mla_perseq_fallback` and `conc_hsd`.
36    fn opt_in_truthy_exact(value: Option<&str>) -> bool {
37        matches!(value, Some("1") | Some("true"))
38    }
39
40    ModelLevers {
41        max_decode_seqs: 1,
42        shadow_topk,
43        kv_poison: opt_in(value("METRALE_KV_POISON").as_deref()),
44        drafter,
45        gdn_regresident: value("METRALE_NO_GDN_REGRESIDENT").as_deref() != Some("1"),
46        gdn_batched_fla: opt_in(value("METRALE_GDN_BATCHED_FLA").as_deref()),
47        gdn_wy17: opt_out(value("METRALE_GDN_WY17").as_deref()),
48        gdn_wyn: opt_out(value("METRALE_GDN_WYN").as_deref()),
49        ffn_small_m: opt_out(value("METRALE_FFN_SMALLM").as_deref()),
50        gemv_sw: gemv_sw::gemv_sw_from(value("METRALE_NO_GEMV_SW").as_deref()),
51        decode_ffn_via_gemm: opt_in(value("METRALE_DECODE_FFN_VIA_GEMM").as_deref()),
52        holo_moe_down_fp4: opt_in_truthy(value("METRALE_HOLO_MOE_DOWN_FP4").as_deref()),
53        holo_moe_gateup_fp4: opt_in_truthy(value("METRALE_HOLO_MOE_GATEUP_FP4").as_deref()),
54        moe_union_stats: opt_in(value("METRALE_MOE_UNION_STATS").as_deref()),
55        fp32_routing: opt_in(value("METRALE_FP32_ROUTING").as_deref()),
56        fp32_gate: opt_in(value("METRALE_FP32_GATE").as_deref()),
57        frankenstein_decode_via_prefill: opt_in(
58            value("METRALE_FRANKENSTEIN_DECODE_VIA_PREFILL").as_deref(),
59        ),
60        k2_diag: opt_in(value("METRALE_K2_DIAG").as_deref()),
61        dflash_debug_dump_full: opt_in(value("METRALE_DFLASH_DEBUG_DUMP_FULL").as_deref()),
62        mtp_debug_norms: opt_in(value("METRALE_MTP_DEBUG_NORMS").as_deref()),
63        mtp_chain_postnorm: opt_in(value("METRALE_MTP_CHAIN_POSTNORM").as_deref()),
64        mtp_target_postnorm: opt_in(value("METRALE_MTP_TARGET_POSTNORM").as_deref()),
65        mtp_kv_exact: opt_in(value("METRALE_MTP_KV_EXACT").as_deref()),
66        moe_fp8_grouped_decode_target: opt_in(value("METRALE_FP8_MOE_GROUPED_DECODE").as_deref()),
67        fp8_attn_m32: opt_in(value("METRALE_FP8_ATTN_M32").as_deref()),
68        // 2026-09-25: The target's declaration (`kernels/<hw>/HARDWARE.toml` `[defaults]
69        // decode_split_silu`; undeclared is on, `crates/kernels/build_defaults.rs`) arrives as
70        // `default_split_silu`. The presence of `METRALE_NO_DECODE_SPLIT_SILU`, any value,
71        // turns it off (`resolve_toggle`'s `legacy_off`).
72        decode_split_silu: crate::layers::ops::target_defaults::resolve_toggle(
73            default_split_silu,
74            None,
75            present("METRALE_NO_DECODE_SPLIT_SILU"),
76        )
77        .value,
78        bf16_tc_prefill: present("METRALE_BF16_TC_PREFILL"),
79        fp8_m64_prefill: present("METRALE_FP8_M64_PREFILL"),
80        int8_prefill: present("METRALE_INT8_PREFILL"),
81        int8_faith5: present("METRALE_INT8_FAITH5"),
82        ffn_nvfp4_mmq: !present("METRALE_NO_FFN_NVFP4_MMQ"),
83        ffn_nvfp4_mmq_down: !present("METRALE_NO_FFN_NVFP4_MMQ_DOWN"),
84        ffn_mmq: present("METRALE_FFN_MMQ"),
85        ffn_mmq_down_q4k: present("METRALE_FFN_MMQ_DOWN_Q4K"),
86        fp4_prefill: present("METRALE_FP4_PREFILL"),
87        prefill_v2: !present("METRALE_DISABLE_PREFILL_V2"),
88        moe_grouped_cutlass: opt_in(value("METRALE_HOLO_MOE_GROUPED_CUTLASS").as_deref()),
89        moe_grouped_down: opt_in(value("METRALE_HOLO_MOE_GROUPED_DOWN").as_deref()),
90        moe_prefill_exact_tiles: match value("METRALE_MOE_PREFILL_EXACT_TILES").as_deref() {
91            Some("0") => Some(false),
92            Some("1") => Some(true),
93            _ => None,
94        },
95        moe_prefill_max_load_factor: value("METRALE_MOE_PREFILL_MAX_LOAD_FACTOR")
96            .as_deref()
97            .and_then(|v| v.parse::<usize>().ok())
98            .filter(|&factor| factor > 0),
99        moe_prefill_zero: opt_in(value("METRALE_MOE_PREFILL_ZERO").as_deref()),
100        moe_prefill_fp8_down: opt_in(value("METRALE_MOE_PREFILL_FP8_DOWN").as_deref()),
101        ssm_w4a4: !present("METRALE_NO_SSM_W4A4"),
102        ssd: !present("METRALE_NO_SSD"),
103        ssm_persistent: !present("METRALE_NO_SSM_PERSISTENT"),
104        moe_zero_intermediates: !present("METRALE_MOE_NO_ZERO_INTERMEDIATES"),
105        moe_max_m_tiles_estimate: present("METRALE_MOE_MAX_M_TILES_ESTIMATE"),
106        moe_w4a4: present("METRALE_MOE_W4A4"),
107        shared_w4a4: !present("METRALE_NO_SHARED_W4A4"),
108        shared_w4a4_down: present("METRALE_SHARED_W4A4_DOWN"),
109        dflash_contig_attn: opt_in(value("METRALE_DFLASH_CONTIG_ATTN").as_deref()),
110        lora_eager: opt_in_truthy(value("METRALE_LORA_EAGER").as_deref()),
111        lora_rotate: opt_in_truthy(value("METRALE_LORA_ROTATE").as_deref()),
112        k4_diag: opt_in(value("METRALE_K4_DIAG").as_deref()),
113        gemma4_diag: opt_in_truthy(value("METRALE_DIAG_GEMMA4").as_deref()),
114        mla_perseq_fallback: opt_in_truthy_exact(value("METRALE_MLA_PERSEQ_FALLBACK").as_deref()),
115        hc_perseq_decode: opt_in(value("METRALE_HC_PERSEQ_DECODE").as_deref()),
116        decode_batch_log: opt_in(value("METRALE_DECODE_BATCH_LOG").as_deref()),
117        ms_profile: opt_in(value("METRALE_MS_PROFILE").as_deref()),
118        conc_hsd: opt_in_truthy_exact(value("METRALE_CONC_HSD").as_deref()),
119        ssm_save_dump: present("METRALE_SSM_SAVE_DUMP"),
120        ep_graphs: opt_in_truthy_exact(value("METRALE_EP_GRAPHS").as_deref()),
121        gdn_decode_graph: opt_in_truthy_exact(value("METRALE_GDN_DECODE_GRAPH").as_deref()),
122        bf16_tc_proj: present("METRALE_BF16_TC_PROJ"),
123        weight_pre_rotated: opt_in_truthy(value("TQ_PLUS_WEIGHT_ROTATION").as_deref()),
124        ssm_ms_profile: opt_in(value("METRALE_SSM_MS_PROFILE").as_deref()),
125        ssm_detail_profile: opt_in(value("METRALE_SSM_DETAIL_PROFILE").as_deref()),
126        ssm_gemv_batch4: opt_out(value("METRALE_SSM_GEMV_BATCH4").as_deref()),
127        gdn_fused_conv: opt_in(value("METRALE_GDN_FUSED_CONV").as_deref()),
128        moe_legacy_pertoken_decode: opt_in(value("METRALE_MOE_LEGACY_PERTOKEN_DECODE").as_deref()),
129    }
130}
131
132impl ModelLevers {
133    /// 2026-09-25: The process-wide levers, resolved from the environment on the first call and
134    /// cached for the life of the process. Prefer this to [`Self::from_env`] outside model
135    /// construction. `max_decode_seqs` is 1 here. `ModelLevers` is `Copy`, so
136    /// `*ModelLevers::get()` gives an owned value.
137    pub fn get() -> &'static Self {
138        static LEVERS: std::sync::OnceLock<ModelLevers> = std::sync::OnceLock::new();
139        LEVERS.get_or_init(Self::from_env)
140    }
141
142    /// 2026-09-25: Resolve from the environment on every call. `TransformerModel::new` uses it so
143    /// each model gets its own copy, whose `max_decode_seqs` it then sets. The
144    /// `from_env_is_called_only_where_it_is_allowed` test fails when another file calls it.
145    pub fn from_env() -> Self {
146        from_values(
147            metrale_config::levers::var,
148            |var| metrale_config::levers::var_os(var).is_some(),
149            crate::speculative::shadow_topk(),
150            crate::drafter_context::resolve_from_env(),
151            crate::layers::ops::target_defaults::declared().decode_split_silu,
152        )
153    }
154
155    /// 2026-09-25: What `from_values` returns with no variable set, a target that declares
156    /// `decode_split_silu` on, `DrafterContext::BOTH` and no shadow
157    /// top-k: every opt-in off, every opt-out on. Tests build a context with this instead of
158    /// mutating the environment.
159    pub fn defaults() -> Self {
160        Self {
161            max_decode_seqs: 1,
162            shadow_topk: 0,
163            kv_poison: false,
164            drafter: crate::drafter_context::DrafterContext::BOTH,
165            gdn_regresident: true,
166            gdn_wy17: true,
167            gdn_wyn: true,
168            ffn_small_m: true,
169            gemv_sw: true,
170            // 2026-09-25: On unless `METRALE_SSM_GEMV_BATCH4=0`. Every opt-out lever must be listed
171            // here or `the_opt_out_lever_is_on_by_default_and_every_opt_in_is_off` fails.
172            ssm_gemv_batch4: true,
173            // 2026-09-25: The dense-FFN opt-outs: on, and turned off by the presence of their
174            // variable at any value, `=0` included.
175            decode_split_silu: true,
176            ffn_nvfp4_mmq: true,
177            ffn_nvfp4_mmq_down: true,
178            prefill_v2: true,
179            // 2026-09-25: The Nemotron prefill opt-outs, presence-gated the same way.
180            ssm_w4a4: true,
181            ssd: true,
182            ssm_persistent: true,
183            moe_zero_intermediates: true,
184            shared_w4a4: true,
185            ..Self::default()
186        }
187    }
188}