metrale_model_layers/layers/ops/
dispatch_config.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: GEMM-path selection: [`GemmDispatch`], resolved from `METRALE_*`
4//! variables and carried to the dispatch sites on
5//! [`crate::layer::ForwardContext`].
6//!
7//! Owner: model-layers ops.
8//! Invariants:
9//! - A switch variable counts as on only when its value is exactly `1`;
10//!   `METRALE_CUBLAS_GEMM` has its own grammar ([`CublasScope`]) and
11//!   `METRALE_W4A16_VARIANT` accepts `v1`, `v2` and `v3`.
12//! - Unknown `METRALE_CUBLAS_GEMM` tokens never add a family.
13
14/// 2026-09-25: Which projection families take a cuBLASLt GEMM arm, so one
15/// family can be armed without the others.
16///
17/// Grammar: `METRALE_CUBLAS_GEMM=<token>[,<token>]*`, ASCII-case-insensitive,
18/// whitespace around a token ignored:
19///
20/// | token | meaning |
21/// |---|---|
22/// | `ffn` | dense-FFN and MoE shared-expert projections |
23/// | `attn` | attention Q/K/V, O and the output gate |
24/// | `ssm` | GDN `in_proj_qkvz` and `out_proj` |
25/// | `head` | LM / MTP head (see [`CublasScope::head`]) |
26/// | `all`, `1`, `true` | every family |
27/// | `off`, `0`, `false`, empty | the empty set |
28///
29/// The result is the union of the tokens, so `off` clears nothing (`ffn,off`
30/// is `ffn`). Unknown tokens are dropped and reported, and never widen the
31/// set.
32#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
33pub struct CublasScope {
34    /// 2026-09-25: Dense-FFN gate/up/down and the MoE shared expert.
35    pub ffn: bool,
36    /// 2026-09-25: Attention Q/K/V, O projection and the output gate.
37    pub attn: bool,
38    /// 2026-09-25: GDN `in_proj_qkvz` and `out_proj`, in prefill and decode.
39    pub ssm: bool,
40    /// 2026-09-25: LM / MTP head. Parsed and set by `all`, but no dispatch site
41    /// reads it, so setting it changes nothing.
42    pub head: bool,
43}
44
45impl CublasScope {
46    /// 2026-09-25: No family armed: what an absent, empty or `off`
47    /// `METRALE_CUBLAS_GEMM` resolves to, and the value in
48    /// [`GemmDispatch::defaults`].
49    pub const OFF: Self = Self {
50        ffn: false,
51        attn: false,
52        ssm: false,
53        head: false,
54    };
55
56    /// 2026-09-25: Every family: what `all`, `1` and `true` resolve to.
57    pub const ALL: Self = Self {
58        ffn: true,
59        attn: true,
60        ssm: true,
61        head: true,
62    };
63
64    /// 2026-09-25: Whether any family is armed.
65    pub fn any(&self) -> bool {
66        self.ffn || self.attn || self.ssm || self.head
67    }
68}
69
70/// 2026-09-25: Parse the [`CublasScope`] grammar. Returns the resolved set and
71/// the lowercased tokens that matched nothing. Pure: the environment read and
72/// the logging are in [`GemmDispatch::from_env`].
73pub fn parse_cublas_scope(raw: Option<&str>) -> (CublasScope, Vec<String>) {
74    let mut scope = CublasScope::OFF;
75    let mut unknown = Vec::new();
76    let Some(raw) = raw else {
77        return (scope, unknown);
78    };
79    for token in raw.split(',') {
80        match token.trim().to_ascii_lowercase().as_str() {
81            "" | "0" | "false" | "off" => {}
82            "1" | "true" | "all" => scope = CublasScope::ALL,
83            "ffn" => scope.ffn = true,
84            "attn" => scope.attn = true,
85            "ssm" => scope.ssm = true,
86            "head" => scope.head = true,
87            other => unknown.push(other.to_owned()),
88        }
89    }
90    (scope, unknown)
91}
92
93/// 2026-09-25: Which GEMM implementation each projection takes, resolved by
94/// [`GemmDispatch::from_env`].
95#[derive(Clone, Copy, Debug, PartialEq, Eq)]
96pub struct GemmDispatch {
97    /// 2026-09-25: Block-scaled FP8 prefill (per-128-block weight scales and
98    /// per-token activation scales). On unless `METRALE_FP8_SINGLE_SCALE=1`.
99    pub fp8_blockscaled_prefill: bool,
100    /// 2026-09-25: Which projection families take a cuBLASLt GEMM arm
101    /// (`METRALE_CUBLAS_GEMM`).
102    pub cublas: CublasScope,
103    /// 2026-09-25: Row-wise FP8 cuBLASLt GEMM for the GDN `in_proj_qkvz`
104    /// prefill (`METRALE_CUBLAS_FP8=1`).
105    pub cublas_fp8: bool,
106    /// 2026-09-25: CUTLASS BF16 GEMM for the GDN `in_proj_qkvz` prefill, on a
107    /// memoized FP8→BF16 dequant of the weight (`METRALE_CUTLASS_GEMM=1`).
108    pub cutlass_gemm: bool,
109    /// 2026-09-25: CUTLASS NVFP4 GEMM (`METRALE_CUTLASS_NVFP4_GEMM=1`). Implies
110    /// the per-projection NVFP4 flags below except `cutlass_nvfp4_ssm_out`.
111    pub cutlass_nvfp4_gemm: bool,
112    pub cutlass_nvfp4_qkvz: bool,
113    pub cutlass_nvfp4_attn_q: bool,
114    pub cutlass_nvfp4_attn_kv: bool,
115    pub cutlass_nvfp4_attn_o: bool,
116    pub cutlass_nvfp4_ssm_out: bool,
117    /// 2026-09-25: `METRALE_W4A16_VARIANT`: `v1`, `v2` and `v3` give 1, 2 and 3;
118    /// any other value, or none, gives 0, which the attention prefill treats
119    /// like 2.
120    pub w4a16_variant: u8,
121}
122
123fn from_values(mut value: impl FnMut(&str) -> Option<String>) -> GemmDispatch {
124    fn on(value: &mut impl FnMut(&str) -> Option<String>, var: &str) -> bool {
125        value(var).as_deref() == Some("1")
126    }
127
128    let all_nvfp4 = on(&mut value, "METRALE_CUTLASS_NVFP4_GEMM");
129    GemmDispatch {
130        w4a16_variant: match value("METRALE_W4A16_VARIANT").as_deref() {
131            Some("v1") => 1,
132            Some("v2") => 2,
133            Some("v3") => 3,
134            _ => 0,
135        },
136        fp8_blockscaled_prefill: !on(&mut value, "METRALE_FP8_SINGLE_SCALE"),
137        cublas: parse_cublas_scope(value("METRALE_CUBLAS_GEMM").as_deref()).0,
138        cublas_fp8: on(&mut value, "METRALE_CUBLAS_FP8"),
139        cutlass_gemm: on(&mut value, "METRALE_CUTLASS_GEMM"),
140        cutlass_nvfp4_gemm: all_nvfp4,
141        cutlass_nvfp4_qkvz: all_nvfp4 || on(&mut value, "METRALE_CUTLASS_NVFP4_QKVZ"),
142        cutlass_nvfp4_attn_q: all_nvfp4 || on(&mut value, "METRALE_CUTLASS_NVFP4_ATTN_Q"),
143        cutlass_nvfp4_attn_kv: all_nvfp4 || on(&mut value, "METRALE_CUTLASS_NVFP4_ATTN_KV"),
144        cutlass_nvfp4_attn_o: all_nvfp4 || on(&mut value, "METRALE_CUTLASS_NVFP4_ATTN_O"),
145        // 2026-09-25: Not implied by `METRALE_CUTLASS_NVFP4_GEMM`.
146        cutlass_nvfp4_ssm_out: on(&mut value, "METRALE_CUTLASS_NVFP4_SSM_OUT"),
147    }
148}
149
150impl GemmDispatch {
151    /// 2026-09-25: Resolve from the environment, and log the resolved
152    /// `METRALE_CUBLAS_GEMM` scope when that variable is set.
153    pub fn from_env() -> Self {
154        let raw = std::env::var("METRALE_CUBLAS_GEMM").ok();
155        let resolved = from_values(metrale_config::levers::var);
156        log_cublas_scope(raw.as_deref(), resolved.cublas);
157        resolved
158    }
159
160    /// 2026-09-25: Everything off except block-scaled FP8 prefill: what an
161    /// environment with none of these variables resolves to.
162    pub fn defaults() -> Self {
163        Self {
164            w4a16_variant: 0,
165            fp8_blockscaled_prefill: true,
166            cublas: CublasScope::OFF,
167            cublas_fp8: false,
168            cutlass_gemm: false,
169            cutlass_nvfp4_gemm: false,
170            cutlass_nvfp4_qkvz: false,
171            cutlass_nvfp4_attn_q: false,
172            cutlass_nvfp4_attn_kv: false,
173            cutlass_nvfp4_attn_o: false,
174            cutlass_nvfp4_ssm_out: false,
175        }
176    }
177
178    /// 2026-09-25: Whether CUTLASS NVFP4 is on for an attention projection:
179    /// `q_proj` reads the Q flag, `k_proj` and `v_proj` the KV flag, and any
180    /// other label the umbrella flag.
181    pub fn cutlass_nvfp4_attn_qkv(&self, label: &str) -> bool {
182        match label {
183            "q_proj" => self.cutlass_nvfp4_attn_q,
184            "k_proj" | "v_proj" => self.cutlass_nvfp4_attn_kv,
185            _ => self.cutlass_nvfp4_gemm,
186        }
187    }
188}
189
190/// 2026-09-25: Log the cuBLASLt families a set `METRALE_CUBLAS_GEMM` resolved
191/// to, and warn about its unknown tokens. Silent when the variable is unset.
192fn log_cublas_scope(raw: Option<&str>, scope: CublasScope) {
193    let Some(raw) = raw else {
194        return;
195    };
196    let (_, unknown) = parse_cublas_scope(Some(raw));
197    if !unknown.is_empty() {
198        tracing::warn!(
199            "METRALE_CUBLAS_GEMM={raw:?}: ignoring unknown families [{}]. The grammar is a \
200             comma-separated subset of all|ffn|attn|ssm|head|off (1/true = all).",
201            unknown.join(", ")
202        );
203    }
204    tracing::info!(
205        "[metrale] METRALE_CUBLAS_GEMM={raw:?} -> cuBLASLt arms ffn={} attn={} ssm={} head={} \
206         (head has no consumer yet){}",
207        scope.ffn,
208        scope.attn,
209        scope.ssm,
210        scope.head,
211        if scope.any() {
212            ""
213        } else {
214            " — no arm enabled"
215        }
216    );
217}
218
219impl Default for GemmDispatch {
220    fn default() -> Self {
221        Self::defaults()
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use std::collections::HashMap;
229
230    fn resolve(values: &[(&str, &str)]) -> GemmDispatch {
231        let values: HashMap<_, _> = values.iter().copied().collect();
232        from_values(|name| values.get(name).map(|value| (*value).to_owned()))
233    }
234
235    #[test]
236    fn defaults_have_only_blockscaled_prefill_on() {
237        let d = GemmDispatch::defaults();
238        assert_eq!(
239            resolve(&[]),
240            d,
241            "absent environment uses the public default"
242        );
243        assert_eq!(
244            d,
245            GemmDispatch {
246                fp8_blockscaled_prefill: true,
247                cublas: CublasScope::OFF,
248                cublas_fp8: false,
249                cutlass_gemm: false,
250                cutlass_nvfp4_gemm: false,
251                cutlass_nvfp4_qkvz: false,
252                cutlass_nvfp4_attn_q: false,
253                cutlass_nvfp4_attn_kv: false,
254                cutlass_nvfp4_attn_o: false,
255                cutlass_nvfp4_ssm_out: false,
256                w4a16_variant: 0,
257            }
258        );
259    }
260
261    #[test]
262    fn the_umbrella_flag_implies_the_per_projection_ones() {
263        let d = resolve(&[("METRALE_CUTLASS_NVFP4_GEMM", "1")]);
264        assert!(d.cutlass_nvfp4_gemm);
265        assert!(d.cutlass_nvfp4_qkvz);
266        assert!(d.cutlass_nvfp4_attn_qkv("q_proj"));
267        assert!(d.cutlass_nvfp4_attn_qkv("k_proj"));
268        assert!(d.cutlass_nvfp4_attn_qkv("v_proj"));
269        assert!(d.cutlass_nvfp4_attn_o);
270        // 2026-09-25: SSM-out is not implied by the umbrella flag.
271        assert!(!d.cutlass_nvfp4_ssm_out);
272    }
273
274    #[test]
275    fn per_projection_flags_are_independent() {
276        let cases = [
277            (
278                "METRALE_CUTLASS_NVFP4_QKVZ",
279                [true, false, false, false, false],
280            ),
281            (
282                "METRALE_CUTLASS_NVFP4_ATTN_Q",
283                [false, true, false, false, false],
284            ),
285            (
286                "METRALE_CUTLASS_NVFP4_ATTN_KV",
287                [false, false, true, false, false],
288            ),
289            (
290                "METRALE_CUTLASS_NVFP4_ATTN_O",
291                [false, false, false, true, false],
292            ),
293            (
294                "METRALE_CUTLASS_NVFP4_SSM_OUT",
295                [false, false, false, false, true],
296            ),
297        ];
298        for (name, expected) in cases {
299            let d = resolve(&[(name, "1")]);
300            assert_eq!(
301                [
302                    d.cutlass_nvfp4_qkvz,
303                    d.cutlass_nvfp4_attn_q,
304                    d.cutlass_nvfp4_attn_kv,
305                    d.cutlass_nvfp4_attn_o,
306                    d.cutlass_nvfp4_ssm_out,
307                ],
308                expected,
309                "{name} must not enable a neighboring projection"
310            );
311        }
312    }
313
314    #[test]
315    fn non_nvfp4_flags_map_independently_and_single_scale_is_inverted() {
316        let cases = [
317            ("METRALE_CUBLAS_GEMM", [true, false, false]),
318            ("METRALE_CUBLAS_FP8", [false, true, false]),
319            ("METRALE_CUTLASS_GEMM", [false, false, true]),
320        ];
321        for (name, expected) in cases {
322            let d = resolve(&[(name, "1")]);
323            assert_eq!(
324                [d.cublas.any(), d.cublas_fp8, d.cutlass_gemm],
325                expected,
326                "{name} must not enable a neighboring GEMM path"
327            );
328            assert!(d.fp8_blockscaled_prefill);
329        }
330        assert!(!resolve(&[("METRALE_FP8_SINGLE_SCALE", "1")]).fp8_blockscaled_prefill);
331    }
332
333    fn scope(raw: &str) -> CublasScope {
334        resolve(&[("METRALE_CUBLAS_GEMM", raw)]).cublas
335    }
336
337    /// 2026-09-25: Each spelling in the [`CublasScope`] table maps to its
338    /// family set.
339    #[test]
340    fn the_scope_grammar_maps_each_spelling_to_its_family_set() {
341        let f = |ffn, attn, ssm, head| CublasScope {
342            ffn,
343            attn,
344            ssm,
345            head,
346        };
347        let cases: [(&str, CublasScope); 13] = [
348            ("all", CublasScope::ALL),
349            ("1", CublasScope::ALL),
350            ("true", CublasScope::ALL),
351            ("ALL", CublasScope::ALL),
352            ("off", CublasScope::OFF),
353            ("0", CublasScope::OFF),
354            ("false", CublasScope::OFF),
355            ("", CublasScope::OFF),
356            ("ffn", f(true, false, false, false)),
357            ("attn", f(false, true, false, false)),
358            ("ssm", f(false, false, true, false)),
359            ("head", f(false, false, false, true)),
360            ("ffn,attn", f(true, true, false, false)),
361        ];
362        for (raw, expected) in cases {
363            assert_eq!(scope(raw), expected, "METRALE_CUBLAS_GEMM={raw:?}");
364        }
365        assert_eq!(scope(" ffn , ssm "), f(true, false, true, false));
366        // 2026-09-25: A union, so `off` clears nothing; a clearing token would
367        // make the result depend on token order.
368        assert_eq!(scope("ffn,off"), f(true, false, false, false));
369    }
370
371    /// 2026-09-25: A typo must not widen the set; it is dropped and reported
372    /// as typed, lowercased.
373    #[test]
374    fn unknown_families_are_dropped_and_reported_never_widening_the_set() {
375        assert_eq!(scope("junk"), CublasScope::OFF);
376        assert_eq!(
377            scope("ffn,junk"),
378            CublasScope {
379                ffn: true,
380                ..CublasScope::OFF
381            }
382        );
383        let (resolved, unknown) = parse_cublas_scope(Some("ffn, FNN ,bogus"));
384        assert_eq!(
385            resolved,
386            CublasScope {
387                ffn: true,
388                ..CublasScope::OFF
389            }
390        );
391        assert_eq!(
392            unknown,
393            vec!["fnn".to_owned(), "bogus".to_owned()],
394            "the warning must name what the operator typed, lowercased"
395        );
396    }
397
398    /// 2026-09-25: An absent variable and `off` both resolve to the empty set
399    /// with no unknown token to warn about.
400    #[test]
401    fn an_absent_variable_resolves_to_the_empty_set_silently() {
402        assert_eq!(parse_cublas_scope(None), (CublasScope::OFF, Vec::new()));
403        assert_eq!(
404            parse_cublas_scope(Some("off")),
405            (CublasScope::OFF, Vec::new())
406        );
407        assert!(!CublasScope::OFF.any());
408        assert!(CublasScope::ALL.any());
409    }
410
411    #[test]
412    fn w4a16_variants_accept_only_documented_spellings() {
413        for (value, expected) in [
414            ("v1", 1),
415            ("v2", 2),
416            ("v3", 3),
417            ("1", 0),
418            ("V1", 0),
419            ("unknown", 0),
420        ] {
421            assert_eq!(
422                resolve(&[("METRALE_W4A16_VARIANT", value)]).w4a16_variant,
423                expected,
424                "value {value}"
425            );
426        }
427    }
428
429    #[test]
430    fn an_unknown_projection_label_falls_back_to_the_umbrella_flag() {
431        assert!(!GemmDispatch::defaults().cutlass_nvfp4_attn_qkv("mystery"));
432        let d = GemmDispatch {
433            cutlass_nvfp4_gemm: true,
434            ..GemmDispatch::defaults()
435        };
436        assert!(d.cutlass_nvfp4_attn_qkv("mystery"));
437    }
438}
metrale_model_layers/layers/ops/dispatch_config.rs

metrale_model_layers/layers/ops/
dispatch_config.rs
