spark_model/layers/ops/
dispatch_helpers.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GEMM-path dispatch helpers + roofline instrumentation. Extracted from the
4//! `ops` module root during the ≤500-line split. Re-exported at
5//! `crate::layers::ops::*` via `ops.rs`.
6
7#![allow(unused_imports)]
8
9use super::*;
10
11// The nine GEMM-path flags that lived here as `OnceLock<bool>` statics are now
12// `layers::ops::GemmDispatch`, resolved once when the model is built and
13// carried on `ForwardContext`. A static outlived the model whose flags it
14// encoded — swap to a model with different levers and the process kept serving
15// the previous model's dispatch decisions, silently. It also hid the
16// dependency: a function reading the environment through a static takes no
17// argument that says so and gives the compiler nothing to check.
18
19use spark_runtime::gpu::GpuBackend;
20
21// The two BATCHED-PREFILL ADMISSION flags below are not GEMM-path dispatch and
22// have no `GemmDispatch` field; they gate whether concurrent prefills co-admit
23// into one forward. They stay env reads for now (flag→lever conversion is
24// per-PR follow-up work, tracked in the integration notes).
25
26/// Whether chunk-zero streams may use the paged batched-prefill path.
27///
28/// `ATLAS_PREFILL_CODISPATCH` is the end-to-end request-admission flag;
29/// keep the older Q12 spelling as a compatibility alias for existing recipes.
30pub fn prefill_batched_first_chunk_enabled() -> bool {
31    prefill_batched_first_chunk_from_values([
32        std::env::var("ATLAS_Q12_BATCHED_FIRST_CHUNK")
33            .ok()
34            .as_deref(),
35        std::env::var("ATLAS_PREFILL_CODISPATCH").ok().as_deref(),
36    ])
37}
38
39fn prefill_batched_first_chunk_from_values(values: [Option<&str>; 2]) -> bool {
40    values.into_iter().any(bool_value_enabled)
41}
42
43/// The resolved VARLEN batched-prefill decision. One cell, three readers
44/// (admission predicate, batched-attention chunk-0 guard, scheduler wave
45/// planner) — a `OnceLock` so the decision cannot change mid-serve.
46static PREFILL_VARLEN: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
47
48/// Publish the command line's `--prefill-varlen-batch` decision. Returns the
49/// value IN FORCE, which differs from `enabled` when something already
50/// resolved the cell (then the command line did NOT take effect — the caller
51/// warns, mirroring `gdn_flags::set_from_cli`). Absent flag ⇒ never called ⇒
52/// the documented `ATLAS_PREFILL_VARLEN` fallback stays reachable.
53pub fn set_prefill_varlen_from_cli(enabled: bool) -> bool {
54    let _ = PREFILL_VARLEN.set(enabled);
55    *PREFILL_VARLEN.get().expect("just set")
56}
57
58/// VARLEN (ragged) batched prefill enabled? (`--prefill-varlen-batch`,
59/// legacy `ATLAS_PREFILL_VARLEN=1`; default OFF).
60///
61/// SSOT for the admission predicate (`check_kernel_batched_eligible`), the
62/// batched-attention layer's chunk-0 guard, and the scheduler's prefill wave
63/// planner. Those must agree: if admission accepts a batch the layer then
64/// rejects, the bail happens mid-Phase-A with streams already mutated, and
65/// the per-stream fallback re-runs setup on dirty state.
66pub fn prefill_varlen_enabled() -> bool {
67    *PREFILL_VARLEN
68        .get_or_init(|| bool_value_enabled(std::env::var("ATLAS_PREFILL_VARLEN").ok().as_deref()))
69}
70
71fn bool_value_enabled(value: Option<&str>) -> bool {
72    matches!(value, Some("1")) || value.is_some_and(|value| value.eq_ignore_ascii_case("true"))
73}
74
75pub fn log_cutlass_nvfp4_route(gpu: &dyn GpuBackend, name: &str, m: u32, n: u32, k: u32) {
76    // Routing telemetry, not a warning: the dedup key includes M, and
77    // prefill produces a new M per token count, so at WARN this spammed the
78    // production channel on every agentic request (and a polluted WARN
79    // stream misdirects real investigations). Skip the dedup probe entirely
80    // unless a subscriber would take the debug event — this runs per routed
81    // GEMM call.
82    if !tracing::enabled!(tracing::Level::DEBUG) {
83        return;
84    }
85    // De-duplicated on the BACKEND (`OpCache::first_shape`), not in a static:
86    // the shapes a model dispatches are its own, and a process-wide set
87    // suppresses the first route line for every shape a previous model
88    // happened to use — the lines that say which kernel this model took.
89    if gpu.op_cache().first_shape(name, m, n, k) {
90        tracing::debug!("CUTLASS_NVFP4_ROUTE {name} M={m} N={n} K={k}");
91    }
92}
93
94/// Roofline instrumentation: log each unique (kernel, M, N, K) GEMM shape once,
95/// gated by `ATLAS_GEMM_SHAPE_LOG=1`. Used to cross-reference nsys per-call
96/// durations → achieved TFLOPS/bandwidth vs GB10 peak.
97pub fn log_gemm_shape(gpu: &dyn GpuBackend, name: &str, m: u32, n: u32, k: u32) {
98    if std::env::var("ATLAS_GEMM_SHAPE_LOG").ok().as_deref() != Some("1") {
99        return;
100    }
101    if gpu.op_cache().first_shape(name, m, n, k) {
102        let flop = 2.0 * m as f64 * n as f64 * k as f64;
103        tracing::warn!("GEMM_SHAPE {name} M={m} N={n} K={k} FLOP={flop:.3e}");
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::{bool_value_enabled, prefill_batched_first_chunk_from_values};
110
111    #[test]
112    fn accepts_boolean_environment_spellings() {
113        assert!(bool_value_enabled(Some("1")));
114        assert!(bool_value_enabled(Some("true")));
115        assert!(bool_value_enabled(Some("TRUE")));
116        assert!(!bool_value_enabled(Some("0")));
117        assert!(!bool_value_enabled(Some("false")));
118        assert!(!bool_value_enabled(None));
119    }
120
121    #[test]
122    fn either_chunk_zero_spelling_enables_admission() {
123        assert!(prefill_batched_first_chunk_from_values([Some("1"), None]));
124        assert!(prefill_batched_first_chunk_from_values([
125            None,
126            Some("true")
127        ]));
128        assert!(!prefill_batched_first_chunk_from_values([None, None]));
129        assert!(!prefill_batched_first_chunk_from_values([
130            Some("0"),
131            Some("false")
132        ]));
133    }
134}
spark_model/layers/ops/dispatch_helpers.rs

spark_model/layers/ops/
dispatch_helpers.rs
