metrale_model_arch/glm5next_skeleton/
budget.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `StateBudget` (per-sequence state bytes) and `StructuralAccounting` (checkpoint names against the skeleton's contract).
4//!
5//! Owner: model-arch (GLM-5.3).
6//! Invariants: none beyond the types.
7
8/// 2026-09-25: Per-sequence state bytes, built by `Glm5NextTextSkeleton::state_budget` from
9/// the config's geometry. `per_rank` divides only the KDA state and the DSA KV.
10#[derive(Debug, Clone, Copy)]
11pub struct StateBudget {
12    /// 2026-09-25: KDA recurrent state, `[heads, head_dim, head_dim]` fp32 per KDA layer,
13    /// summed over the KDA layers. Does not grow with sequence length.
14    pub kda_recurrent: usize,
15    /// 2026-09-25: KDA causal-conv window, bf16, `conv_dim x (kernel - 1 + num_spec)` per KDA
16    /// layer. Does not grow with sequence length.
17    pub kda_conv: usize,
18    /// 2026-09-25: DSA MLA KV per token across the text-stack DSA layers, `kv_lora_rank` x 2
19    /// bytes each. Grows with sequence length.
20    pub dsa_kv_per_token: usize,
21    /// 2026-09-25: Indexer key and gate state per token across the text-stack DSA layers.
22    /// Grows with sequence length. The indexer projections are `DsaShard::Replicated`
23    /// (`glm5next_dsa/tp.rs`), so every rank holds the whole cache and [`Self::per_rank`]
24    /// does not divide it.
25    pub dsa_indexer_per_token: usize,
26    /// 2026-09-25: mHC highway per token, `hc_mult x hidden` fp32. An activation, so neither
27    /// `fixed` nor `per_token` counts it.
28    pub mhc_highway_per_token: usize,
29    /// 2026-09-25: MoE routing scratch per token: f32 logits plus 4-byte top-k ids and weights.
30    pub moe_routing_per_token: usize,
31}
32
33impl StateBudget {
34    /// 2026-09-25: Fixed (sequence-length-independent) bytes per sequence.
35    pub fn fixed(&self) -> usize {
36        self.kda_recurrent + self.kda_conv
37    }
38    /// 2026-09-25: Bytes that grow with every token of context.
39    pub fn per_token(&self) -> usize {
40        self.dsa_kv_per_token + self.dsa_indexer_per_token
41    }
42    /// 2026-09-25: Total persistent state for a sequence of `tokens`.
43    pub fn for_sequence(&self, tokens: usize) -> usize {
44        self.fixed() + tokens * self.per_token()
45    }
46    /// 2026-09-25: Divides `kda_recurrent`, `kda_conv` and `dsa_kv_per_token` by `ep`; the
47    /// indexer cache, the mHC highway and the routing scratch stay whole.
48    pub fn per_rank(&self, ep: usize) -> StateBudget {
49        StateBudget {
50            kda_recurrent: self.kda_recurrent / ep,
51            kda_conv: self.kda_conv / ep,
52            dsa_kv_per_token: self.dsa_kv_per_token / ep,
53            // 2026-09-25: Not divided: the indexer projections are `DsaShard::Replicated`, so
54            // every rank holds the whole cache.
55            dsa_indexer_per_token: self.dsa_indexer_per_token,
56            mhc_highway_per_token: self.mhc_highway_per_token,
57            moe_routing_per_token: self.moe_routing_per_token,
58        }
59    }
60}
61
62#[derive(Debug)]
63pub struct StructuralAccounting {
64    pub required: usize,
65    pub bound: usize,
66    /// 2026-09-25: Structural tensors the checkpoint does not have. `is_complete` requires
67    /// it empty.
68    pub missing: Vec<String>,
69    /// 2026-09-25: Names that are neither structural nor MLP/vision. `is_complete` requires
70    /// it empty.
71    pub unexpected: Vec<String>,
72    /// 2026-09-25: Count of MLP/MoE and vision names, which the skeleton does not bind.
73    pub deferred: usize,
74}
75
76impl StructuralAccounting {
77    pub fn is_complete(&self) -> bool {
78        self.missing.is_empty() && self.unexpected.is_empty() && self.bound == self.required
79    }
80}
metrale_model_arch/glm5next_skeleton/budget.rs

metrale_model_arch/glm5next_skeleton/
budget.rs
