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 projection FAMILIES take a cuBLASLt GEMM arm.
24///
25/// WHY a set and not a `bool` — H100, 2026-09-11, `Qwen/Qwen3.8-27B-FP8`
26/// native FP8, tip `5f78270dc`, "config B" of the round-3 receipt.
27/// `ATLAS_CUBLAS_GEMM=1` resolved to ONE global boolean, so the variable that
28/// arms #917's dense-FFN W8A8 fast path ALSO armed the SSM `in_proj_qkvz` arm,
29/// which materialised a cached BF16 dequant of the fused QKVZ weight
30/// (`[10240,5120] + [6144,5120]` x 2 B = `167772160` bytes per layer, ~10.3
31/// GiB over 48 SSM layers) outside the buffer ledger. One 28-token prefill ate
32/// 6120 MiB and died at layer 36 with `cuMemAlloc_v2 ... status 2`. The FFN arm
33/// under test could not be exercised end-to-end on an 80 GB card, and no knob
34/// separated the two. The CUTLASS family next door already spells its
35/// projections out (`ATLAS_CUTLASS_NVFP4_QKVZ`, `..._ATTN_Q`, `..._ATTN_KV`,
36/// `..._ATTN_O`, `..._SSM_OUT`); this gives cuBLASLt the same property in one
37/// variable instead of five.
38///
39/// GRAMMAR — `ATLAS_CUBLAS_GEMM=<token>[,<token>]*`, ASCII-case-insensitive,
40/// whitespace around a token ignored:
41///
42/// | token | meaning |
43/// |---|---|
44/// | `ffn` | dense-FFN + MoE shared-expert projections |
45/// | `attn` | attention Q/K/V, O and the output gate |
46/// | `ssm` | SSM/GDN `in_proj_qkvz` |
47/// | `head` | LM / MTP head (see [`CublasScope::head`]) |
48/// | `all`, `1`, `true` | every family — the pre-2026-09-11 spelling |
49/// | `off`, `0`, `false`, empty | the empty set |
50///
51/// The result is the UNION of the tokens, so `off` adds nothing rather than
52/// clearing what another token armed (`ffn,off` is `ffn`). Unknown tokens are
53/// dropped with a warning and never widen the set: a typo must not silently arm
54/// an arm, which is the exact failure above.
55#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
56pub struct CublasScope {
57 /// Dense-FFN gate/up/down and the MoE shared expert — #917's W8A8 arm.
58 pub ffn: bool,
59 /// Attention Q/K/V, O projection and the output gate.
60 pub attn: bool,
61 /// SSM/GDN fused `in_proj_qkvz` prefill projection.
62 pub ssm: bool,
63 /// LM / MTP head. Parsed, carried and covered by `all`, but NO dispatch
64 /// site reads it yet — setting `head` is inert today. Named anyway so the
65 /// grammar is the complete family list and `all` has a fixed meaning; a
66 /// lever that silently drops a spelling is worse than one that documents
67 /// an unclaimed slot.
68 pub head: bool,
69}
70
71impl CublasScope {
72 /// No family armed — the shape an absent or `off` `ATLAS_CUBLAS_GEMM`
73 /// resolves to, and the default for every build.
74 pub const OFF: Self = Self {
75 ffn: false,
76 attn: false,
77 ssm: false,
78 head: false,
79 };
80
81 /// Every family — what `all` / `1` / `true` resolve to.
82 pub const ALL: Self = Self {
83 ffn: true,
84 attn: true,
85 ssm: true,
86 head: true,
87 };
88
89 /// Whether any family is armed (for the resolved-set log line).
90 pub fn any(&self) -> bool {
91 self.ffn || self.attn || self.ssm || self.head
92 }
93}
94
95/// Parse the [`CublasScope`] grammar. Returns the resolved set plus the tokens
96/// that matched nothing, so the caller can warn with the operator's own
97/// spelling. Pure — the environment read and the logging both live in
98/// [`GemmDispatch::from_env`], which is what makes the table testable.
99pub fn parse_cublas_scope(raw: Option<&str>) -> (CublasScope, Vec<String>) {
100 let mut scope = CublasScope::OFF;
101 let mut unknown = Vec::new();
102 let Some(raw) = raw else {
103 return (scope, unknown);
104 };
105 for token in raw.split(',') {
106 match token.trim().to_ascii_lowercase().as_str() {
107 "" | "0" | "false" | "off" => {}
108 "1" | "true" | "all" => scope = CublasScope::ALL,
109 "ffn" => scope.ffn = true,
110 "attn" => scope.attn = true,
111 "ssm" => scope.ssm = true,
112 "head" => scope.head = true,
113 other => unknown.push(other.to_owned()),
114 }
115 }
116 (scope, unknown)
117}
118
119/// Which GEMM implementation each projection takes.
120///
121/// Plain `Copy` data, resolved from the environment at model construction.
122#[derive(Clone, Copy, Debug, PartialEq, Eq)]
123pub struct GemmDispatch {
124 /// Block-scaled FP8 prefill (per-128-block weight scales + per-token
125 /// activation scales). The DEFAULT for block-scaled FP8 checkpoints since
126 /// 2026-06-17: it matches vLLM's per-block precision and avoids the
127 /// single-scale path, whose collapse of per-block dynamic range pushed
128 /// long-context tool-arg decode into the FP8 argmax-flip regime (B1 drift
129 /// gauge ~1400 → ~100 once block-scaled prefill is on).
130 /// Opt out with `ATLAS_FP8_SINGLE_SCALE=1` — diagnostic/fallback only.
131 pub fp8_blockscaled_prefill: bool,
132 /// Which projection families take a cuBLASLt GEMM arm
133 /// (`ATLAS_CUBLAS_GEMM`). The hand-written mma.sync projection GEMMs reach
134 /// only ~30% of the cuBLAS bf16 ceiling on GB10, which is why the arms
135 /// exist; [`CublasScope`] is why they are no longer all one switch.
136 pub cublas: CublasScope,
137 /// Native-FP8 cuBLASLt GEMM.
138 pub cublas_fp8: bool,
139 /// CUTLASS BF16 GEMM, scoped to dense projections using the same FP8→BF16
140 /// cached dequant as cuBLASLt.
141 pub cutlass_gemm: bool,
142 /// Native CUTLASS NVFP4 GEMM: quantizes activations to CUTLASS NVFP4 and
143 /// consumes transposed Atlas NVFP4 weights after repacking scales into the
144 /// CUTLASS SM120 layout. Implies every per-projection NVFP4 flag below.
145 pub cutlass_nvfp4_gemm: bool,
146 pub cutlass_nvfp4_qkvz: bool,
147 pub cutlass_nvfp4_attn_q: bool,
148 pub cutlass_nvfp4_attn_kv: bool,
149 pub cutlass_nvfp4_attn_o: bool,
150 pub cutlass_nvfp4_ssm_out: bool,
151 /// `ATLAS_W4A16_VARIANT` — 1/2/3 pin a kernel variant, 0 = auto (v2).
152 /// A dispatch decision like every other field here, so it belongs on the
153 /// struct the forward pass already carries rather than in a `OnceLock`
154 /// that would pin the first model's choice.
155 pub w4a16_variant: u8,
156}
157
158fn from_values(mut value: impl FnMut(&str) -> Option<String>) -> GemmDispatch {
159 fn on(value: &mut impl FnMut(&str) -> Option<String>, var: &str) -> bool {
160 value(var).as_deref() == Some("1")
161 }
162
163 let all_nvfp4 = on(&mut value, "ATLAS_CUTLASS_NVFP4_GEMM");
164 GemmDispatch {
165 w4a16_variant: match value("ATLAS_W4A16_VARIANT").as_deref() {
166 Some("v1") => 1,
167 Some("v2") => 2,
168 Some("v3") => 3,
169 _ => 0,
170 },
171 fp8_blockscaled_prefill: !on(&mut value, "ATLAS_FP8_SINGLE_SCALE"),
172 cublas: parse_cublas_scope(value("ATLAS_CUBLAS_GEMM").as_deref()).0,
173 cublas_fp8: on(&mut value, "ATLAS_CUBLAS_FP8"),
174 cutlass_gemm: on(&mut value, "ATLAS_CUTLASS_GEMM"),
175 cutlass_nvfp4_gemm: all_nvfp4,
176 cutlass_nvfp4_qkvz: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_QKVZ"),
177 cutlass_nvfp4_attn_q: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_ATTN_Q"),
178 cutlass_nvfp4_attn_kv: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_ATTN_KV"),
179 cutlass_nvfp4_attn_o: all_nvfp4 || on(&mut value, "ATLAS_CUTLASS_NVFP4_ATTN_O"),
180 // Deliberately NOT implied by the umbrella flag.
181 cutlass_nvfp4_ssm_out: on(&mut value, "ATLAS_CUTLASS_NVFP4_SSM_OUT"),
182 }
183}
184
185impl GemmDispatch {
186 /// Resolve from the environment. Called once, when the model is built.
187 pub fn from_env() -> Self {
188 let raw = std::env::var("ATLAS_CUBLAS_GEMM").ok();
189 let resolved = from_values(|var| std::env::var(var).ok());
190 log_cublas_scope(raw.as_deref(), resolved.cublas);
191 resolved
192 }
193
194 /// Everything off, block-scaled FP8 prefill on — the shape a build with no
195 /// `ATLAS_*` set in the environment resolves to. Tests construct a context
196 /// with this instead of mutating the process environment.
197 pub fn defaults() -> Self {
198 Self {
199 w4a16_variant: 0,
200 fp8_blockscaled_prefill: true,
201 cublas: CublasScope::OFF,
202 cublas_fp8: false,
203 cutlass_gemm: false,
204 cutlass_nvfp4_gemm: false,
205 cutlass_nvfp4_qkvz: false,
206 cutlass_nvfp4_attn_q: false,
207 cutlass_nvfp4_attn_kv: false,
208 cutlass_nvfp4_attn_o: false,
209 cutlass_nvfp4_ssm_out: false,
210 }
211 }
212
213 /// NVFP4 attention Q/K/V enabled for the named projection.
214 pub fn cutlass_nvfp4_attn_qkv(&self, label: &str) -> bool {
215 match label {
216 "q_proj" => self.cutlass_nvfp4_attn_q,
217 "k_proj" | "v_proj" => self.cutlass_nvfp4_attn_kv,
218 _ => self.cutlass_nvfp4_gemm,
219 }
220 }
221}
222
223/// Say once, at model build, which cuBLASLt arms an `ATLAS_CUBLAS_GEMM` value
224/// actually armed.
225///
226/// A scoped lever is only an improvement if the operator can SEE the scope it
227/// resolved to: the failure this replaces was invisible until a `cuMemAlloc_v2`
228/// error named a layer 36 nobody had aimed at. Silent when the variable is
229/// unset — a serve that never asked for cuBLASLt should not narrate it.
230fn log_cublas_scope(raw: Option<&str>, scope: CublasScope) {
231 let Some(raw) = raw else {
232 return;
233 };
234 let (_, unknown) = parse_cublas_scope(Some(raw));
235 if !unknown.is_empty() {
236 tracing::warn!(
237 "ATLAS_CUBLAS_GEMM={raw:?}: ignoring unknown families [{}]. The grammar is a \
238 comma-separated subset of all|ffn|attn|ssm|head|off (1/true = all).",
239 unknown.join(", ")
240 );
241 }
242 tracing::info!(
243 "[atlas] ATLAS_CUBLAS_GEMM={raw:?} -> cuBLASLt arms ffn={} attn={} ssm={} head={} \
244 (head has no consumer yet){}",
245 scope.ffn,
246 scope.attn,
247 scope.ssm,
248 scope.head,
249 if scope.any() {
250 ""
251 } else {
252 " — no arm enabled"
253 }
254 );
255}
256
257impl Default for GemmDispatch {
258 fn default() -> Self {
259 Self::defaults()
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266 use std::collections::HashMap;
267
268 fn resolve(values: &[(&str, &str)]) -> GemmDispatch {
269 let values: HashMap<_, _> = values.iter().copied().collect();
270 from_values(|name| values.get(name).map(|value| (*value).to_owned()))
271 }
272
273 #[test]
274 fn defaults_have_only_blockscaled_prefill_on() {
275 let d = GemmDispatch::defaults();
276 assert_eq!(
277 resolve(&[]),
278 d,
279 "absent environment uses the public default"
280 );
281 assert_eq!(
282 d,
283 GemmDispatch {
284 fp8_blockscaled_prefill: true,
285 cublas: CublasScope::OFF,
286 cublas_fp8: false,
287 cutlass_gemm: false,
288 cutlass_nvfp4_gemm: false,
289 cutlass_nvfp4_qkvz: false,
290 cutlass_nvfp4_attn_q: false,
291 cutlass_nvfp4_attn_kv: false,
292 cutlass_nvfp4_attn_o: false,
293 cutlass_nvfp4_ssm_out: false,
294 w4a16_variant: 0,
295 }
296 );
297 }
298
299 #[test]
300 fn the_umbrella_flag_implies_the_per_projection_ones() {
301 let d = resolve(&[("ATLAS_CUTLASS_NVFP4_GEMM", "1")]);
302 assert!(d.cutlass_nvfp4_gemm);
303 assert!(d.cutlass_nvfp4_qkvz);
304 assert!(d.cutlass_nvfp4_attn_qkv("q_proj"));
305 assert!(d.cutlass_nvfp4_attn_qkv("k_proj"));
306 assert!(d.cutlass_nvfp4_attn_qkv("v_proj"));
307 assert!(d.cutlass_nvfp4_attn_o);
308 // SSM-out was never implied by the umbrella flag.
309 assert!(!d.cutlass_nvfp4_ssm_out);
310 }
311
312 #[test]
313 fn per_projection_flags_are_independent() {
314 let cases = [
315 (
316 "ATLAS_CUTLASS_NVFP4_QKVZ",
317 [true, false, false, false, false],
318 ),
319 (
320 "ATLAS_CUTLASS_NVFP4_ATTN_Q",
321 [false, true, false, false, false],
322 ),
323 (
324 "ATLAS_CUTLASS_NVFP4_ATTN_KV",
325 [false, false, true, false, false],
326 ),
327 (
328 "ATLAS_CUTLASS_NVFP4_ATTN_O",
329 [false, false, false, true, false],
330 ),
331 (
332 "ATLAS_CUTLASS_NVFP4_SSM_OUT",
333 [false, false, false, false, true],
334 ),
335 ];
336 for (name, expected) in cases {
337 let d = resolve(&[(name, "1")]);
338 assert_eq!(
339 [
340 d.cutlass_nvfp4_qkvz,
341 d.cutlass_nvfp4_attn_q,
342 d.cutlass_nvfp4_attn_kv,
343 d.cutlass_nvfp4_attn_o,
344 d.cutlass_nvfp4_ssm_out,
345 ],
346 expected,
347 "{name} must not enable a neighboring projection"
348 );
349 }
350 }
351
352 #[test]
353 fn non_nvfp4_flags_map_independently_and_single_scale_is_inverted() {
354 let cases = [
355 ("ATLAS_CUBLAS_GEMM", [true, false, false]),
356 ("ATLAS_CUBLAS_FP8", [false, true, false]),
357 ("ATLAS_CUTLASS_GEMM", [false, false, true]),
358 ];
359 for (name, expected) in cases {
360 let d = resolve(&[(name, "1")]);
361 assert_eq!(
362 [d.cublas.any(), d.cublas_fp8, d.cutlass_gemm],
363 expected,
364 "{name} must not enable a neighboring GEMM path"
365 );
366 assert!(d.fp8_blockscaled_prefill);
367 }
368 assert!(!resolve(&[("ATLAS_FP8_SINGLE_SCALE", "1")]).fp8_blockscaled_prefill);
369 }
370
371 // ───────────────── ATLAS_CUBLAS_GEMM scope grammar ─────────────────
372
373 fn scope(raw: &str) -> CublasScope {
374 resolve(&[("ATLAS_CUBLAS_GEMM", raw)]).cublas
375 }
376
377 /// The whole table, in one place, as the doc comment on [`CublasScope`]
378 /// states it. `1`/`true` keep meaning "every arm" so a pre-2026-09-11
379 /// launch script is unchanged; every other spelling is new.
380 #[test]
381 fn the_scope_grammar_maps_each_spelling_to_its_family_set() {
382 let f = |ffn, attn, ssm, head| CublasScope {
383 ffn,
384 attn,
385 ssm,
386 head,
387 };
388 let cases: [(&str, CublasScope); 13] = [
389 ("all", CublasScope::ALL),
390 ("1", CublasScope::ALL),
391 ("true", CublasScope::ALL),
392 ("ALL", CublasScope::ALL),
393 ("off", CublasScope::OFF),
394 ("0", CublasScope::OFF),
395 ("false", CublasScope::OFF),
396 ("", CublasScope::OFF),
397 ("ffn", f(true, false, false, false)),
398 ("attn", f(false, true, false, false)),
399 ("ssm", f(false, false, true, false)),
400 ("head", f(false, false, false, true)),
401 ("ffn,attn", f(true, true, false, false)),
402 ];
403 for (raw, expected) in cases {
404 assert_eq!(scope(raw), expected, "ATLAS_CUBLAS_GEMM={raw:?}");
405 }
406 // Whitespace is an operator typing a list, not a new family.
407 assert_eq!(scope(" ffn , ssm "), f(true, false, true, false));
408 // Union, so `off` subtracts nothing — documented, and the alternative
409 // (a clearing token) makes the meaning depend on token order.
410 assert_eq!(scope("ffn,off"), f(true, false, false, false));
411 }
412
413 /// A typo must not widen the set. The whole point of the change is that
414 /// arming an unintended family costs 10.3 GiB of unledgered weight copies;
415 /// `ATLAS_CUBLAS_GEMM=fnn` resolving to `all` would reintroduce it.
416 #[test]
417 fn unknown_families_are_dropped_and_reported_never_widening_the_set() {
418 assert_eq!(scope("junk"), CublasScope::OFF);
419 assert_eq!(
420 scope("ffn,junk"),
421 CublasScope {
422 ffn: true,
423 ..CublasScope::OFF
424 }
425 );
426 let (resolved, unknown) = parse_cublas_scope(Some("ffn, FNN ,bogus"));
427 assert_eq!(
428 resolved,
429 CublasScope {
430 ffn: true,
431 ..CublasScope::OFF
432 }
433 );
434 assert_eq!(
435 unknown,
436 vec!["fnn".to_owned(), "bogus".to_owned()],
437 "the warning must name what the operator typed, lowercased"
438 );
439 }
440
441 /// An absent variable is not the same input as `off`, and both must land
442 /// on the empty set without allocating an "unknown token" for the caller
443 /// to warn about.
444 #[test]
445 fn an_absent_variable_resolves_to_the_empty_set_silently() {
446 assert_eq!(parse_cublas_scope(None), (CublasScope::OFF, Vec::new()));
447 assert_eq!(
448 parse_cublas_scope(Some("off")),
449 (CublasScope::OFF, Vec::new())
450 );
451 assert!(!CublasScope::OFF.any());
452 assert!(CublasScope::ALL.any());
453 }
454
455 #[test]
456 fn w4a16_variants_accept_only_documented_spellings() {
457 for (value, expected) in [
458 ("v1", 1),
459 ("v2", 2),
460 ("v3", 3),
461 ("1", 0),
462 ("V1", 0),
463 ("unknown", 0),
464 ] {
465 assert_eq!(
466 resolve(&[("ATLAS_W4A16_VARIANT", value)]).w4a16_variant,
467 expected,
468 "value {value}"
469 );
470 }
471 }
472
473 #[test]
474 fn an_unknown_projection_label_falls_back_to_the_umbrella_flag() {
475 assert!(!GemmDispatch::defaults().cutlass_nvfp4_attn_qkv("mystery"));
476 let d = GemmDispatch {
477 cutlass_nvfp4_gemm: true,
478 ..GemmDispatch::defaults()
479 };
480 assert!(d.cutlass_nvfp4_attn_qkv("mystery"));
481 }
482}