1pub const POOL_SHARE_OF_FREE: f64 = 0.60;
4
5pub const PLANNING_SHARE: f64 = 0.85;
6
7pub const WINDOW_FLOOR: u64 = 4096;
8
9pub const WINDOW_GRANULARITY: u64 = 1024;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct Sizing {
13 pub budget_bytes: u64,
14 pub weight_bytes: u64,
15 pub kv_bytes_per_token: u64,
16 pub trained_context: u64,
17}
18
19impl Sizing {
20 pub fn to_json(&self) -> serde_json::Value {
21 serde_json::json!({
22 "budget_bytes": self.budget_bytes,
23 "weight_bytes": self.weight_bytes,
24 "kv_bytes_per_token": self.kv_bytes_per_token,
25 "trained_context": self.trained_context,
26 })
27 }
28
29 pub fn from_json(v: &serde_json::Value) -> Option<Sizing> {
30 Some(Sizing {
31 budget_bytes: v["budget_bytes"].as_u64()?,
32 weight_bytes: v["weight_bytes"].as_u64()?,
33 kv_bytes_per_token: v["kv_bytes_per_token"].as_u64()?,
34 trained_context: v["trained_context"].as_u64()?,
35 })
36 }
37}
38
39pub fn reserve_bytes(budget_bytes: u64) -> u64 {
40 (1u64 << 30).max(budget_bytes / 8)
41}
42
43pub fn window(models: &[Sizing], lanes: u32) -> Option<u64> {
44 let budget = models.iter().map(|m| m.budget_bytes).min()?;
45 let trained = models.iter().map(|m| m.trained_context).min()?;
46 let per_token: u64 = models.iter().map(|m| m.kv_bytes_per_token).sum::<u64>() * u64::from(lanes.max(1));
47 if budget == 0 || trained == 0 || per_token == 0 {
48 return None;
49 }
50 let committed = models.iter().map(|m| m.weight_bytes).sum::<u64>() + reserve_bytes(budget);
51 let pool = (POOL_SHARE_OF_FREE * budget.saturating_sub(committed) as f64) as u64;
52 let mut n = pool / per_token;
53 n -= n % WINDOW_GRANULARITY;
54 Some(n.min(trained).max(WINDOW_FLOOR.min(trained)))
55}
56
57pub fn headroom_bytes(budget_bytes: u64) -> u64 {
58 reserve_bytes(budget_bytes) / 2
59}
60
61pub fn cells_over_budget(budget_bytes: u64, free_after: u64, bytes_per_cell: u64) -> u64 {
62 if budget_bytes == 0 || bytes_per_cell == 0 {
63 return 0;
64 }
65 headroom_bytes(budget_bytes).saturating_sub(free_after).div_ceil(bytes_per_cell)
66}
67
68pub fn tokens(n: u64) -> String {
70 match n {
71 0..=9_999 => n.to_string(),
72 10_000..=999_499 => format!("{}k", (n + 500) / 1000),
73 _ => format!("{:.1}M", n as f64 / 1e6),
74 }
75}
76
77pub fn budgeted_cells(desired: u64, one_sequence: u64, bytes_per_cell: u64, free_bytes: u64) -> u64 {
78 if free_bytes == 0 || bytes_per_cell == 0 {
79 return desired;
80 }
81 let cap = (POOL_SHARE_OF_FREE * free_bytes as f64) as u64 / bytes_per_cell;
82 desired.min(cap).max(one_sequence.min(desired))
83}
84
85#[cfg(test)]
86mod tests {
87 use super::*;
88
89 #[test]
90 fn token_counts_read_as_people_say_them() {
91 assert_eq!(tokens(4096), "4096");
92 assert_eq!(tokens(131_072), "131k");
93 assert_eq!(tokens(905_562), "906k");
94 assert_eq!(tokens(2_097_152), "2.1M");
95 assert_eq!(tokens(999_499), "999k");
96 assert_eq!(tokens(999_500), "1.0M", "never 1000k");
97 }
98
99 const GIB: u64 = 1 << 30;
100
101 #[test]
102 fn an_affordable_shape_is_granted_whole() {
103 assert_eq!(budgeted_cells(65_536, 16_384, 144 << 10, 30 * GIB), 65_536);
104 }
105
106 #[test]
107 fn a_shape_past_the_budget_is_capped_to_it() {
108 let cells = budgeted_cells(262_144, 32_768, 96 << 10, 20 * GIB);
109 assert_eq!(cells, (0.6 * (20 * GIB) as f64) as u64 / (96 << 10));
110 assert!((32_768..262_144).contains(&cells));
111 }
112
113 #[test]
114 fn one_full_sequence_always_fits() {
115 assert_eq!(budgeted_cells(65_536, 16_384, 144 << 10, GIB), 16_384);
116 }
117
118 #[test]
119 fn no_reported_budget_grants_the_request() {
120 assert_eq!(budgeted_cells(65_536, 16_384, 144 << 10, 0), 65_536);
121 assert_eq!(budgeted_cells(65_536, 16_384, 0, 30 * GIB), 65_536);
122 }
123
124 fn model(budget: u64, weights: u64, price: u64, trained: u64) -> Sizing {
125 Sizing { budget_bytes: budget * GIB, weight_bytes: weights * GIB, kv_bytes_per_token: price, trained_context: trained }
126 }
127
128 const KV_8B: u64 = 2 * 32 * 8 * 128 * 2;
129
130 #[test]
131 fn the_native_engines_cases_size_the_same() {
132 let mut m = model(96, 5, KV_8B, 131_072);
133 assert_eq!(window(&[m], 1), Some(131_072));
134 m.trained_context = 262_144;
135 assert_eq!(window(&[m], 1), Some(262_144));
136 let many = window(&[m], 32).unwrap();
137 assert!((WINDOW_FLOOR..262_144).contains(&many), "{many}");
138 let small = window(&[model(6, 2, KV_8B, 131_072)], 1).unwrap();
139 assert!((WINDOW_FLOOR..131_072).contains(&small), "{small}");
140 let m = model(14, 8, KV_8B, 1 << 20);
141 let (one, four) = (window(&[m], 1).unwrap(), window(&[m], 4).unwrap());
142 assert!(four * 4 <= one + 4 * WINDOW_GRANULARITY && four * 4 + 4 * WINDOW_GRANULARITY >= one, "{one} {four}");
143 let f16 = window(&[model(12, 5, KV_8B, 1 << 20)], 1).unwrap();
144 let q8 = window(&[model(12, 5, KV_8B * 34 / 64, 1 << 20)], 1).unwrap();
145 assert!(q8 > f16, "{f16} {q8}");
146 for budget in [4, 8, 16, 24, 48, 96, 192] {
147 let w = window(&[model(budget, 3, KV_8B, 1 << 20)], 1).unwrap();
148 assert!(w.is_multiple_of(WINDOW_GRANULARITY) && (WINDOW_FLOOR..=1 << 20).contains(&w), "{budget}: {w}");
149 }
150 }
151
152 #[test]
153 fn a_window_is_what_the_budget_leaves_for_every_lane() {
154 let m = model(40, 4, 144 << 10, 1 << 20);
155 let raw = (0.6 * (31 * GIB) as f64) as u64 / ((144 << 10) * 4);
156 assert_eq!(window(&[m], 4), Some(raw - raw % WINDOW_GRANULARITY));
157 }
158
159 #[test]
160 fn the_trained_window_caps_it_and_the_floor_holds_it_up() {
161 assert_eq!(window(&[model(128, 4, 16 << 10, 40_960)], 4), Some(40_960));
162 assert_eq!(window(&[model(16, 30, 96 << 10, 262_144)], 8), Some(WINDOW_FLOOR));
163 assert_eq!(window(&[model(16, 30, 96 << 10, 2048)], 8), Some(2048));
164 }
165
166 #[test]
167 fn models_on_one_device_share_it() {
168 let a = model(40, 4, 64 << 10, 1 << 20);
169 let b = model(48, 8, 32 << 10, 32_768);
170 let both = window(&[a, b], 2).unwrap();
171 let left = 40 * GIB - (12 * GIB + reserve_bytes(40 * GIB));
172 let raw = (0.6 * left as f64) as u64 / ((96 << 10) * 2);
173 assert_eq!(both, (raw - raw % 1024).min(32_768));
174 assert!(both < window(&[a], 2).unwrap());
175 }
176
177 #[test]
178 fn unknown_inputs_size_nothing() {
179 assert_eq!(window(&[], 4), None);
180 assert_eq!(window(&[model(0, 4, 1024, 32_768)], 4), None);
181 assert_eq!(window(&[model(40, 4, 0, 32_768)], 4), None);
182 assert_eq!(window(&[model(40, 4, 1024, 0)], 4), None);
183 }
184
185 #[test]
186 fn sizing_travels_as_json() {
187 let m = model(40, 4, 144 << 10, 40_960);
188 assert_eq!(Sizing::from_json(&m.to_json()), Some(m));
189 assert_eq!(Sizing::from_json(&serde_json::json!({"budget_bytes": 1})), None);
190 }
191
192 #[test]
193 fn a_context_that_ate_the_headroom_gives_cells_back() {
194 let budget = 40 * GIB;
195 let headroom = headroom_bytes(budget);
196 assert_eq!(headroom, 5 * GIB / 2);
197 assert_eq!(headroom_bytes(4 * GIB), GIB / 2, "never under half a GiB");
198 assert_eq!(cells_over_budget(budget, headroom, 1024), 0);
199 assert_eq!(cells_over_budget(budget, headroom + GIB, 1024), 0);
200 assert_eq!(cells_over_budget(budget, headroom - 1024 * 10, 1024), 10);
201 assert_eq!(cells_over_budget(budget, headroom - 1, 1024), 1);
202 assert_eq!(cells_over_budget(budget, 0, 1 << 20), headroom.div_ceil(1 << 20));
203 assert_eq!(cells_over_budget(0, 0, 1024), 0);
204 assert_eq!(cells_over_budget(budget, 0, 0), 0);
205 }
206}