metrale_model_arch/
seq_state_reserve.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Per-sequence device state the serve must reserve, computed from
4//! the config before the model exists.
5//!
6//! The server's `preflight_reserve` adds this charge (`per_sequence_reserve`)
7//! before the model is built, when there is no layer list, proposer or
8//! `GpuBackend`, so it is a function of the config alone, like the
9//! `ssm_reserve` terms.
10//!
11//! Charged, per sequence:
12//!   * the target stack's DSA indexer caches (`Glm5NextDsaState`), one per text DSA layer;
13//!   * the draft proposer's own state (what `Glm5NextMtpHead`'s `alloc_state` allocates),
14//!     in its own field: the proposer is not a `TransformerLayer` and is not in the layer list.
15//!
16//! Not charged here:
17//!   * KDA recurrent and conv state: pool-owned (`ssm_reserve::PoolPlan`);
18//!     `meta.rs` gives pool-backed linear-attention layers pool addresses instead of calling
19//!     `alloc_state`;
20//!   * the paged KV pool: this reserve is subtracted from the KV budget.
21//!
22//! The result is per-rank bytes; nothing here divides by a world size.
23//!
24//! Owner: model-arch (serve memory reserve).
25//! Invariants:
26//! - Any `model_type` other than `glm5_next` is charged zero.
27
28use anyhow::Result;
29use metrale_config::{LayerType, ModelConfig};
30
31use crate::glm5next_dsa::state::{dsa_capacity, indexer_state_bytes};
32use crate::glm5next_skeleton::{Glm5NextTextSkeleton, Mixer};
33
34/// 2026-09-25: Per-sequence, per-rank device state, split by owner.
35#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
36pub struct PerSequenceState {
37    /// 2026-09-25: Indexer caches owned by the target stack's DSA layers.
38    pub target_layers: usize,
39    /// 2026-09-25: State owned by the draft proposer (a `DraftProposer`, not a `TransformerLayer`).
40    pub proposer: usize,
41}
42
43impl PerSequenceState {
44    pub fn total(&self) -> usize {
45        self.target_layers + self.proposer
46    }
47
48    /// 2026-09-25: What the reserve must hold for `max_batch_size` concurrent
49    /// sequences (0 counts as 1).
50    pub fn for_batch(&self, max_batch_size: usize) -> usize {
51        self.total() * max_batch_size.max(1)
52    }
53}
54
55/// 2026-09-25: Per-sequence owned device state for `config`, at a context of `max_seq_len`
56/// tokens.
57///
58/// Returns zeros for any `model_type` other than `glm5_next`. The proposer term needs
59/// `spec_on` and a sparse-attention MTP layer. Errors when the GLM skeleton cannot be built
60/// from `config`.
61pub fn per_sequence_state_bytes(
62    config: &ModelConfig,
63    max_seq_len: usize,
64    spec_on: bool,
65) -> Result<PerSequenceState> {
66    if config.model_type != "glm5_next" {
67        return Ok(PerSequenceState::default());
68    }
69    let capacity = dsa_capacity(max_seq_len, config.index_kpool);
70    let per_layer = indexer_state_bytes(capacity, config.index_head_dim);
71
72    // 2026-09-25: The skeleton, built from the config alone, says which layers carry a DSA
73    // mixer. The per-layer bytes use the pool-rounded capacity (`dsa_capacity`), the same
74    // capacity `Glm5NextDsaState::alloc` allocates.
75    let skeleton = Glm5NextTextSkeleton::from_config(config)?;
76    let dsa_layers = skeleton
77        .layers
78        .iter()
79        .filter(|l| l.mixer == Mixer::Dsa)
80        .count();
81
82    let target_layers = dsa_layers * per_layer;
83
84    // 2026-09-25: The GLM MTP head's `alloc_state` (`glm5next_mtp_head/proposer.rs`)
85    // allocates one DSA indexer block for its drafter layer and five buffers: `concat`,
86    // `x`, `logits`, `arg`, `head_xchg`, in that order below. Its `block_table` is a host
87    // `Vec<u32>`, not device memory.
88    let proposer = if spec_on && config.mtp_layer_types.contains(&LayerType::SparseAttention) {
89        per_layer
90            + 2 * config.hidden_size * 2
91            + config.hidden_size * 2
92            + config.vocab_size * 2
93            + 4
94            + 16
95    } else {
96        0
97    };
98
99    Ok(PerSequenceState {
100        target_layers,
101        proposer,
102    })
103}
104
105#[cfg(test)]
106#[path = "seq_state_reserve_tests.rs"]
107mod tests;
metrale_model_arch/seq_state_reserve.rs

metrale_model_arch/
seq_state_reserve.rs
