avarok_core/config/kv_completeness.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The two KV-completeness capability gates.
4//!
5//! Split out of `methods.rs` (the 500-LoC cap) when the GLM validation switch
6//! pushed that file over, and placed beside `kv_completeness_tests.rs` so the
7//! predicates and the tests that pin them sit together: a new model type must
8//! be taught to both gates, and forgetting one is the defect this file exists
9//! to make visible.
10
11use super::ModelConfig;
12
13impl ModelConfig {
14 /// Whether every byte of a sequence's per-layer state is represented by
15 /// its KV blocks.
16 ///
17 /// False for models whose PREFILL builds per-sequence state that KV pages
18 /// do not carry: GLM-5.3's DSA indexer rows (`Glm5NextDsaState`) and
19 /// compressed DeepSeek V4's compressor pool/ring. Every KV-only mechanism
20 /// โ radix prefix reuse and the `--swap-space-gb` spill image alike โ is
21 /// unsafe for those models, and this is the single fact both gates below
22 /// are asking about.
23 ///
24 /// ๐ `glm5_next` stays `false` even though the machinery to make it true now
25 /// exists. `Glm5NextLayer` implements the Marconi aux hooks
26 /// (`has_aux_state`/`snapshot_aux`/`restore_aux`, `layers/glm5next_dsa/aux_state.rs`),
27 /// so a snapshot CAN now carry the DSA indexer rows the KV pages do not โ and
28 /// KDA never needed carrying, because `uses_ssm_pool()` is true for it and
29 /// `SsmSnapshotPool` already captures that region device-to-device. What is
30 /// missing is not code, it is PROOF: the round trip is verified on CPU against
31 /// `MockGpuBackend` only, and nothing has yet shown that a warm-cache GLM-5.3
32 /// serve produces the same tokens as a cold one on real hardware. Flipping this
33 /// arm before that measurement exists would trade a correct-but-slow serve for
34 /// an unproven one. Flip it in the same change that lands the GPU evidence.
35 ///
36 /// ๐งช The measurement itself needs the arm OPEN, so
37 /// `kv_only_prefix_cache_is_safe` below carries an env-only validation switch
38 /// (`AVAROK_GLM53_PREFIX_CACHE_UNPROVEN`) that opens the prefix-cache arm and
39 /// nothing else. This `match` โ the compiled default and the swap-out answer โ
40 /// is untouched by it. `scripts/glm53-pc/` in the spark-bench repo is the
41 /// harness that turns the switch into the evidence.
42 fn per_sequence_state_is_kv_complete(&self) -> bool {
43 match self.model_type.as_str() {
44 "glm5_next" | "glm5_next_text" => false,
45 "deepseek_v4" => self.compress_ratios.iter().all(|&ratio| ratio == 0),
46 _ => true,
47 }
48 }
49
50 /// Whether the radix prefix cache captures every state needed to resume
51 /// this model exactly. Preflight SSOT for `build_prefix_cache`.
52 ///
53 /// ๐งช `AVAROK_GLM53_PREFIX_CACHE_UNPROVEN` (PRESENCE, house convention โ `=0`
54 /// is NOT "off") opens this arm for `glm5_next` so the GPU evidence the gate
55 /// above asks for can actually be MEASURED. It is a validation switch, not a
56 /// feature flag: the compiled default is unchanged, and unsetting the
57 /// variable is the whole rollback.
58 ///
59 /// ๐ด Scoped to the PREFIX-CACHE arm on purpose. `kv_only_swap_out_is_safe()`
60 /// deliberately does not consult it, because the `--swap-space-gb` image is
61 /// structurally incapable of carrying aux blobs โ `save_sequence_state_dispatch`
62 /// writes KV blocks plus `SsmLayerState` and calls no aux hook โ so a shared
63 /// override would unlock a path the DSA codec never reaches.
64 ///
65 /// Both call sites (`preflight_reserve` and `build_prefix_cache`) reach the
66 /// predicate through this method, so the reservation and the allocation cannot
67 /// disagree about whether the Marconi slots are live (ANOMALIES A68's class).
68 pub fn kv_only_prefix_cache_is_safe(&self) -> bool {
69 self.kv_only_prefix_cache_is_safe_with(Self::glm53_prefix_cache_validation_env())
70 }
71
72 /// Pure core of [`Self::kv_only_prefix_cache_is_safe`] (env-free, unit-testable
73 /// โ the `marconi_snapshot_slots_with` pattern, so no test has to mutate a
74 /// process-global variable that a sibling test is reading).
75 ///
76 /// The override is narrow by construction: it can only ever open `glm5_next`,
77 /// so a stray export cannot silently re-enable a compressed DeepSeek-V4 serve.
78 pub(crate) fn kv_only_prefix_cache_is_safe_with(&self, validation_override: bool) -> bool {
79 self.per_sequence_state_is_kv_complete()
80 || (validation_override
81 && matches!(self.model_type.as_str(), "glm5_next" | "glm5_next_text"))
82 }
83
84 /// The `AVAROK_GLM53_PREFIX_CACHE_UNPROVEN` validation switch (PRESENCE).
85 fn glm53_prefix_cache_validation_env() -> bool {
86 std::env::var_os("AVAROK_GLM53_PREFIX_CACHE_UNPROVEN").is_some()
87 }
88
89 /// Whether a sequence may be swapped out to the `--swap-space-gb` pool and
90 /// restored from it. Preflight SSOT for `resolve_swap_space_gb`.
91 ///
92 /// `save_sequence_state_dispatch` writes KV blocks plus the `SsmLayerState`
93 /// of each `LayerType::LinearAttention` layer, and nothing else; the
94 /// swap-out then calls `free_sequence`, which hands every remaining
95 /// per-layer state to #821's `release_state`. A model that is not
96 /// KV-complete therefore resumes with a freshly ZEROED pool behind a KV
97 /// image that assumes a populated one โ a silently wrong answer, not a
98 /// crash. Distinct from the prefix-cache predicate because they are
99 /// distinct guarantees; they happen to have the same answer today.
100 pub fn kv_only_swap_out_is_safe(&self) -> bool {
101 self.per_sequence_state_is_kv_complete()
102 }
103}