Skip to main content

cherenkov/qwen4_exp/gpu/
budget.rs

1//! Expert-pool sizing after fixed buffers and host reservations.
2
3use super::*;
4
5pub(super) struct PoolMemory {
6    pub device: usize,
7    pub fixed: usize,
8    pub host: usize,
9    pub allocation_limit: Option<usize>,
10    pub prefill: usize,
11}
12
13impl PoolMemory {
14    pub fn bytes(&self, request: PoolBudget) -> Result<usize> {
15        let available = self
16            .allocation_limit
17            .unwrap_or(usize::MAX)
18            .saturating_sub(self.fixed)
19            .saturating_sub(self.prefill);
20
21        if let PoolBudget::Gb(gb) = request {
22            let bytes = gb * BYTES_PER_GB as f64;
23
24            ensure!(
25                bytes.is_finite() && bytes > 0.0 && bytes < usize::MAX as f64,
26                "expert pool size is not representable in bytes"
27            );
28
29            let bytes = bytes as usize;
30
31            ensure!(
32                bytes <= available,
33                "explicit expert pool exceeds server memory budget after fixed buffers and prefill scratch"
34            );
35
36            return Ok(bytes);
37        }
38
39        // Adaptive keeps the existing 6 GB margin on the measured Air, but
40        // no longer forces an 8 GB minimum or a 20 GB maximum on other Macs.
41        let margin = if matches!(request, PoolBudget::Adaptive) {
42            6 * BYTES_PER_GB
43        } else {
44            BYTES_PER_GB / 2
45        };
46        let headroom = self.prefill.max(margin);
47        let bytes = self
48            .device
49            .saturating_sub(self.host)
50            .saturating_sub(self.fixed)
51            .saturating_sub(headroom)
52            .min(available);
53
54        ensure!(
55            bytes > 0,
56            "fixed buffers and reservations leave no expert pool capacity"
57        );
58
59        Ok(bytes)
60    }
61}
62
63#[cfg(test)]
64#[path = "../../../tests/unit/qwen4_exp/gpu/budget.rs"]
65mod tests;