spark_runtime/lib.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3#![deny(warnings)]
4#![deny(clippy::all)]
5
6pub mod buffers;
7#[cfg(feature = "cuda")]
8pub mod cublaslt;
9// Metal/no-cuda builds get unreachable stubs so spark-model's unconditional
10// references to these cuda-only entry points still resolve (compile-only).
11#[cfg(not(feature = "cuda"))]
12#[path = "cublaslt_metal_stub.rs"]
13pub mod cublaslt;
14#[cfg(feature = "cuda")]
15pub mod cuda_backend;
16#[cfg(feature = "cuda")]
17pub mod cutlass;
18#[cfg(not(feature = "cuda"))]
19#[path = "cutlass_metal_stub.rs"]
20pub mod cutlass;
21#[cfg(unix)]
22pub mod fast_weights;
23#[cfg(feature = "cuda")]
24pub mod flashinfer;
25#[cfg(not(feature = "cuda"))]
26#[path = "flashinfer_metal_stub.rs"]
27pub mod flashinfer;
28pub mod gpu;
29#[path = "gpu_args.rs"]
30mod gpu_args;
31pub mod kernel_args;
32pub mod kernel_audit;
33pub mod kv_cache;
34pub mod kv_dequant;
35pub mod kv_spill;
36pub mod launch_trace;
37#[cfg(feature = "metal")]
38pub mod metal_backend;
39pub mod op_cache;
40pub mod pinned_hosts;
41pub mod prefix_cache;
42pub mod progress;
43pub mod radix_tree;
44pub mod run_metrics;
45pub mod sampler;
46pub mod weights;
47
48/// Last paged-KV block boundary strictly below `total_tokens`.
49///
50/// A warm multi-turn hit can never match past this point: the chat template's
51/// generation-prompt suffix (assistant header, and the empty `<think></think>`
52/// block emitted when thinking is disabled) is not reproduced when the next
53/// turn re-renders the *completed* assistant message, so the longest common
54/// prefix diverges inside the prompt's final block. `RadixTree::walk` then
55/// floors `matched_tokens` to this boundary. Placing an SSM snapshot here makes
56/// the next turn's restore exact (zero recurrence replay); without it the
57/// lookup falls back to the coarse `--ssm-checkpoint-interval` grid.
58///
59/// Returns `None` when the prompt is too short to have such a boundary.
60pub fn ssm_tail_boundary(total_tokens: usize, block_size: usize) -> Option<usize> {
61 if block_size == 0 || total_tokens <= block_size {
62 return None;
63 }
64 let boundary = ((total_tokens - 1) / block_size) * block_size;
65 (boundary > 0).then_some(boundary)
66}
67
68/// OPT-IN switch for the tail checkpoint (`ATLAS_SSM_TAIL_CKPT=1`).
69///
70/// Default OFF. The 3-traj A/B (2026-07-10, 174 samples/arm) showed it is
71/// perf-NEUTRAL: it removes the SSM replay on ~89% of warm turns (mean 254 -> 25
72/// tokens), but the prefill-chunk split needed to land a snapshot on
73/// `ssm_tail_boundary` costs a median 868 ms extra forward pass for a median of 8
74/// trailing tokens, which cancels the ~1374 ms of replay it saves. It becomes a
75/// clear win only once the SSM state can be captured MID-CHUNK (in the GDN prefill
76/// kernel) instead of via an extra pass. Until then it stays off by default and
77/// ungated for accuracy.
78pub fn ssm_tail_ckpt_enabled() -> bool {
79 matches!(std::env::var("ATLAS_SSM_TAIL_CKPT").as_deref(), Ok("1"))
80}
81
82/// Default-ON switch for MID-CHUNK tail SSM capture (opt-out `ATLAS_SSM_TAIL_MIDCHUNK=0`).
83///
84/// Default ON => mid-chunk capture fires on prefill passes spanning the
85/// block-floored matched-prefix boundary. When disabled, the prefill
86/// chunk is NOT clamped to `ssm_tail_boundary`; instead each GDN layer's
87/// recurrent (h_state) and conv (conv_state) kernels are split at the block-
88/// floored matched-prefix boundary and the @tb state is copied into a reserved
89/// Marconi snapshot slot in-pass, removing the ~868 ms extra forward pass the
90/// clamp-based `ATLAS_SSM_TAIL_CKPT` path costs.
91/// Publish the command line's `--ssm-tail-midchunk`. Call once, at serve time,
92/// before any prefill runs.
93///
94/// `None` means THE FLAG WAS NOT GIVEN, and is not the same as `Some(default)`.
95/// Publishing the clap default sealed this cell on every `spark serve`, which
96/// made the documented `ATLAS_SSM_TAIL_MIDCHUNK=0` opt-out a silent no-op — an
97/// operator could set it, see the flag echoed in the startup log, and get the
98/// opposite behaviour with nothing anywhere saying so. A knob that looks like an
99/// opt-out and is not costs more than no knob at all, so an absent flag now
100/// publishes nothing and leaves the environment fallback below to decide.
101pub fn set_ssm_tail_midchunk(on: Option<bool>) {
102 if let Some(on) = on {
103 let _ = SSM_TAIL_MIDCHUNK.set(on);
104 }
105}
106
107static SSM_TAIL_MIDCHUNK: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
108
109pub fn ssm_tail_midchunk_enabled() -> bool {
110 // Default ON (2026-07-19): mid-chunk GDN tail capture eliminates the warm-turn
111 // SSM replay (~1.17s component of warm TTFT) by capturing state in-pass at the
112 // block-floored matched-prefix boundary.
113 //
114 // ★ `--ssm-tail-midchunk` WINS when it is given, and only then. It used to
115 // win unconditionally — serve.rs published the clap default on every boot,
116 // sealing this cell before anything asked, so `ATLAS_SSM_TAIL_MIDCHUNK=0`
117 // did NOTHING under `spark serve` while still being documented as the
118 // opt-out. `set_ssm_tail_midchunk` now takes an `Option` and an absent flag
119 // publishes nothing, so the read below is live again for the CLI, for tests
120 // and for examples alike.
121 //
122 // ★ The 2026-07-19 validation did not cover what it appeared to. It read
123 // "BFCL e2e 1007/1007" — a COMPLETION count, not an accuracy score — and
124 // warm-TTFT, which is a timing signal. Neither can see a wrong recurrent
125 // state, and on NVIDIA the captured h_state was in fact never written at
126 // all (see `prepare_midchunk_capture`, which now refuses the plan off
127 // `atlas_scale`). "flag-off byte-identical" held; it just was not evidence
128 // that flag-ON was correct.
129 *SSM_TAIL_MIDCHUNK
130 .get_or_init(|| !matches!(std::env::var("ATLAS_SSM_TAIL_MIDCHUNK").as_deref(), Ok("0")))
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136
137 #[test]
138 fn an_absent_flag_does_not_seal_the_midchunk_cell() {
139 // The defect this shape fixes: `set_ssm_tail_midchunk(bool)` was called
140 // with the clap default on every `spark serve`, sealing the cell before
141 // anything read it — so `ATLAS_SSM_TAIL_MIDCHUNK=0` was documented,
142 // echoed back in the startup log, and inert.
143 //
144 // ★ The cell is process-global with no reset, so this is the only test
145 // in this crate that may touch it: a second one would be
146 // order-dependent on this.
147 for _ in 0..3 {
148 set_ssm_tail_midchunk(None);
149 }
150 set_ssm_tail_midchunk(Some(false));
151 assert!(
152 !ssm_tail_midchunk_enabled(),
153 "an absent flag must leave the cell open for the next writer"
154 );
155 set_ssm_tail_midchunk(Some(true));
156 assert!(!ssm_tail_midchunk_enabled(), "and a SET one is final");
157 }
158
159 #[test]
160 fn the_tail_boundary_is_the_last_block_strictly_below_the_prompt() {
161 // `None` where no such boundary exists, rather than 0 — a snapshot at
162 // token 0 is not a cheap restore, it is a full replay wearing one.
163 assert_eq!(ssm_tail_boundary(0, 16), None);
164 assert_eq!(ssm_tail_boundary(16, 16), None, "not the prompt's own end");
165 assert_eq!(ssm_tail_boundary(17, 16), Some(16));
166 assert_eq!(ssm_tail_boundary(32, 16), Some(16), "strictly below");
167 assert_eq!(ssm_tail_boundary(33, 16), Some(32));
168 assert_eq!(ssm_tail_boundary(100, 0), None, "no division by zero");
169 }
170}