metrale_model_layers/layers/ops/model_levers.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Model-side kernel-path levers: one `Copy` struct, resolved from the environment and carried on `ForwardContext`.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - [`ModelLevers::get`] resolves the environment once per process (a `OnceLock` in
8//! `model_levers_resolve.rs`).
9//! - `from_values` reads nothing but its two closures and its arguments, so the tests drive the
10//! production resolution without touching the process environment.
11//! - Every field but `max_decode_seqs` comes from `from_values`. It sets `max_decode_seqs` to 1,
12//! and `TransformerModel::new` overwrites it with the configured max batch.
13//!
14//! Read a lever here once and pass the resolved field down; do not read `METRALE_*` in code that
15//! runs per token, per layer or per request. The `hot_path_env_guards` test fails when a guarded
16//! hot-path file reads the environment outside its allow list.
17//!
18//! [`super::GemmDispatch`], the other lever struct on [`crate::layer::ForwardContext`], chooses
19//! the GEMM implementation of each projection. [`ModelLevers`] holds every other kernel-path
20//! switch.
21
22/// 2026-09-25: Kernel-path levers for one loaded model. `TransformerModel::new` resolves its own
23/// copy with [`ModelLevers::from_env`]; other readers use the process-wide [`ModelLevers::get`].
24#[derive(Clone, Copy, Debug, PartialEq, Default)]
25pub struct ModelLevers {
26 /// 2026-09-25: On unless `METRALE_NO_GDN_REGRESIDENT=1`. Allows the register-resident GDN
27 /// prefill recurrence on the replay path when the K and V head dims are 128 and the kernel
28 /// resolved (`qwen3_ssm/trait_prefill_recur.rs`).
29 pub gdn_regresident: bool,
30 /// 2026-09-25: `METRALE_GDN_BATCHED_FLA=1`: the batched FLA kernels for multi-stream GDN
31 /// prefill when the K and V head dims are 128 (`qwen3_ssm/trait_prefill_gdn/batched.rs`).
32 pub gdn_batched_fla: bool,
33 /// 2026-09-25: On unless `METRALE_GDN_WY17=0`. Allows the WY17 kernel for a 17-token batched
34 /// GDN decode, except under the FP16 SSM state (`trait_decode_batched_conv_gdn.rs`).
35 pub gdn_wy17: bool,
36 /// 2026-09-25: On unless `METRALE_GDN_WYN=0`. Passed to `wyn_kernel` / `wyn_table_kernel`,
37 /// which choose the WY-N GDN decode kernels.
38 pub gdn_wyn: bool,
39
40 /// 2026-09-25: On unless `METRALE_NO_GEMV_SW=1` (`gemv_sw::gemv_sw_from`). Selects the
41 /// single-warp decode GEMVs (`w4a16_gemv_sw`, `w4a16_gemv_dual_sw`) where their kernel
42 /// resolved (`gemv_sw::use_gemv_sw`).
43 pub gemv_sw: bool,
44 /// 2026-09-25: `METRALE_DECODE_FFN_VIA_GEMM=1`: a SiLU dense-FFN decode runs gate, up and down
45 /// through `w4a16_prefill_gemm` at M=1 when the transposed `gate_proj_t` / `up_proj_t` weights
46 /// are present; without them it warns once and keeps the GEMV path (`dense_ffn.rs`).
47 pub decode_ffn_via_gemm: bool,
48 /// 2026-09-25: On unless `METRALE_FFN_SMALLM=0`. Allows the small-M GEMM tiles for
49 /// `m <= 64` with K a multiple of 32 (`dense_ffn.rs`, `multi_seq/qkv.rs`).
50 pub ffn_small_m: bool,
51 /// 2026-09-25: `METRALE_HOLO_MOE_DOWN_FP4=1` or `true` (any case). Only the weight loader's
52 /// warning reads this field (`weight_loader/qwen35/load_layers.rs`); `MoeLayer` reads the
53 /// variable itself into `down_fp4` (`moe/init.rs`).
54 pub holo_moe_down_fp4: bool,
55 /// 2026-09-25: `METRALE_HOLO_MOE_GATEUP_FP4=1` or `true` (any case). Only the weight loader's
56 /// warning reads this field; `MoeLayer` reads the variable itself into `gateup_fp4`.
57 pub holo_moe_gateup_fp4: bool,
58 /// 2026-09-25: `METRALE_MOE_UNION_STATS=1`: sample per-layer MoE expert-union statistics
59 /// into `ModelStats::moe_union` (`moe/union_stats.rs`).
60 pub moe_union_stats: bool,
61 /// 2026-09-25: `METRALE_FP32_ROUTING=1`: the MoE-input norm writes an FP32 router input and
62 /// the gate GEMM reads it at full precision. The last term of
63 /// `MoeLayer::fp32_routing_active`, whose other four terms are properties of the layer's
64 /// weights and kernels.
65 pub fp32_routing: bool,
66 /// 2026-09-25: `METRALE_FP32_GATE=1`: the batched gate writes FP32 logits from the BF16
67 /// router input (`moe/forward_batched_gate.rs`). Implied there by an active
68 /// [`Self::fp32_routing`].
69 pub fp32_gate: bool,
70 /// 2026-09-25: `METRALE_FRANKENSTEIN_DECODE_VIA_PREFILL=1`: DFlash capture layers
71 /// (`is_dflash_capture_layer`) run their MoE decode through `forward_prefill` with one row;
72 /// other layers are unchanged (`moe/forward.rs`). Diagnostic.
73 pub frankenstein_decode_via_prefill: bool,
74 /// 2026-09-25: `METRALE_K2_DIAG=1`: synchronising checkpoints in `forward_k2`, each labelled
75 /// with its stage, so the first failing stage names itself.
76 pub k2_diag: bool,
77
78 /// 2026-09-25: On when the compiled target's `[defaults] decode_split_silu` is on (a target
79 /// that does not declare it inherits on, `crates/kernels/build_defaults.rs`), unless
80 /// `METRALE_NO_DECODE_SPLIT_SILU` is present with any value. Decode runs `silu_mul` and then
81 /// a separate down GEMV instead of the fused SiLU+down kernel. An installed LoRA adapter takes
82 /// the split path whatever this says (`levers.decode_split_silu || self.lora.is_some()` in
83 /// `dense_ffn.rs`).
84 pub decode_split_silu: bool,
85 /// 2026-09-25: `METRALE_BF16_TC_PREFILL` (presence): the BF16 tensor-core dense-FFN prefill
86 /// GEMM. The call site honours it only when the BF16 kernel it selected (v2, or v1 without
87 /// [`Self::prefill_v2`]) is loaded.
88 pub bf16_tc_prefill: bool,
89 /// 2026-09-25: `METRALE_FP8_M64_PREFILL` (presence): the FP8 M64 dense-FFN prefill arm, when
90 /// `w4a16_gemm_t` is loaded.
91 pub fp8_m64_prefill: bool,
92 /// 2026-09-25: `METRALE_INT8_PREFILL` (presence): dense-FFN prefill through an int8 requant
93 /// and `int8_gemm_faith2`, when that kernel is loaded.
94 pub int8_prefill: bool,
95 /// 2026-09-25: `METRALE_INT8_FAITH5` (presence): the int8 arm launches the faith5 kernel
96 /// instead of faith2, when faith5 is loaded.
97 pub int8_faith5: bool,
98 /// 2026-09-25: On unless `METRALE_NO_FFN_NVFP4_MMQ` is present. NVFP4 W4A4 MMQ for the
99 /// gate/up prefill GEMMs of a SiLU dense FFN, when its kernels are loaded and no LoRA adapter
100 /// is installed (`dense_ffn.rs`).
101 pub ffn_nvfp4_mmq: bool,
102 /// 2026-09-25: On unless `METRALE_NO_FFN_NVFP4_MMQ_DOWN` is present. The same MMQ arm for the
103 /// down projection. It applies only when the gate/up arm ([`Self::ffn_nvfp4_mmq`]) is active,
104 /// but the two variables are independent.
105 pub ffn_nvfp4_mmq_down: bool,
106 /// 2026-09-25: `METRALE_FFN_MMQ` (presence): the Q4_K MMQ dense-FFN prefill arm, when its
107 /// kernels are loaded and the NVFP4 MMQ arm is not active.
108 pub ffn_mmq: bool,
109 /// 2026-09-25: `METRALE_FFN_MMQ_DOWN_Q4K` (presence): under [`Self::ffn_mmq`], keep the down
110 /// projection on Q4_K. Unset, down goes through the int8 faith2 path (`down_faith2` in
111 /// `dense_ffn.rs`, which reads `!levers.ffn_mmq_down_q4k`).
112 pub ffn_mmq_down_q4k: bool,
113 /// 2026-09-25: `METRALE_FP4_PREFILL` (presence): dense-FFN prefill through `w4a4_gemm`, with
114 /// the activations quantized to NVFP4 per GEMM, when both kernels are loaded.
115 pub fp4_prefill: bool,
116 /// 2026-09-25: On unless `METRALE_DISABLE_PREFILL_V2` is present. Prefer the v2 BF16 t_m128
117 /// prefill kernel over v1 when v2 is loaded.
118 pub prefill_v2: bool,
119
120 /// 2026-09-25: `METRALE_HOLO_MOE_GROUPED_CUTLASS=1`: the routed MoE prefill runs gate/up
121 /// through the single-launch CUTLASS grouped NVFP4 GEMM when the layer holds the load-time
122 /// host tables (`cutlass_grouped_host`, `moe/forward_prefill_routed.rs`).
123 pub moe_grouped_cutlass: bool,
124 /// 2026-09-25: `METRALE_HOLO_MOE_GROUPED_DOWN=1`: the down projection too. It also needs
125 /// [`Self::moe_grouped_cutlass`] and the host down tables.
126 pub moe_grouped_down: bool,
127 /// 2026-09-25: `METRALE_MOE_PREFILL_EXACT_TILES=1` or `=0`; any other value, or unset, is
128 /// `None`, which the call site resolves to on for NVFP4 experts and off otherwise
129 /// (`forward_prefill_routed.rs`). When on, the per-expert tile bound is read back from device
130 /// memory, so graph capture forces it off.
131 pub moe_prefill_exact_tiles: Option<bool>,
132 /// 2026-09-25: `METRALE_MOE_PREFILL_MAX_LOAD_FACTOR=<n>`: with exact tiles off, cap the
133 /// per-expert tile bound at n times the average rows per expert. Unset, unparseable or `0` is
134 /// `None`, the worst-case bound.
135 pub moe_prefill_max_load_factor: Option<usize>,
136 /// 2026-09-25: `METRALE_MOE_PREFILL_ZERO=1`: zero the routed-prefill gate/up/down scratch
137 /// before the grouped GEMMs. Expert parallelism (`ctx.comm`) zeroes it regardless.
138 pub moe_prefill_zero: bool,
139 /// 2026-09-25: `METRALE_MOE_PREFILL_FP8_DOWN=1`: the FP8 grouped GEMM for the routed-prefill
140 /// down projection, when `moe_fp8_grouped_gemm_t` and `bf16_to_fp8` are loaded.
141 pub moe_prefill_fp8_down: bool,
142
143 // 2026-09-25: The Nemotron prefill levers (`nemotron_mamba2/prefill.rs`,
144 // `nemotron_moe/prefill_sorted.rs`, `prefill_shared_up.rs`). All eight are presence-gated:
145 // any value, `0` included, counts as set.
146 /// 2026-09-25: On unless `METRALE_NO_SSM_W4A4` is present. The W4A4 native-FP4 GEMM for the
147 /// Mamba2 SSM prefill projections at n >= 512, when the kernels are loaded and the `fp8_act`
148 /// buffer is large enough.
149 pub ssm_w4a4: bool,
150 /// 2026-09-25: On unless `METRALE_NO_SSD` is present. The chunked SSD scan for Mamba2
151 /// prefill, when `ops::ssd_scan_fits(state_size)` also holds.
152 pub ssd: bool,
153 /// 2026-09-25: On unless `METRALE_NO_SSM_PERSISTENT` is present. The persistent SSM prefill
154 /// kernel, tried only when the SSD scan is not taken; otherwise the sequential scan runs.
155 pub ssm_persistent: bool,
156 /// 2026-09-25: On unless `METRALE_MOE_NO_ZERO_INTERMEDIATES` is present. Zero
157 /// `expert_up_out` and `expert_down_out` before the sorted-prefill grouped GEMMs.
158 pub moe_zero_intermediates: bool,
159 /// 2026-09-25: `METRALE_MOE_MAX_M_TILES_ESTIMATE` (presence): size the grouped GEMM's
160 /// `max_m_tiles` (its grid.y) from twice the average rows per expert instead of the worst
161 /// case, every routed row on one expert. Rows of an expert beyond that bound get no block and
162 /// are not computed.
163 pub moe_max_m_tiles_estimate: bool,
164 /// 2026-09-25: `METRALE_MOE_W4A4` (presence): the W4A4 grouped up projection for a latent
165 /// MoE input at n >= 512.
166 pub moe_w4a4: bool,
167 /// 2026-09-25: On unless `METRALE_NO_SHARED_W4A4` is present. W4A4 for the shared-expert up
168 /// projection at n >= 512, when no native FP8 shared-up weight is used.
169 pub shared_w4a4: bool,
170 /// 2026-09-25: `METRALE_SHARED_W4A4_DOWN` (presence): W4A4 for the shared-expert down
171 /// projection at n >= 512. Independent of [`Self::shared_w4a4`].
172 pub shared_w4a4_down: bool,
173
174 /// 2026-09-25: `METRALE_DFLASH_CONTIG_ATTN=1`: the DFlash head's paged block layer takes
175 /// `forward_block_layer_attention_contig`.
176 pub dflash_contig_attn: bool,
177
178 /// 2026-09-25: `METRALE_LORA_EAGER=1` or `true` (any case): decode and verify run without
179 /// CUDA graphs while a LoRA pool is loaded (`model/trait_impl/decode_a.rs`, `decode_a2.rs`,
180 /// `verify_b.rs`, `verify_c.rs`).
181 pub lora_eager: bool,
182 /// 2026-09-25: `METRALE_LORA_ROTATE=1` or `true` (any case): permits runtime adapter
183 /// rotation and swap (`TransformerModel::lora_rotatable`, which `METRALE_LORA_PEER` also
184 /// sets).
185 pub lora_rotate: bool,
186
187 /// 2026-09-25: `METRALE_K4_DIAG=1`: synchronise after each named phase of the batched GDN
188 /// decode so an illegal access is attributed to its phase; skipped under graph capture
189 /// (`qwen3_ssm/trait_decode_batched.rs`).
190 pub k4_diag: bool,
191 /// 2026-09-25: `METRALE_DIAG_GEMMA4=1` or `true` (any case): per-layer hidden-state norm dumps
192 /// on the Gemma-4 attention decode path; off under graph capture (`decode_inner.rs`).
193 pub gemma4_diag: bool,
194 /// 2026-09-25: `METRALE_DFLASH_DEBUG_DUMP_FULL=1`: the model-side half of the DFlash full
195 /// dump, which writes the sequence's tokens once per model (`model-engine` `impl_b3.rs`).
196 ///
197 /// The drafter half is `DFlashLevers::debug_dump_full` in metrale-model-arch's
198 /// `dflash_head/levers.rs`, which resolves the same variable: `TransformerModel` holds the
199 /// drafter only as a `dyn DraftProposer`. The test
200 /// `the_two_halves_of_the_dflash_dump_name_the_same_flag` checks that both resolutions name
201 /// one variable.
202 pub dflash_debug_dump_full: bool,
203 /// 2026-09-25: `METRALE_MTP_DEBUG_NORMS=1`: per-stage norm dumps in the MTP drafter's
204 /// `forward_one`.
205 pub mtp_debug_norms: bool,
206 /// 2026-09-25: `METRALE_MTP_CHAIN_POSTNORM=1`: draft `j > 0` of an MTP propose chain reads the
207 /// drafter's final-normed hidden (`norm_output`) instead of its pre-norm residual stream
208 /// (`hidden_states`). Both chains read it through `MtpHead::chain_hidden`.
209 pub mtp_chain_postnorm: bool,
210 /// 2026-09-25: `METRALE_MTP_TARGET_POSTNORM=1`: drafter rows built from a target hidden (the
211 /// first draft, drafter prefill, the exact-KV catch-up) go through the target's final norm
212 /// before `pre_fc_norm_hidden` (`MtpHead::target_postnorm_rows`). No effect when the head holds
213 /// no target final-norm weight.
214 pub mtp_target_postnorm: bool,
215 /// 2026-09-25: `METRALE_MTP_KV_EXACT=1`: after a verify, the MTP drafter keeps the KV row built
216 /// from the verified input token and trims only the chain rows (`draft_proposer.rs`); the
217 /// model then appends one catch-up row per accepted draft from the verify forward's stashed
218 /// hidden (`model/trait_impl/speculative_mtp.rs`).
219 pub mtp_kv_exact: bool,
220 /// 2026-09-25: `METRALE_FP8_MOE_GROUPED_DECODE=1`: the target model's multi-row FP8 MoE decode
221 /// runs through the cross-row grouped kernels (`FfnComponent::fp8_grouped_decode_ok`). The MTP
222 /// drafter calls `MoeLayer::fp8_grouped_decode_ok` directly and does not read this lever;
223 /// `METRALE_NO_FP8_MOE_GROUPED_DECODE` (presence) disables both
224 /// (`moe/forward_fp8_grouped_decode.rs`).
225 pub moe_fp8_grouped_decode_target: bool,
226 /// 2026-09-26: `METRALE_FP8_ATTN_M32=1`: attention layers load the 32-row tensor-core twin
227 /// `w8a16_gemm_pipelined_m32` at construction (`qwen3_attention/init_proj_kernels.rs`),
228 /// which batches of more than 16 rows can then take. Without it the handle is zero.
229 pub fp8_attn_m32: bool,
230 /// 2026-09-25: `METRALE_SSM_SAVE_DUMP` (presence): the scratch/SSM-state fingerprint probe on a
231 /// sequence's first decode step, which also runs that step without graphs
232 /// (`model/trait_impl/decode_a.rs`).
233 pub ssm_save_dump: bool,
234
235 // 2026-09-25: The batched-decode levers, read in `model/trait_impl/decode_a2.rs`.
236 // `mla_perseq_fallback` and `conc_hsd` accept `1` or `true` (case-sensitive); the other
237 // three accept only `1`.
238 /// 2026-09-25: `METRALE_MLA_PERSEQ_FALLBACK`: on an MLA model, decode the batch one sequence
239 /// at a time.
240 pub mla_perseq_fallback: bool,
241 /// 2026-09-25: `METRALE_HC_PERSEQ_DECODE=1`: on a model with `hc_mult > 0`, decode the batch
242 /// one sequence at a time, as an active QSA indexer already does. Decided before the
243 /// expert-parallel branch, so an EP batch takes the per-sequence path too.
244 pub hc_perseq_decode: bool,
245 /// 2026-09-25: `METRALE_DECODE_BATCH_LOG=1`: log the batch's SSM slot vector each step.
246 pub decode_batch_log: bool,
247 /// 2026-09-25: `METRALE_MS_PROFILE=1`: multi-sequence decode profiling. It disables the
248 /// batched decode graph so the profiling syncs are legal. A different variable from
249 /// [`Self::ssm_ms_profile`] (`METRALE_SSM_MS_PROFILE`).
250 pub ms_profile: bool,
251 /// 2026-09-25: `METRALE_CONC_HSD`: per-row hidden-state dumps after the embed and after each
252 /// layer, for batches of two or more rows without a communicator.
253 pub conc_hsd: bool,
254
255 /// 2026-09-25: `METRALE_EP_GRAPHS=1` or `true` (case-sensitive): allow decode CUDA-graph
256 /// capture when a communicator is present (`model/trait_impl/decode_a.rs`).
257 pub ep_graphs: bool,
258 /// 2026-09-25: `METRALE_GDN_DECODE_GRAPH=1` or `true` (case-sensitive): the same permission as
259 /// [`Self::ep_graphs`]; `decode_a.rs` ORs the two.
260 pub gdn_decode_graph: bool,
261
262 /// 2026-09-25: `METRALE_BF16_TC_PROJ` (presence): attention prefill projections use
263 /// `w4a16_gemm_n128_m128_bf16` when that kernel is loaded (`qwen3_attention/prefill_weights.rs`).
264 pub bf16_tc_proj: bool,
265 /// 2026-09-25: `TQ_PLUS_WEIGHT_ROTATION=1` or `true` (any case): the weight loader rotates the
266 /// Q/K/V projection weights at load (head_dim 128 only, `load_layers/attention_arms.rs`), and
267 /// the attention paths skip their runtime rotation (`if !weight_pre_rotated`, e.g.
268 /// `decode/write_kv_cache.rs`).
269 pub weight_pre_rotated: bool,
270
271 /// 2026-09-25: `METRALE_SSM_MS_PROFILE=1`: per-phase timing of the multi-sequence SSM decode;
272 /// off under graph capture (`qwen3_ssm/trait_decode_multi_seq.rs`).
273 pub ssm_ms_profile: bool,
274 /// 2026-09-25: `METRALE_SSM_DETAIL_PROFILE=1`: per-sub-step timing inside the batched SSM
275 /// recurrence; off under graph capture (`trait_decode_multi_seq/ssm_batched.rs`).
276 pub ssm_detail_profile: bool,
277 /// 2026-09-25: On unless `METRALE_SSM_GEMV_BATCH4=0`. For n <= 16 the SSM projections use the
278 /// contiguous-batch GEMV tier (batch4 for n <= 4, batch16 above) when its kernel resolved
279 /// (`trait_decode_multi_seq/ssm_batched_proj.rs`).
280 pub ssm_gemv_batch4: bool,
281 /// 2026-09-25: `METRALE_GDN_FUSED_CONV=1`: the fused conv+GDN+norm decode kernel when
282 /// `nv == 2 * nk`, the K and V head dims are 128 and the kernel resolved
283 /// (`trait_decode_multi_seq/ssm_batched_recurrent.rs`).
284 pub gdn_fused_conv: bool,
285 /// 2026-09-25: `METRALE_MOE_LEGACY_PERTOKEN_DECODE=1`: the multi-sequence MoE decode takes the
286 /// per-token path instead of the token-major one. The call site branches on
287 /// `!moe_legacy_pertoken_decode` (`qwen3_ssm/trait_decode_multi_seq.rs`).
288 pub moe_legacy_pertoken_decode: bool,
289 /// 2026-09-25: The configured max decode batch. Not from the environment: `from_values` sets
290 /// 1, and `TransformerModel::new` overwrites it with `max_batch_size` (at least 1). The paged
291 /// decode attention pins its split-K split count to it (`decode/attention_forward.rs`), so a
292 /// sequence's attention reduction does not depend on how many sequences share its batch. It
293 /// also bounds the verify WY-table stash (`model/trait_impl/gdn_woa.rs`).
294 pub max_decode_seqs: u32,
295 /// 2026-09-25: `METRALE_MTP_SHADOW_TOPK=k` (`speculative::shadow_topk`): unset or unparseable
296 /// is 0 (off), and it is capped at 8. The MTP drafter copies its logits to the host and logs
297 /// its top-k candidates (`mtp_head/forward.rs`); the server's `verify_k3_step.rs` and
298 /// `verify_k4_step.rs` also read it.
299 pub shadow_topk: usize,
300 /// 2026-09-25: `METRALE_KV_POISON=1`: fresh KV blocks are filled with `0xFF`
301 /// (`PagedKvCache::poison_block`) instead of zeroed (`model/block_mgmt.rs`). Diagnostic.
302 pub kv_poison: bool,
303 /// 2026-09-25: The MTP drafter context policy from `drafter_context::resolve_from_env`
304 /// (`METRALE_NO_MTP_DRAFTER_CONTEXT`, `METRALE_MTP_DRAFTER_CONTEXT_PREFILL_ONLY_UNSAFE`). One
305 /// value, because its `carry` implies its `prefill`.
306 pub drafter: crate::drafter_context::DrafterContext,
307}
308
309/// 2026-09-25: How the levers are read: `from_values` and the `get`, `from_env` and `defaults`
310/// constructors.
311#[path = "model_levers_resolve.rs"]
312mod resolve;
313
314#[cfg(test)]
315#[path = "model_levers_tests.rs"]
316mod tests;
317
318/// 2026-09-25: The source-level test of which hot-path files may read the environment. Declared
319/// here rather than in `tests` because it scans files outside this module.
320#[cfg(test)]
321#[path = "hot_path_env_guards.rs"]
322mod hot_path_env_guards;