metrale_model_layers/layers/ops/
model_levers_resolve.rs1use 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 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 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 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 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 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 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 ssm_gemv_batch4: true,
173 decode_split_silu: true,
176 ffn_nvfp4_mmq: true,
177 ffn_nvfp4_mmq_down: true,
178 prefill_v2: true,
179 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}