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}