metrale_model_layers/layers/ops/
dispatch_helpers.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Batched-prefill admission switches (codispatch, varlen) and GEMM
4//! route and shape logging.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - The codispatch and varlen decisions are each resolved at most once per
9//!   process, by whichever comes first: the command line or the first read
10//!   from the environment. Every later read returns that value.
11
12#![allow(unused_imports)]
13
14use super::*;
15
16use metrale_gpu_runtime::gpu::GpuBackend;
17
18// 2026-09-25: The codispatch and varlen switches gate batched prefill
19// admission, not a GEMM path, and have no `GemmDispatch` field.
20// `METRALE_Q12_BATCHED_FIRST_CHUNK` has no command-line flag.
21
22/// 2026-09-25: The resolved codispatch decision, set once so it cannot change
23/// mid-serve.
24static PREFILL_CODISPATCH: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25
26/// 2026-09-25: Publish the command line's `--prefill-codispatch` decision.
27/// Returns the value in force, which differs from `enabled` when the cell was
28/// already resolved; the command line then did not take effect. The serve
29/// calls this only when the flag is given, so without it the
30/// `METRALE_PREFILL_CODISPATCH` fallback applies.
31pub fn set_prefill_codispatch_from_cli(enabled: bool) -> bool {
32    let _ = PREFILL_CODISPATCH.set(enabled);
33    *PREFILL_CODISPATCH.get().expect("just set")
34}
35
36/// 2026-09-25: Whether cross-request co-dispatch of fresh prompts is on:
37/// `--prefill-codispatch`, else `METRALE_PREFILL_CODISPATCH` (`1` or `true`,
38/// any case); off when neither is set. The scheduler's levers and
39/// [`prefill_batched_first_chunk_enabled`] both read it.
40pub fn prefill_codispatch_enabled() -> bool {
41    *PREFILL_CODISPATCH.get_or_init(|| {
42        bool_value_enabled(std::env::var("METRALE_PREFILL_CODISPATCH").ok().as_deref())
43    })
44}
45
46/// 2026-09-25: Whether chunk-zero streams may use the paged batched-prefill
47/// path: codispatch on, or `METRALE_Q12_BATCHED_FIRST_CHUNK` set to `1` or
48/// `true`. The scheduler reads codispatch alone.
49pub fn prefill_batched_first_chunk_enabled() -> bool {
50    prefill_batched_first_chunk_from_parts(
51        prefill_codispatch_enabled(),
52        std::env::var("METRALE_Q12_BATCHED_FIRST_CHUNK")
53            .ok()
54            .as_deref(),
55    )
56}
57
58/// 2026-09-25: The pure OR, testable without the process-wide cell or the
59/// environment. `codispatch` is already resolved; `q12` is the variable's raw
60/// value.
61fn prefill_batched_first_chunk_from_parts(codispatch: bool, q12: Option<&str>) -> bool {
62    codispatch || bool_value_enabled(q12)
63}
64
65/// 2026-09-28: Whether a batched prefill wave may carry chunk-zero streams: the
66/// first-chunk lever ([`prefill_batched_first_chunk_enabled`]) or varlen batching
67/// ([`prefill_varlen_enabled`]). The single rule for the model engine's admission
68/// (`kernel_batched_eligible`), its paged metadata upload, and the attention layer's
69/// routing; before 2026-09-28 the layer read the lever alone and refused a
70/// varlen-admitted chunk-zero wave after admission, which failed every stream in it.
71pub fn prefill_batched_chunk_zero_admitted() -> bool {
72    batched_chunk_zero_admitted(
73        prefill_batched_first_chunk_enabled(),
74        prefill_varlen_enabled(),
75    )
76}
77
78/// 2026-09-28: The pure rule behind [`prefill_batched_chunk_zero_admitted`].
79pub fn batched_chunk_zero_admitted(first_chunk_lever: bool, varlen: bool) -> bool {
80    first_chunk_lever || varlen
81}
82
83/// 2026-09-25: The resolved varlen decision, set once so it cannot change
84/// mid-serve.
85static PREFILL_VARLEN: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
86
87/// 2026-09-25: Publish the command line's `--prefill-varlen-batch` decision.
88/// Returns the value in force, which differs from `enabled` when the cell was
89/// already resolved; the command line then did not take effect. The serve
90/// calls this only when the flag is given, so without it the
91/// `METRALE_PREFILL_VARLEN` fallback applies.
92pub fn set_prefill_varlen_from_cli(enabled: bool) -> bool {
93    let _ = PREFILL_VARLEN.set(enabled);
94    *PREFILL_VARLEN.get().expect("just set")
95}
96
97/// 2026-09-25: Whether varlen (ragged) batched prefill is on:
98/// `--prefill-varlen-batch`, else `METRALE_PREFILL_VARLEN` (`1` or `true`, any
99/// case); off when neither is set. The admission predicate
100/// (`check_kernel_batched_eligible`), the batched attention layer and the
101/// scheduler's levers all read this one value.
102pub fn prefill_varlen_enabled() -> bool {
103    *PREFILL_VARLEN
104        .get_or_init(|| bool_value_enabled(std::env::var("METRALE_PREFILL_VARLEN").ok().as_deref()))
105}
106
107fn bool_value_enabled(value: Option<&str>) -> bool {
108    matches!(value, Some("1")) || value.is_some_and(|value| value.eq_ignore_ascii_case("true"))
109}
110
111pub fn log_cutlass_nvfp4_route(gpu: &dyn GpuBackend, name: &str, m: u32, n: u32, k: u32) {
112    // 2026-09-25: Debug level, because the dedup key includes M and every
113    // prefill length is a new M. The dedup probe is skipped when no subscriber
114    // takes debug events.
115    if !tracing::enabled!(tracing::Level::DEBUG) {
116        return;
117    }
118    // 2026-09-25: Deduplicated per backend by `OpCache::first_shape`.
119    if gpu.op_cache().first_shape(name, m, n, k) {
120        tracing::debug!("CUTLASS_NVFP4_ROUTE {name} M={m} N={n} K={k}");
121    }
122}
123
124/// 2026-09-25: With `METRALE_GEMM_SHAPE_LOG=1`, log each (kernel, M, N, K) GEMM
125/// shape once per backend at warn level, with its FLOP count.
126pub fn log_gemm_shape(gpu: &dyn GpuBackend, name: &str, m: u32, n: u32, k: u32) {
127    if std::env::var("METRALE_GEMM_SHAPE_LOG").ok().as_deref() != Some("1") {
128        return;
129    }
130    if gpu.op_cache().first_shape(name, m, n, k) {
131        let flop = 2.0 * m as f64 * n as f64 * k as f64;
132        tracing::warn!("GEMM_SHAPE {name} M={m} N={n} K={k} FLOP={flop:.3e}");
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::{bool_value_enabled, prefill_batched_first_chunk_from_parts};
139
140    #[test]
141    fn accepts_boolean_environment_spellings() {
142        assert!(bool_value_enabled(Some("1")));
143        assert!(bool_value_enabled(Some("true")));
144        assert!(bool_value_enabled(Some("TRUE")));
145        assert!(!bool_value_enabled(Some("0")));
146        assert!(!bool_value_enabled(Some("false")));
147        assert!(!bool_value_enabled(None));
148    }
149
150    /// 2026-09-25: Either spelling enables the path, and codispatch off does
151    /// not veto `METRALE_Q12_BATCHED_FIRST_CHUNK`.
152    #[test]
153    fn either_chunk_zero_spelling_enables_admission() {
154        assert!(prefill_batched_first_chunk_from_parts(true, None));
155        assert!(prefill_batched_first_chunk_from_parts(false, Some("true")));
156        assert!(prefill_batched_first_chunk_from_parts(false, Some("1")));
157        assert!(!prefill_batched_first_chunk_from_parts(false, None));
158        assert!(!prefill_batched_first_chunk_from_parts(false, Some("0")));
159        assert!(!prefill_batched_first_chunk_from_parts(
160            false,
161            Some("false")
162        ));
163        assert!(prefill_batched_first_chunk_from_parts(false, Some("1")));
164    }
165}
metrale_model_layers/layers/ops/dispatch_helpers.rs

metrale_model_layers/layers/ops/
dispatch_helpers.rs
