ferrum_types/
startup_memory.rs

1//! Device-aware startup policy. Costs come from the compiled execution plan,
2//! including its provider workspaces and physical storage layouts.
3
4use crate::{RuntimeConfigSnapshot, RuntimeConfigSource};
5use serde::{Deserialize, Serialize};
6
7/// Sampled before model upload. Available bytes are additional allocatable
8/// memory, already accounting for other processes and platform ceilings.
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
10pub struct DeviceMemorySnapshot {
11    pub capacity_bytes: u64,
12    pub available_bytes: u64,
13    pub source: String,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17pub struct StartupMemoryRequest {
18    pub device: DeviceMemorySnapshot,
19    pub usable_capacity_bytes: u64,
20    pub context_is_explicit: bool,
21    pub sequences_is_explicit: bool,
22    pub batch_is_explicit: bool,
23}
24
25impl StartupMemoryRequest {
26    pub fn from_snapshot(
27        device: DeviceMemorySnapshot,
28        utilization: f32,
29        requested: &RuntimeConfigSnapshot,
30    ) -> Result<Self, String> {
31        if device.capacity_bytes == 0 || device.available_bytes > device.capacity_bytes {
32            return Err("invalid device memory sample: expected available <= capacity > 0".into());
33        }
34        if !utilization.is_finite() || utilization <= 0.0 || utilization > 1.0 {
35            return Err("memory utilization must be in (0, 1]".into());
36        }
37        let explicit = |key: &str| {
38            requested.entries.iter().any(|entry| {
39                entry.key == key
40                    && matches!(
41                        entry.source,
42                        RuntimeConfigSource::Cli
43                            | RuntimeConfigSource::Env
44                            | RuntimeConfigSource::ConfigFile
45                            | RuntimeConfigSource::ScriptCase
46                    )
47            })
48        };
49        let budget = requested
50            .entries
51            .iter()
52            .find(|entry| entry.key == "FERRUM_RUNTIME_MEMORY_BUDGET_BYTES")
53            .map(|entry| {
54                entry
55                    .effective_value
56                    .parse::<u64>()
57                    .map_err(|_| "runtime memory budget must be a positive byte count".to_owned())
58            })
59            .transpose()?;
60        let usable_capacity_bytes = budget.unwrap_or_else(|| {
61            ((device.available_bytes as f64) * f64::from(utilization)).floor() as u64
62        });
63        if usable_capacity_bytes == 0 || usable_capacity_bytes > device.available_bytes {
64            return Err(format!(
65                "runtime memory budget {usable_capacity_bytes} bytes cannot fit the sampled {} available bytes (device ceiling {} bytes; {})",
66                device.available_bytes, device.capacity_bytes, device.source
67            ));
68        }
69        Ok(Self {
70            device,
71            usable_capacity_bytes,
72            context_is_explicit: explicit("FERRUM_MAX_MODEL_LEN") || explicit("FERRUM_KV_CAPACITY"),
73            sequences_is_explicit: explicit("FERRUM_PAGED_MAX_SEQS"),
74            batch_is_explicit: explicit("FERRUM_MAX_BATCHED_TOKENS"),
75        })
76    }
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
80pub struct StartupResourceLimits {
81    pub context_tokens: usize,
82    pub max_sequences: usize,
83    pub max_batch_tokens: usize,
84}
85
86/// Independent capacity probes, not a promise that every admitted sequence
87/// can simultaneously occupy the maximum context. Runtime admission remains
88/// responsible for the actual mixture of live requests.
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum StartupWorkload {
91    Prefill {
92        context_tokens: usize,
93        chunk_tokens: usize,
94    },
95    Decode {
96        /// Request allocation ceiling; sequence state is probed at a one-token
97        /// committed frontier, not a full context for every sequence.
98        context_tokens: usize,
99        active_sequences: usize,
100    },
101}
102
103#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
104pub struct StartupMemoryPlan {
105    pub request: StartupMemoryRequest,
106    pub requested: StartupResourceLimits,
107    pub selected: StartupResourceLimits,
108    pub context_peak_bytes: u64,
109    pub decode_peak_bytes: u64,
110    pub reasons: Vec<String>,
111}
112
113impl StartupMemoryPlan {
114    pub fn apply_to_engine_config(&self, engine: &mut crate::EngineConfig) -> Result<(), String> {
115        engine.runtime.max_model_len = Some(self.selected.context_tokens);
116        engine.scheduler.max_running_requests = self.selected.max_sequences;
117        engine.batching.max_num_batched_tokens = self.selected.max_batch_tokens;
118        engine.memory.usable_capacity_bytes = Some(
119            usize::try_from(self.request.usable_capacity_bytes)
120                .map_err(|_| "runtime memory budget exceeds the process address space")?,
121        );
122        engine.runtime.startup_memory_plan = Some(self.clone());
123        Ok(())
124    }
125}
126
127/// Find supported limits using a monotone upper bound supplied by the actual
128/// compiled plan. No tensors are allocated by this search. Provider errors are
129/// errors, never interpreted as evidence that a smaller request will work.
130pub fn fit_startup_resources(
131    request: &StartupMemoryRequest,
132    target: StartupResourceLimits,
133    mut peak_bytes: impl FnMut(StartupWorkload) -> Result<u64, String>,
134) -> Result<StartupMemoryPlan, String> {
135    if target.context_tokens == 0 || target.max_sequences == 0 || target.max_batch_tokens == 0 {
136        return Err("startup context, sequence and batch limits must be positive".into());
137    }
138    let budget = request.usable_capacity_bytes;
139    let prefill = |context_tokens: usize, chunk: usize| StartupWorkload::Prefill {
140        context_tokens,
141        chunk_tokens: chunk.min(context_tokens).max(1),
142    };
143    let batch_floor = if request.sequences_is_explicit {
144        target.max_sequences
145    } else {
146        1
147    };
148    let batch_target = if request.batch_is_explicit {
149        target.max_batch_tokens
150    } else {
151        target.max_batch_tokens.max(batch_floor)
152    };
153    if batch_target < batch_floor {
154        return Err("explicit batch token limit is smaller than explicit sequence limit".into());
155    }
156    let batch_cost =
157        |batch: usize, cost: &mut dyn FnMut(StartupWorkload) -> Result<u64, String>| {
158            let context = if request.context_is_explicit {
159                target.context_tokens
160            } else {
161                batch.min(target.context_tokens)
162            };
163            cost(prefill(context, batch))
164        };
165    let batch = if request.batch_is_explicit {
166        require_fit(
167            "explicit batch/context",
168            batch_cost(batch_target, &mut peak_bytes)?,
169            budget,
170        )?;
171        batch_target
172    } else {
173        largest_fit(batch_floor, batch_target, budget, |value| {
174            batch_cost(value, &mut peak_bytes)
175        })?
176    };
177    let context = if request.context_is_explicit {
178        require_fit(
179            "explicit context",
180            peak_bytes(prefill(target.context_tokens, batch))?,
181            budget,
182        )?;
183        target.context_tokens
184    } else {
185        largest_fit(1, target.context_tokens, budget, |tokens| {
186            let prefill_peak = peak_bytes(prefill(tokens, batch))?;
187            let decode_peak = peak_bytes(StartupWorkload::Decode {
188                context_tokens: tokens,
189                active_sequences: if request.sequences_is_explicit {
190                    target.max_sequences
191                } else {
192                    1
193                },
194            })?;
195            Ok(prefill_peak.max(decode_peak))
196        })?
197    };
198    let decode = |active_sequences| StartupWorkload::Decode {
199        context_tokens: context,
200        active_sequences,
201    };
202    let sequences = if request.sequences_is_explicit {
203        require_fit(
204            "explicit concurrency",
205            peak_bytes(decode(target.max_sequences))?,
206            budget,
207        )?;
208        target.max_sequences
209    } else {
210        largest_fit(1, target.max_sequences.min(batch), budget, |count| {
211            peak_bytes(decode(count))
212        })?
213    };
214    let selected = StartupResourceLimits {
215        context_tokens: context,
216        max_sequences: sequences,
217        max_batch_tokens: batch,
218    };
219    let mut reasons = Vec::new();
220    for (name, before, after) in [
221        ("context tokens", target.context_tokens, context),
222        ("concurrent sequences", target.max_sequences, sequences),
223        ("batch tokens", target.max_batch_tokens, batch),
224    ] {
225        if before != after {
226            reasons.push(format!("{name}: {before} -> {after}, compiled resource plan fitted to {budget} usable bytes"));
227        }
228    }
229    Ok(StartupMemoryPlan {
230        request: request.clone(),
231        requested: target,
232        selected,
233        context_peak_bytes: peak_bytes(prefill(context, batch))?,
234        decode_peak_bytes: peak_bytes(decode(sequences))?,
235        reasons,
236    })
237}
238
239fn require_fit(label: &str, required: u64, budget: u64) -> Result<(), String> {
240    if required > budget {
241        return Err(format!(
242            "{label} requires {required} bytes, exceeding the {budget}-byte runtime budget"
243        ));
244    }
245    Ok(())
246}
247
248fn largest_fit(
249    minimum: usize,
250    maximum: usize,
251    budget: u64,
252    mut cost: impl FnMut(usize) -> Result<u64, String>,
253) -> Result<usize, String> {
254    require_fit("minimum runnable configuration", cost(minimum)?, budget)?;
255    let mut low = minimum;
256    let mut high = maximum;
257    while low < high {
258        let middle = low + (high - low) / 2 + (high - low) % 2;
259        if cost(middle)? <= budget {
260            low = middle;
261        } else {
262            high = middle - 1;
263        }
264    }
265    Ok(low)
266}
267
268#[cfg(test)]
269mod tests;
ferrum_types/startup_memory.rs

ferrum_types/
startup_memory.rs
