diff --git a/mesh/mesh-rs/crates/mesh-router/src/lib.rs b/mesh/mesh-rs/crates/mesh-router/src/lib.rs index 5d304639b..d689a4dc9 100644 --- a/mesh/mesh-rs/crates/mesh-router/src/lib.rs +++ b/mesh/mesh-rs/crates/mesh-router/src/lib.rs @@ -14,36 +14,145 @@ use mesh_schema::{BackendId, Fidelity, RouteCandidate, RouteDecision, ServeRequest}; +/// How a backend's energy scales with the context it must attend over. +/// +/// # Why a scalar is not enough +/// +/// The router previously costed every backend as `uj_per_token × tokens`: one +/// scalar, linear in token count, blind to how much context precedes each +/// token. For a backend whose per-token cost genuinely does not depend on the +/// prefix, that is correct and remains expressible ([`CostProfile::linear`]). +/// +/// For one that re-reads its context per token it is not, and the error has a +/// **direction**. A linear model applied to a super-linear backend +/// *under*-estimates, always. At a gate, under-estimating means admitting work +/// the budget cannot cover, so the request is accepted and then overruns its +/// ceiling — which is the one failure a joule ceiling exists to prevent. A +/// ceiling checked against a model that cannot represent the expensive case is +/// not a ceiling; it is a formality that happens to pass. +/// +/// # Why no profile ships as a default +/// +/// The context term is a property of a *particular backend on particular +/// hardware*, and must be calibrated by measurement there. Published figures +/// are not a substitute: the study that motivated this change measured one +/// model per attention condition, so architecture is confounded with model +/// family and parameter count, and its slopes cannot be lifted into a constant +/// without inheriting that confound. So this type carries the *shape* of the +/// cost and the operator supplies the numbers — the same reason +/// `BandingScale` ships no canonical scale. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CostProfile { + /// Fixed microjoules to serve the request at all, independent of size. + pub base_uj: u64, + /// Microjoules per token, independent of context length. + pub uj_per_token: u64, + /// Additional microjoules per token, per 1024 tokens of preceding context. + /// + /// Zero reproduces the historical scalar model exactly. A non-zero value + /// says this backend gets more expensive per token as the prefix grows, + /// which is what makes a long-context request refusable *before* it runs. + pub uj_per_token_per_kctx: u64, +} + +/// Tokens of context per unit of [`CostProfile::uj_per_token_per_kctx`]. +const CTX_UNIT: u64 = 1024; + +impl CostProfile { + /// The historical model: a flat per-token cost with no context term. + /// + /// Retained so existing callers are unchanged and the scalar case stays + /// honestly expressible — some backends really are flat. + pub fn linear(uj_per_token: u64) -> Self { + CostProfile { base_uj: 0, uj_per_token, uj_per_token_per_kctx: 0 } + } + + /// A profile whose per-token cost grows with the context length. + pub fn context_scaled(base_uj: u64, uj_per_token: u64, uj_per_token_per_kctx: u64) -> Self { + CostProfile { base_uj, uj_per_token, uj_per_token_per_kctx } + } + + /// Whether this profile can represent a context-dependent cost at all. + /// + /// A router gating long-context work against flat profiles is deciding with + /// a model that cannot express the case it is meant to catch. + pub fn is_context_aware(&self) -> bool { + self.uj_per_token_per_kctx > 0 + } + + /// Estimated microjoules for `tokens` tokens over `context_tokens` of + /// preceding context. + /// + /// Saturating throughout: an estimate that wrapped would come back small, + /// and a small estimate is precisely what admits an unaffordable request. + pub fn estimate_uj(&self, tokens: u64, context_tokens: u64) -> u64 { + let flat = self.uj_per_token.saturating_mul(tokens); + let ctx_units = context_tokens / CTX_UNIT; + let ctx = self + .uj_per_token_per_kctx + .saturating_mul(tokens) + .saturating_mul(ctx_units); + self.base_uj.saturating_add(flat).saturating_add(ctx) + } + + /// The profile as it is sealed into a receipt, so a route decision can be + /// re-derived rather than trusted. + pub fn wire(&self) -> String { + format!( + "base={};per_token={};per_token_per_kctx={}", + self.base_uj, self.uj_per_token, self.uj_per_token_per_kctx + ) + } +} + /// A backend the router can route to: its identity, the recipe it serves, the -/// fidelity it delivers, and its modeled per-token energy cost. +/// fidelity it delivers, and its modeled energy cost. #[derive(Clone, Debug, PartialEq, Eq)] pub struct BackendSpec { pub id: BackendId, pub recipe: String, pub fidelity: Fidelity, - pub uj_per_token: u64, + /// How this backend's cost scales. See [`CostProfile`]. + pub cost: CostProfile, } impl BackendSpec { + /// A backend with a flat per-token cost — the historical constructor, + /// unchanged in behaviour. pub fn new( id: impl Into, recipe: impl Into, fidelity: Fidelity, uj_per_token: u64, ) -> Self { - Self { - id: BackendId::new(id), - recipe: recipe.into(), - fidelity, - uj_per_token, - } + Self::with_cost(id, recipe, fidelity, CostProfile::linear(uj_per_token)) } - /// Estimated total microjoules to serve `req` on this backend: per-token - /// cost over input + output tokens. + /// A backend with an explicit cost profile. + pub fn with_cost( + id: impl Into, + recipe: impl Into, + fidelity: Fidelity, + cost: CostProfile, + ) -> Self { + Self { id: BackendId::new(id), recipe: recipe.into(), fidelity, cost } + } + + /// The flat per-token term, for callers that only ever set one. + pub fn uj_per_token(&self) -> u64 { + self.cost.uj_per_token + } + + /// Estimated total microjoules to serve `req` on this backend. + /// + /// The context term is charged against `prompt_tokens`: the prefix every + /// generated token attends over. Output tokens extend that prefix as they + /// are produced, so this remains a lower bound on a growing-context + /// backend — deliberately, since the alternative is to model the growth and + /// then defend a second uncalibrated constant. fn estimate(&self, req: &ServeRequest) -> u64 { let tokens = req.prompt_tokens as u64 + req.max_tokens as u64; - self.uj_per_token.saturating_mul(tokens) + self.cost.estimate_uj(tokens, req.prompt_tokens as u64) } } @@ -168,6 +277,116 @@ mod tests { r } + fn req_ctx(floor: Fidelity, ceiling: u64, prompt_tokens: u32, max_tokens: u32) -> ServeRequest { + ServeRequest { + prompt_hash: ContentId::of_bytes(b"q"), + prompt_tokens, + max_tokens, + fidelity_floor: floor, + joule_ceiling_uj: ceiling, + } + } + + /// The historical scalar model must be reproduced exactly, so existing + /// backends and their receipts do not shift under this change. + #[test] + fn a_linear_profile_is_the_old_model_exactly() { + let p = CostProfile::linear(7); + assert_eq!(p.estimate_uj(100, 0), 700); + // Context is irrelevant to a flat backend, however long it gets. + assert_eq!(p.estimate_uj(100, 1_000_000), 700); + assert!(!p.is_context_aware()); + + // And via the untouched constructor. + let b = BackendSpec::new("b", "r", Fidelity::F1, 7); + assert_eq!(b.estimate(&req_ctx(Fidelity::F1, u64::MAX, 10, 90)), 700); + assert_eq!(b.uj_per_token(), 7); + } + + /// The defect this change exists for: a flat model applied to a + /// context-scaling backend under-estimates, and the error grows with the + /// context. Under-estimating at a gate means *admitting* work the ceiling + /// cannot cover. + #[test] + fn a_flat_model_under_estimates_a_context_scaling_backend() { + let flat = CostProfile::linear(10); + let scaled = CostProfile::context_scaled(0, 10, 2); + + // Short prompt: the two models agree, which is why the defect hides. + assert_eq!(flat.estimate_uj(1_000, 512), scaled.estimate_uj(1_000, 512)); + + // Long prompt: they diverge, and always in the same direction. + let ctx = 16_384; + let f = flat.estimate_uj(1_000, ctx); + let s = scaled.estimate_uj(1_000, ctx); + assert!(s > f, "the context term must cost more, not less"); + // 16_384 / 1024 = 16 context units -> 10*1000 + 2*1000*16. + assert_eq!(s, 10_000 + 32_000); + assert_eq!(f, 10_000); + } + + /// The capability the docs imply and the scalar model could not deliver: + /// refuse a long-context route *before* it runs. + #[test] + fn a_long_context_route_is_refused_before_it_burns_the_budget() { + let ceiling = 20_000u64; + // 100 output tokens over a 1_500-token prompt. + let r = req_ctx(Fidelity::F1, ceiling, 1_500, 100); + + // Described flat, the same backend is costed at 10 * 1_600 = 16_000 and + // is admitted. + let mut blind = Router::new(); + blind.register(BackendSpec::new("b", "r", Fidelity::F1, 10)); + let blind_est = blind.spec(&BackendId::new("b")).unwrap().estimate(&r); + assert_eq!(blind_est, 16_000); + assert!(blind_est <= ceiling, "the flat model admits it"); + assert!(blind.route(&r).chosen.is_some(), "and the router therefore routes it"); + + // Described with a calibrated context term, the same request costs more + // than the ceiling and is refused before anything runs. + let mut aware = Router::new(); + aware.register(BackendSpec::with_cost( + "b", + "r", + Fidelity::F1, + CostProfile::context_scaled(0, 10, 200), + )); + let aware_est = aware.spec(&BackendId::new("b")).unwrap().estimate(&r); + assert!(aware_est > ceiling, "the calibrated model refuses it ({aware_est})"); + assert!( + aware.route(&r).chosen.is_none(), + "the router must not route what it cannot afford" + ); + } + + /// A profile that cannot express context is detectable, so a deployment + /// gating long-context work on flat profiles can notice it is deciding with + /// a model that cannot represent the case. + #[test] + fn context_awareness_is_inspectable() { + assert!(!CostProfile::linear(10).is_context_aware()); + assert!(CostProfile::context_scaled(0, 10, 1).is_context_aware()); + assert!(!CostProfile::context_scaled(500, 10, 0).is_context_aware()); + } + + /// An estimate must never wrap: a wrapped estimate comes back small, and a + /// small estimate is exactly what admits an unaffordable request. + #[test] + fn an_overflowing_estimate_saturates_rather_than_wrapping() { + let p = CostProfile::context_scaled(u64::MAX, u64::MAX, u64::MAX); + assert_eq!(p.estimate_uj(u64::MAX, u64::MAX), u64::MAX); + // And a saturated estimate exceeds any real ceiling, so it refuses. + assert!(p.estimate_uj(1_000_000, 1_000_000) > 1_000_000_000); + } + + /// The profile is sealable, so a route decision can be re-derived from the + /// receipt rather than trusted. + #[test] + fn the_profile_is_carried_on_the_wire() { + let w = CostProfile::context_scaled(5, 10, 2).wire(); + assert_eq!(w, "base=5;per_token=10;per_token_per_kctx=2"); + } + fn req(floor: Fidelity, ceiling: u64, max_tokens: u32) -> ServeRequest { ServeRequest { prompt_hash: ContentId::of_bytes(b"q"),