spark_model/layers/ops/
dispatch_config.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GEMM-path selection, resolved once and then **carried**.
4//!
5//! These flags used to be nine `OnceLock` statics that read `ATLAS_*` at first
6//! touch. A static is the wrong home for them twice over:
7//!
8//! * **It outlives the model whose flags it encodes.** Swap to a model whose
9//!   recipe sets different levers and the process keeps serving the previous
10//!   model's dispatch decisions — silently, because a cached `bool` has no way
11//!   to say it is stale.
12//! * **It hides a dependency.** A function that reads the environment through a
13//!   static takes no argument that says so, cannot be tested with a different
14//!   configuration without mutating the process, and gives the compiler nothing
15//!   to check.
16//!
17//! Carrying it on [`crate::layer::ForwardContext`] — which already reaches
18//! every dispatch site — fixes both. The value is resolved once when the model
19//! is built, borrowed for the duration of that model's run, and dropped with
20//! it. If a future context is missed, the build fails; there is no runtime
21//! check to forget.
22
23/// Which GEMM implementation each projection takes.
24///
25/// Plain `Copy` data, resolved from the environment at model construction.
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub struct GemmDispatch {
28    /// Block-scaled FP8 prefill (per-128-block weight scales + per-token
29    /// activation scales). The DEFAULT for block-scaled FP8 checkpoints since
30    /// 2026-06-17: it matches vLLM's per-block precision and avoids the
31    /// single-scale path, whose collapse of per-block dynamic range pushed
32    /// long-context tool-arg decode into the FP8 argmax-flip regime (B1 drift
33    /// gauge ~1400 → ~100 once block-scaled prefill is on).
34    /// Opt out with `ATLAS_FP8_SINGLE_SCALE=1` — diagnostic/fallback only.
35    pub fp8_blockscaled_prefill: bool,
36    /// cuBLASLt BF16 GEMM. The hand-written mma.sync projection GEMMs reach
37    /// only ~30% of the cuBLAS bf16 ceiling on GB10.
38    pub cublas_gemm: bool,
39    /// Native-FP8 cuBLASLt GEMM.
40    pub cublas_fp8: bool,
41    /// CUTLASS BF16 GEMM, scoped to dense projections using the same FP8→BF16
42    /// cached dequant as cuBLASLt.
43    pub cutlass_gemm: bool,
44    /// Native CUTLASS NVFP4 GEMM: quantizes activations to CUTLASS NVFP4 and
45    /// consumes transposed Atlas NVFP4 weights after repacking scales into the
46    /// CUTLASS SM120 layout. Implies every per-projection NVFP4 flag below.
47    pub cutlass_nvfp4_gemm: bool,
48    pub cutlass_nvfp4_qkvz: bool,
49    pub cutlass_nvfp4_attn_q: bool,
50    pub cutlass_nvfp4_attn_kv: bool,
51    pub cutlass_nvfp4_attn_o: bool,
52    pub cutlass_nvfp4_ssm_out: bool,
53    /// `ATLAS_W4A16_VARIANT` — 1/2/3 pin a kernel variant, 0 = auto (v2).
54    /// A dispatch decision like every other field here, so it belongs on the
55    /// struct the forward pass already carries rather than in a `OnceLock`
56    /// that would pin the first model's choice.
57    pub w4a16_variant: u8,
58}
59
60fn from_values(mut value: impl FnMut(&str) -> Option<String>) -> GemmDispatch {
61    fn on(value: &mut impl FnMut(&str) -> Option<String>, var: &str) -> bool {
62        value(var).as_deref() == Some("1")
63    }
64
65    let all_nvfp4 = on(&mut value, "ATLAS_CUTLASS_NVFP4_GEMM");
66    GemmDispatch {
67        w4a16_variant: match value("ATLAS_W4A16_VARIANT").as_deref() {
68            Some("v1") => 1,
69            Some("v2") => 2,
70            Some("v3") => 3,
71            _ => 0,
72        },
73        fp8_blockscaled_prefill: !on(&mut value, "ATLAS_FP8_SINGLE_SCALE"),
74        cublas_gemm: on(&mut value, "ATLAS_CUBLAS_GEMM"),
75        cublas_fp8: on(&mut value, "ATLAS_CUBLAS_FP8"),
76        cutlass_gemm: on(&mut value, "ATLAS_CUTLASS_GEMM"),
77        cutlass_nvfp4_gemm: all_nvfp4,
78        cutlass_nvfp4_qkvz: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_QKVZ"),
79        cutlass_nvfp4_attn_q: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_ATTN_Q"),
80        cutlass_nvfp4_attn_kv: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_ATTN_KV"),
81        cutlass_nvfp4_attn_o: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_ATTN_O"),
82        // Deliberately NOT implied by the umbrella flag.
83        cutlass_nvfp4_ssm_out: on(&mut value, "ATLAS_CUTLASS_NVFP4_SSM_OUT"),
84    }
85}
86
87impl GemmDispatch {
88    /// Resolve from the environment. Called once, when the model is built.
89    pub fn from_env() -> Self {
90        from_values(|var| std::env::var(var).ok())
91    }
92
93    /// Everything off, block-scaled FP8 prefill on — the shape a build with no
94    /// `ATLAS_*` set in the environment resolves to. Tests construct a context
95    /// with this instead of mutating the process environment.
96    pub fn defaults() -> Self {
97        Self {
98            w4a16_variant: 0,
99            fp8_blockscaled_prefill: true,
100            cublas_gemm: false,
101            cublas_fp8: false,
102            cutlass_gemm: false,
103            cutlass_nvfp4_gemm: false,
104            cutlass_nvfp4_qkvz: false,
105            cutlass_nvfp4_attn_q: false,
106            cutlass_nvfp4_attn_kv: false,
107            cutlass_nvfp4_attn_o: false,
108            cutlass_nvfp4_ssm_out: false,
109        }
110    }
111
112    /// NVFP4 attention Q/K/V enabled for the named projection.
113    pub fn cutlass_nvfp4_attn_qkv(&self, label: &str) -> bool {
114        match label {
115            "q_proj" => self.cutlass_nvfp4_attn_q,
116            "k_proj" | "v_proj" => self.cutlass_nvfp4_attn_kv,
117            _ => self.cutlass_nvfp4_gemm,
118        }
119    }
120}
121
122impl Default for GemmDispatch {
123    fn default() -> Self {
124        Self::defaults()
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use std::collections::HashMap;
132
133    fn resolve(values: &[(&str, &str)]) -> GemmDispatch {
134        let values: HashMap<_, _> = values.iter().copied().collect();
135        from_values(|name| values.get(name).map(|value| (*value).to_owned()))
136    }
137
138    #[test]
139    fn defaults_have_only_blockscaled_prefill_on() {
140        let d = GemmDispatch::defaults();
141        assert_eq!(
142            resolve(&[]),
143            d,
144            "absent environment uses the public default"
145        );
146        assert_eq!(
147            d,
148            GemmDispatch {
149                fp8_blockscaled_prefill: true,
150                cublas_gemm: false,
151                cublas_fp8: false,
152                cutlass_gemm: false,
153                cutlass_nvfp4_gemm: false,
154                cutlass_nvfp4_qkvz: false,
155                cutlass_nvfp4_attn_q: false,
156                cutlass_nvfp4_attn_kv: false,
157                cutlass_nvfp4_attn_o: false,
158                cutlass_nvfp4_ssm_out: false,
159                w4a16_variant: 0,
160            }
161        );
162    }
163
164    #[test]
165    fn the_umbrella_flag_implies_the_per_projection_ones() {
166        let d = resolve(&[("ATLAS_CUTLASS_NVFP4_GEMM", "1")]);
167        assert!(d.cutlass_nvfp4_gemm);
168        assert!(d.cutlass_nvfp4_qkvz);
169        assert!(d.cutlass_nvfp4_attn_qkv("q_proj"));
170        assert!(d.cutlass_nvfp4_attn_qkv("k_proj"));
171        assert!(d.cutlass_nvfp4_attn_qkv("v_proj"));
172        assert!(d.cutlass_nvfp4_attn_o);
173        // SSM-out was never implied by the umbrella flag.
174        assert!(!d.cutlass_nvfp4_ssm_out);
175    }
176
177    #[test]
178    fn per_projection_flags_are_independent() {
179        let cases = [
180            (
181                "ATLAS_CUTLASS_NVFP4_QKVZ",
182                [true, false, false, false, false],
183            ),
184            (
185                "ATLAS_CUTLASS_NVFP4_ATTN_Q",
186                [false, true, false, false, false],
187            ),
188            (
189                "ATLAS_CUTLASS_NVFP4_ATTN_KV",
190                [false, false, true, false, false],
191            ),
192            (
193                "ATLAS_CUTLASS_NVFP4_ATTN_O",
194                [false, false, false, true, false],
195            ),
196            (
197                "ATLAS_CUTLASS_NVFP4_SSM_OUT",
198                [false, false, false, false, true],
199            ),
200        ];
201        for (name, expected) in cases {
202            let d = resolve(&[(name, "1")]);
203            assert_eq!(
204                [
205                    d.cutlass_nvfp4_qkvz,
206                    d.cutlass_nvfp4_attn_q,
207                    d.cutlass_nvfp4_attn_kv,
208                    d.cutlass_nvfp4_attn_o,
209                    d.cutlass_nvfp4_ssm_out,
210                ],
211                expected,
212                "{name} must not enable a neighboring projection"
213            );
214        }
215    }
216
217    #[test]
218    fn non_nvfp4_flags_map_independently_and_single_scale_is_inverted() {
219        let cases = [
220            ("ATLAS_CUBLAS_GEMM", [true, false, false]),
221            ("ATLAS_CUBLAS_FP8", [false, true, false]),
222            ("ATLAS_CUTLASS_GEMM", [false, false, true]),
223        ];
224        for (name, expected) in cases {
225            let d = resolve(&[(name, "1")]);
226            assert_eq!(
227                [d.cublas_gemm, d.cublas_fp8, d.cutlass_gemm],
228                expected,
229                "{name} must not enable a neighboring GEMM path"
230            );
231            assert!(d.fp8_blockscaled_prefill);
232        }
233        assert!(!resolve(&[("ATLAS_FP8_SINGLE_SCALE", "1")]).fp8_blockscaled_prefill);
234        assert!(
235            !resolve(&[("ATLAS_CUBLAS_GEMM", "true")]).cublas_gemm,
236            "dispatch booleans accept only the documented literal 1"
237        );
238    }
239
240    #[test]
241    fn w4a16_variants_accept_only_documented_spellings() {
242        for (value, expected) in [
243            ("v1", 1),
244            ("v2", 2),
245            ("v3", 3),
246            ("1", 0),
247            ("V1", 0),
248            ("unknown", 0),
249        ] {
250            assert_eq!(
251                resolve(&[("ATLAS_W4A16_VARIANT", value)]).w4a16_variant,
252                expected,
253                "value {value}"
254            );
255        }
256    }
257
258    #[test]
259    fn an_unknown_projection_label_falls_back_to_the_umbrella_flag() {
260        assert!(!GemmDispatch::defaults().cutlass_nvfp4_attn_qkv("mystery"));
261        let d = GemmDispatch {
262            cutlass_nvfp4_gemm: true,
263            ..GemmDispatch::defaults()
264        };
265        assert!(d.cutlass_nvfp4_attn_qkv("mystery"));
266    }
267}
spark_model/layers/ops/dispatch_config.rs

spark_model/layers/ops/
dispatch_config.rs
