metrale_model_layers/circuit_exec/
policy.rs1use 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
30pub 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
40pub 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 (
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 ("rms_norm_act_quant".to_string(), "off".to_string()),
79 ]);
80 Policy {
81 opt_in_levers: Default::default(),
82 settings,
83 }
84}
85
86pub 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 ("METRALE_W4A4_PROJ_AUDIT", "ops/w4a4_proj.rs"),
134];
135
136pub 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}