metrale_gpu_runtime/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: GPU runtime: the `GpuBackend` trait with its CUDA, Metal and mock
4//! implementations, device buffers, streams, the kernel registry, the cuBLASLt,
5//! CUTLASS and FlashInfer bridges, and the process-wide SSM tail-snapshot and
6//! hermetic switches.
7//!
8//! Owner: gpu-runtime.
9//! Invariants: none beyond the types.
10
11#![deny(warnings)]
12#![deny(clippy::all)]
13
14pub mod buffers;
15#[cfg(feature = "cuda")]
16pub mod cublaslt;
17// 2026-09-25: Without the `cuda` feature these modules come from the
18// `*_metal_stub.rs` files, which keep the entry points metrale-model-layers calls
19// without a `cfg`; their GPU entry points are `unreachable!`.
20#[cfg(not(feature = "cuda"))]
21#[path = "cublaslt_metal_stub.rs"]
22pub mod cublaslt;
23#[cfg(feature = "cuda")]
24pub mod cuda_backend;
25#[cfg(feature = "cuda")]
26pub mod cutlass;
27#[cfg(not(feature = "cuda"))]
28#[path = "cutlass_metal_stub.rs"]
29pub mod cutlass;
30#[cfg(feature = "cuda")]
31pub mod flashinfer;
32#[cfg(not(feature = "cuda"))]
33#[path = "flashinfer_metal_stub.rs"]
34pub mod flashinfer;
35pub mod gpu;
36#[path = "gpu_args.rs"]
37mod gpu_args;
38pub mod host_heap;
39pub mod kernel_args;
40#[cfg(feature = "metal")]
41pub mod metal_backend;
42pub mod op_cache;
43pub mod pinned_hosts;
44pub mod timing;
45
46/// 2026-09-25: The last multiple of `block_size` strictly below `total_tokens`, or
47/// `None` when `block_size` is 0 or `total_tokens <= block_size`.
48///
49/// The SSM tail snapshot is placed here: in-pass by `prepare_midchunk_capture`
50/// (metrale-model-engine), or, with `ssm_tail_ckpt_enabled` on and
51/// `ssm_tail_midchunk_enabled` off, by ending a prefill chunk here
52/// (`run_standard` and `run_batched_mixed` in metrale-server).
53pub fn ssm_tail_boundary(total_tokens: usize, block_size: usize) -> Option<usize> {
54    if block_size == 0 || total_tokens <= block_size {
55        return None;
56    }
57    let boundary = ((total_tokens - 1) / block_size) * block_size;
58    (boundary > 0).then_some(boundary)
59}
60
61/// 2026-09-25: Tail-checkpoint chunk split, on only with `METRALE_SSM_TAIL_CKPT=1`
62/// (see `ssm_tail_boundary`). Resolved on first call and cached.
63///
64/// Off by default: an A/B measured 2026-07-10 (174 samples per arm) found it
65/// perf-neutral, because the extra chunk split (median 868 ms) cancelled the SSM
66/// replay it saved (about 1374 ms).
67pub fn ssm_tail_ckpt_enabled() -> bool {
68    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
69    *ON.get_or_init(|| matches!(std::env::var("METRALE_SSM_TAIL_CKPT").as_deref(), Ok("1")))
70}
71
72/// 2026-09-25: Publish `met serve --hermetic`, before the first `hermetic_enabled`
73/// call, which caches its answer. Only `true` is published, so without the flag
74/// `METRALE_HERMETIC=1` still decides.
75pub fn set_hermetic(on: bool) {
76    if on {
77        let _ = HERMETIC.set(true);
78    }
79}
80
81static HERMETIC: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
82
83/// 2026-09-25: Whether hermetic mode is on: `set_hermetic(true)` was called or
84/// `METRALE_HERMETIC=1`. Resolved on first call and cached; read by the SSM
85/// snapshot lookups in metrale-cache (`snapshot.rs`, `snapshot_tier.rs`).
86pub fn hermetic_enabled() -> bool {
87    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
88    *ON.get_or_init(|| {
89        HERMETIC
90            .get()
91            .copied()
92            .unwrap_or_else(|| matches!(std::env::var("METRALE_HERMETIC").as_deref(), Ok("1")))
93    })
94}
95
96/// 2026-09-25: Publish `--no-ssm-tail-midchunk`: the serve plan passes
97/// `Some(false)` when the flag is given and `None` otherwise. `None` publishes
98/// nothing, so `METRALE_SSM_TAIL_MIDCHUNK` still decides. The cell is set once:
99/// the first published value, or the first `ssm_tail_midchunk_enabled` read, wins.
100pub fn set_ssm_tail_midchunk(on: Option<bool>) {
101    if let Some(on) = on {
102        let _ = SSM_TAIL_MIDCHUNK.set(on);
103    }
104}
105
106static SSM_TAIL_MIDCHUNK: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
107
108/// 2026-09-25: Mid-chunk SSM tail capture: on unless `--no-ssm-tail-midchunk` or
109/// `METRALE_SSM_TAIL_MIDCHUNK=0`. While it is on the scheduler does not end chunks
110/// at `ssm_tail_boundary`, and `prepare_midchunk_capture` (metrale-model-engine)
111/// plans an in-pass capture there, which it does only on `metrale_scale` builds.
112/// Resolved on first call and cached.
113pub fn ssm_tail_midchunk_enabled() -> bool {
114    *SSM_TAIL_MIDCHUNK.get_or_init(|| {
115        !matches!(
116            std::env::var("METRALE_SSM_TAIL_MIDCHUNK").as_deref(),
117            Ok("0")
118        )
119    })
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn an_absent_flag_does_not_seal_the_midchunk_cell() {
128        // 2026-09-25: `None` must leave the cell unset so a later `Some` applies.
129        // The cell is process-global with no reset, so no other test in this
130        // crate may touch it.
131        for _ in 0..3 {
132            set_ssm_tail_midchunk(None);
133        }
134        set_ssm_tail_midchunk(Some(false));
135        assert!(
136            !ssm_tail_midchunk_enabled(),
137            "an absent flag must leave the cell open for the next writer"
138        );
139        set_ssm_tail_midchunk(Some(true));
140        assert!(!ssm_tail_midchunk_enabled(), "and a SET one is final");
141    }
142
143    #[test]
144    fn the_tail_boundary_is_the_last_block_strictly_below_the_prompt() {
145        // 2026-09-25: No boundary gives `None`, not 0: a snapshot at token 0
146        // would restore nothing.
147        assert_eq!(ssm_tail_boundary(0, 16), None);
148        assert_eq!(ssm_tail_boundary(16, 16), None, "not the prompt's own end");
149        assert_eq!(ssm_tail_boundary(17, 16), Some(16));
150        assert_eq!(ssm_tail_boundary(32, 16), Some(16), "strictly below");
151        assert_eq!(ssm_tail_boundary(33, 16), Some(32));
152        assert_eq!(ssm_tail_boundary(100, 0), None, "no division by zero");
153    }
154}
155
156#[cfg(feature = "cuda")]
157pub mod cuda_host;
158#[cfg(feature = "cuda")]
159pub mod kernel;
160pub mod mlx_int8;
161#[cfg(feature = "cuda")]
162pub mod registry;
163#[cfg(feature = "cuda")]
164pub mod stream;
metrale_gpu_runtime/lib.rs

metrale_gpu_runtime/
lib.rs
