gglib_core/domain/context_fit.rs
1//! How large a context this machine can actually serve.
2//!
3//! `gglib up` already does this arithmetic: it picks a model by asking what
4//! fits in VRAM at 32k and prints the answer as the number that earns the
5//! user's trust. The launch path then resolved its context from a chain whose
6//! lowest reachable rung was a flat 4096 and served that instead — so the
7//! number shown and the number used were unrelated.
8//!
9//! This is the launch-time half of that arithmetic. `up`'s shortlist still
10//! asks a different question — "does this model fit *at* 32k?" — so the two
11//! are not one calculation and cannot be made one cheaply: they are answered
12//! against different budgets, and the shortlist runs before the model is
13//! downloaded, when its real KV geometry is not yet readable.
14//!
15//! What changed is that the banner no longer implies otherwise. It reports 32k
16//! as the bar the model had to clear, and says the served context is sized at
17//! launch, which is true and knowable. Naming a rung there would have been the
18//! same error in the other direction.
19//!
20//! ## Why it snaps to rungs
21//!
22//! A resident is identified partly by the context it was launched with, and a
23//! request that resolves to a different one evicts and relaunches. A value
24//! computed from a live free-memory reading would wobble between requests and
25//! recycle the server — blowing the prefix cache and every saved slot file —
26//! on essentially every turn. Snapping to a fixed ladder makes the result a
27//! step function that changes only when the machine genuinely changes.
28
29use crate::cache_config::KvCacheType;
30use crate::domain::kv_estimate::{KvElemsPerToken, kv_bytes_per_token};
31use crate::domain::recommendation::BUDGET_UTILISATION;
32
33use crate::settings::DEFAULT_CONTEXT_SIZE;
34
35/// `bytes * factor`, saturating and rounding down.
36#[allow(clippy::cast_precision_loss, clippy::cast_sign_loss)]
37#[allow(clippy::cast_possible_truncation)]
38fn scale(bytes: u64, factor: f64) -> u64 {
39 (bytes as f64 * factor) as u64
40}
41
42/// The context sizes a fitted value is allowed to take.
43///
44/// Powers of two from the built-in default upward. Deliberately coarse: the
45/// point is stability, not squeezing out the last few thousand tokens.
46const RUNGS: [u64; 6] = [4096, 8192, 16_384, 32_768, 65_536, 131_072];
47
48// The ladder starts at the built-in default, so a fitted value can never be
49// worse than serving no fit at all. Compile-time because both operands are
50// constants and the mistake would be a source edit.
51const _: () = assert!(RUNGS[0] == DEFAULT_CONTEXT_SIZE);
52
53/// The largest context `weights_bytes` can serve inside `budget_bytes`.
54///
55/// `None` whenever the answer cannot be computed from facts — an unknown
56/// trained context, unknown KV shape, or no memory reading. That is a refusal
57/// rather than an optimistic guess, matching `SlotFootprint::new`: a caller
58/// that gets `None` falls back down its own chain instead of launching against
59/// a number nobody stands behind.
60///
61/// Also `None` when the machine cannot fit even the smallest rung. Returning
62/// something smaller would be inventing a context this module has no basis
63/// for; the built-in default is the right thing to fall back to, and it will
64/// fail honestly if it does not fit either.
65///
66/// Pure. `GGLIB_DISABLE_CONTEXT_FIT` is read by the caller, not here — the
67/// switch belongs with the other runtime switches at the admission site, and a
68/// domain function that read the environment could not be tested without
69/// mutating process-global state.
70///
71/// `budget_bytes` must be a figure that does not move between requests: this
72/// value ends up in a resident's identity, so a budget that drifts evicts and
73/// relaunches the model it just sized.
74///
75/// A live free-memory reading is therefore wrong — on Apple it is a fraction
76/// of available system RAM, and it moves with whatever else is open. So is
77/// netting out the current resident set, which was tried and removed: it moved
78/// whenever a co-resident loaded or was evicted, and made one model's budget
79/// depend on whether another model's KV shape was readable. What the caller
80/// supplies is total device capacity less a fixed reservation for the second
81/// resident slot.
82#[must_use]
83pub fn fit_context(
84 trained_ctx: Option<u64>,
85 weights_bytes: Option<u64>,
86 kv: Option<KvElemsPerToken>,
87 k: KvCacheType,
88 v: KvCacheType,
89 budget_bytes: Option<u64>,
90) -> Option<u64> {
91 fit_context_explained(trained_ctx, weights_bytes, kv, k, v, budget_bytes).0
92}
93
94/// What [`fit_context`] worked from, for a person reading a launch log.
95///
96/// The two constants governing this — `BUDGET_UTILISATION` and the caller's
97/// co-resident reservation — are judgement calls, not measurements. Nothing
98/// acts on this record; it exists so the numbers behind a fitted context are
99/// visible when someone asks whether those judgements were right, rather than
100/// having to be re-derived from the rung alone.
101///
102/// `gglib model explain` is where a person reads it, through
103/// `residency::explain::explain_fit`, which supplies the same two budgets
104/// `admit` does. Before that it reached only a `debug!` line written after a
105/// launch, which is not a reading anyone could take across a catalog — and
106/// ADR 0009's first kill criterion needs exactly that.
107#[derive(Debug, Clone, Copy)]
108pub struct FitInputs {
109 /// Device memory the fit was allowed to spend against.
110 pub budget_bytes: Option<u64>,
111 /// The model's weights, as summed across shards.
112 pub weights_bytes: Option<u64>,
113 /// Bytes of KV cache each token of context costs at the resolved types.
114 pub kv_bytes_per_token: Option<u64>,
115 /// The model's trained ceiling.
116 pub trained_ctx: Option<u64>,
117 /// The context that fit before snapping to a rung — the difference between
118 /// this and the chosen rung is what the ladder costs.
119 pub unsnapped: Option<u64>,
120}
121
122/// The same calculation, reporting what it worked from.
123#[must_use]
124pub fn fit_context_explained(
125 trained_ctx: Option<u64>,
126 weights_bytes: Option<u64>,
127 kv: Option<KvElemsPerToken>,
128 k: KvCacheType,
129 v: KvCacheType,
130 budget_bytes: Option<u64>,
131) -> (Option<u64>, FitInputs) {
132 let per_token = kv
133 .map(|elems| kv_bytes_per_token(elems, k, v))
134 .filter(|&b| b > 0);
135 // `0` is this codebase's sentinel for "size could not be read"
136 // (`total_model_bytes`), and it is the most dangerous value to take
137 // literally: weights that cost nothing hand the entire budget to the KV
138 // cache. `launch_deadline_secs` reads the same field and refuses it the
139 // same way.
140 let weights = weights_bytes.filter(|&b| b > 0);
141 // Never claim more than the model was trained for, *before* snapping, so
142 // the result is always both a rung and within the model's range.
143 let unsnapped = (|| {
144 let usable = scale(budget_bytes?, BUDGET_UTILISATION);
145 let for_kv = usable.checked_sub(weights?)?;
146 Some((for_kv / per_token?).min(trained_ctx?))
147 })();
148
149 let inputs = FitInputs {
150 budget_bytes,
151 weights_bytes: weights,
152 kv_bytes_per_token: per_token,
153 trained_ctx,
154 unsnapped,
155 };
156 let fitted = unsnapped.and_then(|u| RUNGS.iter().rev().find(|&&rung| rung <= u).copied());
157 (fitted, inputs)
158}
159
160#[cfg(test)]
161#[path = "context_fit_tests.rs"]
162mod tests;