metrale_model_layers/ssm_reserve/
pool_plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The SSM state pool as one plan, which both the pre-load reserve (the server's
4//! `preflight_reserve`) and the allocation (`SsmStatePool::new`) read, so the two cannot
5//! differ (M5, LIFECYCLE-DESIGN.md section 3.3: "preflight == allocation").
6//!
7//! - [`pool_counts`]: how many units the pool holds: every slot plus the padding dummy, and
8//!   with speculation the verify slots (plus their dummy) with their h and conv intermediates
9//!   and one checkpoint each.
10//! - [`PoolPlan::new`]: the bytes, from `metrale_circuit::state::StatePlan` over the recurrent
11//!   states the model's circuit declares (`[[block.gdn.state]]`, `[[block.mamba.state]]`),
12//!   evaluated under this process's (tensor-parallel local) config.
13//!
14//! Owner: model-layers (SSM reserve).
15//! Invariants:
16//! - A model whose `model_type` has a circuit is sized from the circuit's declarations only.
17//!   A model without one is sized from `ModelConfig::ssm_h_state_bytes` /
18//!   `ssm_conv_state_bytes`: the transitional source, removed at M3 when every family has a
19//!   circuit. One source per family, never both.
20//! - The dummy slots are counted: the allocator allocates them, so the reserve does too.
21
22use std::collections::BTreeMap;
23
24use anyhow::{Context, Result, bail, ensure};
25use metrale_circuit::state::{
26    Holding, StateDecl, StateDtype, StateFormat, StateInputs, StateKind, StatePlan, VerifyInputs,
27    VerifySteps,
28};
29use metrale_config::ModelConfig;
30
31use super::{SsmRollbackMode, mtp_state_slots, verify_slot_h_intermediates};
32
33/// 2026-09-30: What decides the pool's unit counts.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct PoolShape {
36    /// 2026-09-30: Claimable slots (`--max-batch-size`).
37    pub max_slots: usize,
38    /// 2026-09-30: The verify pools exist (MTP, DFlash or self-speculation).
39    pub spec: bool,
40    /// 2026-09-30: Rows of the widest verify (`K`, the drafts plus one).
41    pub num_intermediates: usize,
42    /// 2026-09-30: Drafts per verify.
43    pub num_drafts: usize,
44    /// 2026-09-30: Every verify slot holds `K - 1` h intermediates (DFlash, or a K that is not
45    /// the MTP `num_drafts + 1`); otherwise the MTP ladder tiers them.
46    pub uniform_h: bool,
47    /// 2026-09-30: `--ssm-rollback-mode`.
48    pub rollback: SsmRollbackMode,
49}
50
51/// 2026-09-30: The verify pools' unit counts.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct VerifyCounts {
54    /// 2026-09-30: Per verify slot, the dummy last: h intermediates.
55    pub h_steps: Vec<usize>,
56    /// 2026-09-30: Conv intermediates per verify slot (0 under replay).
57    pub conv_steps: usize,
58}
59
60impl VerifyCounts {
61    /// 2026-09-30: Verify slots, the dummy included.
62    pub fn slots(&self) -> usize {
63        self.h_steps.len()
64    }
65}
66
67/// 2026-09-30: The pool's unit counts.
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct PoolCounts {
70    /// 2026-09-30: Live slots, the padding dummy (`max_slots`) included.
71    pub slots: usize,
72    /// 2026-09-30: The verify pools; `None` without speculation.
73    pub verify: Option<VerifyCounts>,
74}
75
76/// 2026-09-30: [`pool_counts_with`] with the verify slots `mtp_state_slots` and the per-slot
77/// h intermediates `verify_slot_h_intermediates` (the MTP ladder).
78///
79/// 2026-10-02: Under a published prompt-lookup copy tier (`copy_tier`) the copy slots hold
80/// the copy depth and the verify rows widen to it ([`pool_counts_tiered`]).
81pub fn pool_counts(shape: &PoolShape) -> PoolCounts {
82    pool_counts_tiered(shape, super::copy_tier())
83}
84
85/// 2026-10-02: [`pool_counts`] with the copy tier passed in, so tests need not publish it.
86/// The tier widens `num_intermediates` (every slot's conv intermediates and the dummy's h)
87/// and raises the copy slots' h; a uniform-h pool takes the widened depth everywhere.
88pub fn pool_counts_tiered(shape: &PoolShape, tier: Option<super::CopyTier>) -> PoolCounts {
89    let widened = PoolShape {
90        num_intermediates: super::tier_rows(shape.num_intermediates, tier),
91        ..*shape
92    };
93    pool_counts_with(&widened, mtp_state_slots(shape.max_slots), |s| {
94        super::tier_h(
95            verify_slot_h_intermediates(s, shape.num_drafts, false),
96            s,
97            tier,
98        )
99    })
100}
101
102/// 2026-09-30: Pure: the counts for `mtp_slots` verify slots, slot `s` tiered to `slot_h(s)`
103/// h intermediates (capped at `K - 1`). The verify dummy, at index `mtp_slots`, always gets
104/// `K - 1`: batched-verify padding rows may write any live token index. Replay keeps no
105/// intermediates, only the checkpoints.
106pub fn pool_counts_with(
107    shape: &PoolShape,
108    mtp_slots: usize,
109    slot_h: impl Fn(usize) -> usize,
110) -> PoolCounts {
111    let full = shape.num_intermediates.saturating_sub(1);
112    let verify = shape.spec.then(|| {
113        let replay = shape.rollback == SsmRollbackMode::Replay;
114        VerifyCounts {
115            h_steps: (0..=mtp_slots)
116                .map(|s| match (replay, s == mtp_slots || shape.uniform_h) {
117                    (true, _) => 0,
118                    (false, true) => full,
119                    (false, false) => slot_h(s).min(full),
120                })
121                .collect(),
122            conv_steps: if replay { 0 } else { shape.num_intermediates },
123        }
124    });
125    PoolCounts {
126        slots: shape.max_slots + 1,
127        verify,
128    }
129}
130
131/// 2026-09-30: Where a pool's per-unit sizes come from.
132#[derive(Debug, Clone, Copy, PartialEq, Eq)]
133pub enum UnitSource {
134    /// 2026-09-30: The recurrent states the model's circuit declares.
135    Circuit,
136    /// 2026-09-30: `ModelConfig::ssm_h_state_bytes` / `ssm_conv_state_bytes`, for a
137    /// `model_type` no circuit serves. Transitional: removed at M3.
138    Transitional,
139}
140
141/// 2026-09-30: The dims the circuits' recurrent state shapes read, from this process's config
142/// (after tensor-parallel sharding, `topology.rs`).
143pub fn state_dims(c: &ModelConfig) -> BTreeMap<String, u64> {
144    [
145        ("lin_k_heads", c.linear_num_key_heads),
146        ("lin_k_dim", c.linear_key_head_dim),
147        ("lin_v_heads", c.linear_num_value_heads),
148        ("lin_v_dim", c.linear_value_head_dim),
149        ("mamba_heads", c.mamba_num_heads),
150        ("mamba_head_dim", c.mamba_head_dim),
151        ("ssm_state", c.ssm_state_size),
152        ("ssm_groups", c.n_groups),
153        ("conv_kernel", c.linear_conv_kernel_dim),
154    ]
155    .into_iter()
156    .map(|(k, v)| (k.to_string(), v as u64))
157    .collect()
158}
159
160/// 2026-09-30: One recurrent layer's h and conv declarations, and where they came from.
161pub fn recurrent_units(c: &ModelConfig) -> Result<(Vec<StateDecl>, UnitSource)> {
162    if let Some(decls) = metrale_circuit::recurrent_states(&c.model_type, &state_dims(c))
163        .with_context(|| format!("the circuit of `{}` declares no usable state", c.model_type))?
164    {
165        ensure!(
166            !decls.is_empty() || c.num_ssm_layers() == 0,
167            "`{}` has {} recurrent layers and its circuit declares no recurrent state",
168            c.model_type,
169            c.num_ssm_layers()
170        );
171        return Ok((decls, UnitSource::Circuit));
172    }
173    let decl = |local: &str, bytes: usize, format, verify| StateDecl {
174        id: format!("transitional.{local}"),
175        local: local.to_string(),
176        block: "transitional".to_string(),
177        layer: None,
178        section: metrale_circuit::Section::Main,
179        kind: StateKind::Recurrent,
180        format,
181        elements: (bytes / 4) as u64,
182        verify: Some(verify),
183        lifetime: metrale_circuit::state::Lifetime::Sequence,
184    };
185    Ok((
186        vec![
187            decl(
188                "h",
189                c.ssm_h_state_bytes(),
190                StateFormat::Keyed("ssm_h_storage".into()),
191                VerifySteps::H,
192            ),
193            decl(
194                "conv",
195                c.ssm_conv_state_bytes(),
196                StateFormat::Fixed(StateDtype::F32),
197                VerifySteps::Conv,
198            ),
199        ],
200        UnitSource::Transitional,
201    ))
202}
203
204/// 2026-09-30: The SSM pool of one model, sized.
205#[derive(Debug, Clone)]
206pub struct PoolPlan {
207    /// 2026-09-30: Where the unit sizes came from.
208    pub source: UnitSource,
209    /// 2026-09-30: Recurrent layers.
210    pub layers: usize,
211    /// 2026-09-30: One h unit at FP32 width (the prefill staging blob).
212    pub h_f32_unit: usize,
213    /// 2026-09-30: One h unit as stored (half under `--ssm-h-dtype f16-pool`).
214    pub h_stored_unit: usize,
215    /// 2026-09-30: One conv unit.
216    pub conv_unit: usize,
217    /// 2026-09-30: The counts the plan was made for.
218    pub counts: PoolCounts,
219    layer: StatePlan,
220    h_id: String,
221    conv_id: String,
222}
223
224/// 2026-09-30: A recurrent state of the pool.
225#[derive(Debug, Clone, Copy, PartialEq, Eq)]
226pub enum PoolState {
227    H,
228    Conv,
229}
230
231impl PoolPlan {
232    /// 2026-09-30: The pool of `config` holding `counts`.
233    pub fn new(config: &ModelConfig, counts: &PoolCounts, h_f16_pool: bool) -> Result<Self> {
234        let (decls, source) = recurrent_units(config)?;
235        let find = |v: VerifySteps| {
236            decls
237                .iter()
238                .find(|d| d.verify == Some(v))
239                .with_context(|| format!("no recurrent state keeps {v:?} intermediates"))
240        };
241        let layers = config.num_ssm_layers();
242        if layers == 0 {
243            return Ok(Self::empty(source, counts));
244        }
245        let (h, conv) = (find(VerifySteps::H)?, find(VerifySteps::Conv)?);
246        let storage = if h_f16_pool {
247            StateDtype::F16
248        } else {
249            StateDtype::F32
250        };
251        let inputs = StateInputs {
252            formats: BTreeMap::from([("ssm_h_storage".to_string(), storage)]),
253            slots: counts.slots as u64,
254            verify: counts.verify.as_ref().map(|v| VerifyInputs {
255                h_steps: v.h_steps.iter().map(|&n| n as u64).collect(),
256                conv_steps: v.conv_steps as u64,
257            }),
258            kv: None,
259            draft_kv: None,
260        };
261        let layer = StatePlan::new(&decls, &inputs)?;
262        let unit = |d: &StateDecl, dtype: StateDtype| (d.elements * dtype.size()) as usize;
263        let conv_dtype = match conv.format {
264            StateFormat::Fixed(d) => d,
265            StateFormat::Keyed(_) => bail!("the conv state's format is not fixed"),
266        };
267        Ok(Self {
268            source,
269            layers,
270            h_f32_unit: unit(h, StateDtype::F32),
271            h_stored_unit: unit(h, storage),
272            conv_unit: unit(conv, conv_dtype),
273            counts: counts.clone(),
274            layer,
275            h_id: h.id.clone(),
276            conv_id: conv.id.clone(),
277        })
278    }
279
280    fn empty(source: UnitSource, counts: &PoolCounts) -> Self {
281        Self {
282            source,
283            layers: 0,
284            h_f32_unit: 0,
285            h_stored_unit: 0,
286            conv_unit: 0,
287            counts: counts.clone(),
288            layer: StatePlan { terms: Vec::new() },
289            h_id: String::new(),
290            conv_id: String::new(),
291        }
292    }
293
294    /// 2026-09-30: Bytes one layer holds of `state` for `holding`.
295    pub fn layer_bytes(&self, state: PoolState, holding: Holding) -> usize {
296        let id = match state {
297            PoolState::H => &self.h_id,
298            PoolState::Conv => &self.conv_id,
299        };
300        self.layer
301            .bytes_where(|t| &t.state == id && t.holding == holding) as usize
302    }
303
304    /// 2026-09-30: Bytes of the whole pool: every layer.
305    pub fn total(&self) -> usize {
306        self.layers * self.layer.bytes() as usize
307    }
308}
metrale_model_layers/ssm_reserve/pool_plan.rs

metrale_model_layers/ssm_reserve/
pool_plan.rs
