metrale_model_engine/factory/build/
slots.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-01: How many sequence slots a model is built with. `--max-batch-size N` asks for N;
4//! `--max-batch-size auto` asks for the count that admits the most sequences at the full
5//! `--max-seq-len`, sizing the per-slot state (the SSM pools and every other slot-scaled reserve
6//! term) and the KV pool together: each slot of state costs KV blocks, so past some count more
7//! slots admit fewer sequences, not more.
8//!
9//! Owner: metrale-model-engine.
10//! Invariants:
11//! - [`balance_slots`] is pure: the reserve and the pool size come in as functions of the slot
12//!   count, so the boot and the tests evaluate the same arithmetic.
13//! - An explicit count is never changed; `auto` is the only request this module resolves.
14
15use anyhow::{Result, ensure};
16
17/// 2026-10-01: The `--max-batch-size` request.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum SlotRequest {
20    /// 2026-10-01: Exactly this many slots.
21    Count(usize),
22    /// 2026-10-01: The count [`balance_slots`] chooses, at most [`AUTO_MAX_SLOTS`].
23    Auto,
24}
25
26/// 2026-10-01: The most slots `auto` considers: the widest rung of the decode batch ladder
27/// (`traits::DECODE_BATCH_LADDER`), the widest batch a decode step is padded to and captured at.
28pub const AUTO_MAX_SLOTS: usize =
29    crate::traits::DECODE_BATCH_LADDER[crate::traits::DECODE_BATCH_LADDER.len() - 1];
30
31impl SlotRequest {
32    /// 2026-10-01: The slots sized before the KV pool is (the pre-load reserve, the buffer
33    /// arena): the count, or for `auto` its ceiling [`AUTO_MAX_SLOTS`].
34    pub fn ceiling(self) -> usize {
35        match self {
36            SlotRequest::Count(n) => n,
37            SlotRequest::Auto => AUTO_MAX_SLOTS,
38        }
39    }
40}
41
42impl std::str::FromStr for SlotRequest {
43    type Err = String;
44
45    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
46        if s == "auto" {
47            return Ok(SlotRequest::Auto);
48        }
49        s.parse::<usize>()
50            .map(SlotRequest::Count)
51            .map_err(|_| format!("`{s}` is neither a slot count nor `auto`"))
52    }
53}
54
55impl std::fmt::Display for SlotRequest {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        match self {
58            SlotRequest::Count(n) => write!(f, "{n}"),
59            SlotRequest::Auto => f.write_str("auto"),
60        }
61    }
62}
63
64/// 2026-10-01: The slot request and the pre-load reserve as a function of the slot count (the
65/// preflight's plan, `inference_reserve(slots)`), which KV sizing evaluates at the count it
66/// builds.
67pub struct SlotPlan<'a> {
68    pub request: SlotRequest,
69    pub reserve_for: &'a dyn Fn(usize) -> Result<usize>,
70}
71
72/// 2026-10-01: A built model and the slot count it was built with (the request, or what `auto`
73/// resolved to), which the scheduler and the serve's disclosure take from here.
74pub struct BuiltModel {
75    pub model: Box<dyn crate::traits::Model>,
76    pub max_batch_size: usize,
77    /// 2026-10-04: sha256 over the drafter's stored weights (`WeightStore::drafter_weights_sha256`),
78    /// `None` for a checkpoint with no `mtp.*` tensors. The `--spec-cost-model measured` drafter
79    /// key.
80    pub drafter_weights_sha256: Option<String>,
81}
82
83/// 2026-10-01: One candidate slot count's outcome.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub struct Balance {
86    /// 2026-10-01: The slot count.
87    pub slots: usize,
88    /// 2026-10-01: KV blocks the pool gets with that many slots.
89    pub kv_blocks: usize,
90    /// 2026-10-01: Sequences admitted at the full `--max-seq-len`: the slots, or the full-length
91    /// sequences the pool holds, whichever is fewer.
92    pub admitted: usize,
93}
94
95/// 2026-10-01: The slot count in `1..=max_slots` that admits the most sequences at
96/// `blocks_per_seq` blocks each, `kv_blocks_at(slots)` giving the pool size for a count. On a tie
97/// the smaller count wins, so the chosen count never exceeds what its pool holds at full length
98/// (a count above it ties with the count one below, whose pool is at least as large) and the boot's
99/// overcommit check passes. An error when `max_slots` is 0, when `kv_blocks_at` fails for every
100/// count, or when no count's pool holds one full-length sequence.
101pub fn balance_slots(
102    max_slots: usize,
103    blocks_per_seq: usize,
104    kv_blocks_at: impl Fn(usize) -> Result<usize>,
105) -> Result<Balance> {
106    ensure!(max_slots > 0, "no slot count to balance (ceiling 0)");
107    let per_seq = blocks_per_seq.max(1);
108    let mut best: Option<Balance> = None;
109    for slots in 1..=max_slots {
110        let Ok(kv_blocks) = kv_blocks_at(slots) else {
111            continue;
112        };
113        let b = Balance {
114            slots,
115            kv_blocks,
116            admitted: slots.min(kv_blocks / per_seq),
117        };
118        if best.is_none_or(|x| b.admitted > x.admitted) {
119            best = Some(b);
120        }
121    }
122    let best = best.ok_or_else(|| {
123        anyhow::anyhow!("no slot count in 1..={max_slots} leaves memory for the KV pool")
124    })?;
125    ensure!(
126        best.admitted > 0,
127        "no slot count in 1..={max_slots} leaves a KV pool that holds one sequence at the full \
128         --max-seq-len ({per_seq} blocks; {} at 1 slot): lower --max-seq-len",
129        best.kv_blocks
130    );
131    Ok(best)
132}
133
134#[cfg(test)]
135#[path = "slots_tests.rs"]
136mod slots_tests;
metrale_model_engine/factory/build/slots.rs

metrale_model_engine/factory/build/
slots.rs
