spark_model/
seq_state_reserve.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Per-sequence device state the serve must RESERVE, computed before the model exists.
4//!
5//! # Why this is config-keyed and not a layer-trait sum
6//!
7//! `preflight_reserve(args, config, free_mem)` runs long before `build_model`, and the model
8//! builder *consumes* the reserve it returns. There is no layer list, no proposer, and no
9//! `GpuBackend` at that point, so a per-sequence term cannot be `Σ_layers layer.some_method()`.
10//! Every sibling reserve term (`ssm_reserve::*`) is a free function over config-derived scalars
11//! for exactly this reason; this module follows that precedent.
12//!
13//! # What is charged, and what deliberately is not
14//!
15//! Charged — device memory a sequence OWNS, that no other reserve term covers:
16//!   * the target stack's DSA indexer caches (`Glm5NextDsaState`), one per text DSA layer;
17//!   * the draft proposer's own per-sequence state (`Glm5NextMtpHead::alloc_state`).
18//!
19//! The proposer is a SEPARATE OWNER — it is not a `TransformerLayer` and never appears in the
20//! layer list, so a layer-side sum cannot see it. It is returned separately and must not be
21//! folded into the target-layer term.
22//!
23//! NOT charged, each because something else already accounts for it:
24//!   * KDA recurrent + conv state — pool-owned, covered by `ssm_reserve::ssm_pool_reserve_bytes`
25//!     (and `meta.rs` never calls `alloc_state` for a pool-backed mixer);
26//!   * the paged KV pool — it *is* the KV budget this reserve is subtracted from; charging it
27//!     here would be circular;
28//!   * the buffer arena — sized by `max_batch_tokens`, not per sequence (`buffer_arena_bytes`);
29//!   * SSM snapshot / replay-ring / h-stage — already reserve terms, and the snapshot term
30//!     already scales by `max_batch_size`;
31//!   * CUDA-graph exec memory — per SLOT, not per sequence, and LRU-bounded;
32//!   * pad-row dummy states — transient per step, and GLM declines every batched path.
33//!
34//! 🔴 The indexer cache is REPLICATED across ranks. EP does not halve it, so this returns
35//! per-rank bytes directly and callers must not divide again.
36
37use anyhow::Result;
38use atlas_core::config::{LayerType, ModelConfig};
39
40use crate::layers::glm5next_dsa::state::{dsa_capacity, indexer_state_bytes};
41use crate::layers::glm5next_skeleton::{Glm5NextTextSkeleton, Mixer};
42
43/// Per-sequence, per-rank device state, split by OWNER.
44#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
45pub struct PerSequenceState {
46    /// Indexer caches owned by the target stack's DSA layers.
47    pub target_layers: usize,
48    /// State owned by the draft proposer (a `DraftProposer`, never a `TransformerLayer`).
49    pub proposer: usize,
50}
51
52impl PerSequenceState {
53    pub fn total(&self) -> usize {
54        self.target_layers + self.proposer
55    }
56
57    /// What the reserve must hold for `max_batch_size` concurrent sequences.
58    /// Multiplied here, once, so no caller can apply it twice.
59    pub fn for_batch(&self, max_batch_size: usize) -> usize {
60        self.total() * max_batch_size.max(1)
61    }
62}
63
64/// Per-sequence owned device state for `config`, at a context of `max_seq_len` tokens.
65///
66/// Returns zeros for any model that owns no per-sequence device state outside the pools —
67/// which is every model except GLM-5.3 today. `spec_on` gates the proposer term: with
68/// speculation off no proposer is constructed and no proposer state is ever allocated.
69pub fn per_sequence_state_bytes(
70    config: &ModelConfig,
71    max_seq_len: usize,
72    spec_on: bool,
73) -> Result<PerSequenceState> {
74    if config.model_type != "glm5_next" {
75        return Ok(PerSequenceState::default());
76    }
77    let capacity = dsa_capacity(max_seq_len, config.index_kpool);
78    let per_layer = indexer_state_bytes(capacity, config.index_head_dim);
79
80    // The skeleton is the SSOT for which layers carry a DSA mixer, and it builds from the
81    // config alone — no weights, no checkpoint, no GPU. `state_budget().dsa_indexer_per_token`
82    // is the same arithmetic per token; we multiply by the POOL-ROUNDED capacity rather than a
83    // raw token count, which is what keeps the reserve equal to the allocation (`dsa_capacity`).
84    let skeleton = Glm5NextTextSkeleton::from_config(config)?;
85    let dsa_layers = skeleton
86        .layers
87        .iter()
88        .filter(|l| l.mixer == Mixer::Dsa)
89        .count();
90
91    let target_layers = dsa_layers * per_layer;
92
93    // The GLM MTP head allocates, per sequence: one DSA indexer block for its own drafter
94    // layer, plus five small fixed buffers. `block_table` is a host `Vec<u32>`, not device
95    // memory, and the drafter's private KV pool is a construction-time singleton — neither
96    // belongs here. Mirrors `Glm5NextMtpHead::alloc_state`.
97    let proposer = if spec_on && config.mtp_layer_types.contains(&LayerType::SparseAttention) {
98        per_layer                              // drafter DSA indexer block, one layer
99            + 2 * config.hidden_size * 2       // concat
100            + config.hidden_size * 2           // x
101            + config.vocab_size * 2            // logits
102            + 4                                // arg
103            + 16 // head_xchg
104    } else {
105        0
106    };
107
108    Ok(PerSequenceState {
109        target_layers,
110        proposer,
111    })
112}
113
114#[cfg(test)]
115#[path = "seq_state_reserve_tests.rs"]
116mod tests;
spark_model/seq_state_reserve.rs

spark_model/
seq_state_reserve.rs
