metrale_model_layers/circuit_exec/
policy.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The live policy a plan is fused under: every setting the rules read, taken from
4//! the switch the legacy dispatch reads for it (the sources INSTANCES.toml cites), plus the
5//! switches no rule models, which the executor refuses when they are away from their default.
6//!
7//! Owner: model-layers circuit executor.
8//! Invariants:
9//! - Each setting is read from the one place legacy dispatch reads it; nothing here restates a
10//!   default.
11//! - A switch that changes legacy decode but that no rule reads is reported by
12//!   [`unmodelled_switches`], never silently ignored.
13//! - The environment switches of the multi-sequence dispatch ([`MULTI_SEQ_ENV_SWITCHES`]) are
14//!   refused when present at all, whatever their value: over-refusal, never a misrun. Their
15//!   readers parse them; this list only names them.
16
17use std::collections::BTreeMap;
18
19use anyhow::{Result, bail};
20use metrale_cache::kv_cache::KvCacheDtype;
21use metrale_circuit::Policy;
22
23use crate::layers::ops::ModelLevers;
24use crate::layers::{RowTiers, row_tiers};
25
26fn on_off(b: bool) -> String {
27    if b { "on" } else { "off" }.to_string()
28}
29
30/// 2026-09-28: The spelling the rules use for a KV-cache dtype; an error for one the circuit
31/// has no rule for.
32pub fn kv_dtype_name(d: KvCacheDtype) -> Result<&'static str> {
33    Ok(match d {
34        KvCacheDtype::Bf16 => "bf16",
35        KvCacheDtype::Fp8 => "fp8",
36        other => bail!("KV-cache dtype {other:?} has no circuit spelling"),
37    })
38}
39
40/// 2026-09-28: The policy of this process: `kv_cache_dtype` and `lm_head_dtype` come from the
41/// model being built, the rest from the process-wide switches.
42pub fn live_policy(levers: &ModelLevers, kv_cache_dtype: &str, lm_head_dtype: &str) -> Policy {
43    let tiers = match row_tiers() {
44        RowTiers::ByRows => "by_rows",
45        RowTiers::Exact => "exact",
46        RowTiers::Canonical => "canonical",
47    };
48    let ssm_h = if crate::layers::qwen3_ssm::ssm_h_fp16_enabled() {
49        "f16"
50    } else {
51        "f32"
52    };
53    let settings = BTreeMap::from([
54        ("row_tiers".to_string(), tiers.to_string()),
55        ("kv_cache_dtype".to_string(), kv_cache_dtype.to_string()),
56        ("lm_head_dtype".to_string(), lm_head_dtype.to_string()),
57        ("ssm_h_dtype".to_string(), ssm_h.to_string()),
58        // 2026-09-30: The GDN flag cell, which a pinned `--ssm-batched-recurrent` sets and the
59        // multi-sequence dispatch reads; the target default alone misses the pin.
60        (
61            "ssm_batched_recurrent".to_string(),
62            on_off(crate::layers::qwen3_ssm::ssm_batched_recurrent_enabled()),
63        ),
64        (
65            "ssm_ba_gates_hopper".to_string(),
66            on_off(crate::layers::ops::ssm_ba_gates_hopper_enabled()),
67        ),
68        ("gemv_sw".to_string(), on_off(levers.gemv_sw)),
69        (
70            "w4a16_tc".to_string(),
71            on_off(crate::layers::ops::gemv_tc::tc_enabled()),
72        ),
73        (
74            "decode_split_silu".to_string(),
75            on_off(levers.decode_split_silu),
76        ),
77        // 2026-09-30: No fused norm-quantize launch exists in the engine.
78        ("rms_norm_act_quant".to_string(), "off".to_string()),
79    ]);
80    Policy {
81        opt_in_levers: Default::default(),
82        settings,
83    }
84}
85
86/// 2026-09-28: Environment switches of the multi-sequence decode and (2026-09-29) MTP verify
87/// dispatch whose defaults the rules' row bands and arms encode, with the file that reads
88/// each.
89pub const MULTI_SEQ_ENV_SWITCHES: [(&str, &str); 16] = [
90    (
91        "METRALE_GDN_FUSED_VERIFY",
92        "qwen3_ssm/trait_decode_batched_conv_gdn.rs",
93    ),
94    (
95        "METRALE_NO_BATCHED_BA_GATES",
96        "qwen3_ssm/trait_decode_batched.rs",
97    ),
98    (
99        "METRALE_NO_BATCHED_GDN_NORM",
100        "qwen3_ssm/trait_decode_batched.rs",
101    ),
102    (
103        "METRALE_SSM_TC_PROJ",
104        "qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs",
105    ),
106    ("METRALE_NO_SSM_M128", "qwen3_ssm/gdn_flags.rs"),
107    (
108        "METRALE_SSM_FFN_PREFILL_MIN_N",
109        "qwen3_ssm/trait_decode_multi_seq.rs",
110    ),
111    (
112        "METRALE_NO_SSM_FFN_PREFILL",
113        "qwen3_ssm/trait_decode_multi_seq.rs",
114    ),
115    (
116        "METRALE_NO_QK_NORM_STRIDED",
117        "qwen3_attention/trait_impl/multi_seq/qkv.rs",
118    ),
119    (
120        "METRALE_NO_ROPE_STRIDED",
121        "qwen3_attention/trait_impl/multi_seq/attn.rs",
122    ),
123    (
124        "METRALE_NO_ATTN_BATCH_CACHE_WRITE",
125        "qwen3_attention/trait_impl/multi_seq/attn.rs",
126    ),
127    ("METRALE_W4A16_K64_MIN_K", "layers/mod.rs (w4a16_k64_min_k)"),
128    ("METRALE_NO_W4A16_K64", "layers/mod.rs (w4a16_k64_min_k)"),
129    ("METRALE_NO_MMQ_SMALL_TILE", "layers/dense_ffn.rs"),
130    ("METRALE_NO_MMQ_TILE64", "layers/dense_ffn.rs"),
131    ("METRALE_W4A16_TC_WIDE", "ops/gemv_tc.rs"),
132    // 2026-09-30: A diagnostic that runs the W4A16 twin of every W4A4 projection.
133    ("METRALE_W4A4_PROJ_AUDIT", "ops/w4a4_proj.rs"),
134];
135
136/// 2026-09-28: Switches that change legacy decode which no rule reads, when they are set.
137pub fn unmodelled_switches(levers: &ModelLevers) -> Vec<String> {
138    let mut out: Vec<String> = MULTI_SEQ_ENV_SWITCHES
139        .iter()
140        .filter(|(var, _)| std::env::var_os(var).is_some())
141        .map(|(var, reader)| format!("{var} (read in {reader})"))
142        .collect();
143    if crate::layers::qwen3_ssm::verify_exact_enabled() {
144        out.push("the exact MTP verify chain (--exact-verify)".to_string());
145    }
146    if !levers.ffn_small_m {
147        out.push("the small-M projection GEMMs off (METRALE_FFN_SMALLM)".to_string());
148    }
149    if levers.ssm_ms_profile || levers.ssm_detail_profile || levers.ms_profile || levers.conc_hsd {
150        out.push(
151            "multi-sequence profiling or hidden dumps (METRALE_MS_PROFILE, METRALE_CONC_HSD, \
152             METRALE_SSM_*_PROFILE)"
153                .to_string(),
154        );
155    }
156    if crate::layers::qwen3_ssm::gdn_fused_norm_enabled() {
157        out.push("fused GDN output norm (METRALE_GDN_FUSED_NORM)".to_string());
158    }
159    if levers.decode_ffn_via_gemm {
160        out.push("dense FFN decode through the GEMM (METRALE_DECODE_FFN_VIA_GEMM)".to_string());
161    }
162    if levers.fp32_routing {
163        out.push("FP32 MoE routing (METRALE_FP32_ROUTING)".to_string());
164    }
165    if levers.ssm_save_dump {
166        out.push("SSM save dump (METRALE_SSM_SAVE_DUMP)".to_string());
167    }
168    out
169}