cortiq_engine/
router.rs

1//! Recon-argmin skill routing (spec §9, P1 signal-consistency): the
2//! container's selection descriptors define per-skill affine subspaces
3//! over φ(x); the winner is the skill that reconstructs φ best. No
4//! trained gate — routing is a property of the skills themselves.
5//!
6//! The decision layer is the debugged cortiq-router recipe (the task-routing
7//! service, `cortiq-bot/cortiq-router/src/router.rs`): raw squared
8//! reconstruction error per skill, a temperature-calibrated softmax over
9//! −error for the confidence, and a NOVELTY ENSEMBLE of three independent
10//! OOD signals — the winner's error as a z-score against its own training
11//! shell, the leader margin, and the calibrated confidence — thresholded by
12//! θ that was set to the (1−fpr) quantile of in-scope held-out scores.
13//! Files without the calibration fall back to the normalized error E with a
14//! fixed threshold (the pre-calibration behaviour, unchanged).
15
16use crate::pipeline::Pipeline;
17use base64::Engine as _;
18use cortiq_core::CmfModel;
19use cortiq_core::quant::f16_to_f32;
20
21/// Ensemble weights and margin sharpness (cortiq-router constants).
22pub const NOVELTY_W_ENERGY: f32 = 0.5;
23pub const NOVELTY_W_MARGIN: f32 = 0.25;
24pub const NOVELTY_W_CONF: f32 = 0.25;
25pub const NOVELTY_MARGIN_K: f32 = 8.0;
26
27#[derive(Debug, Clone)]
28pub struct SkillRoute {
29    pub id: String,
30    /// Normalized reconstruction error E = ‖r − BBᵀr‖²/‖φ‖² ∈ [0, 1]; lower = closer.
31    pub error: f32,
32    /// Raw squared reconstruction error (the calibrated recipe's quantity).
33    pub raw_error: f32,
34    /// Calibrated probability (temperature softmax over −raw_error) — 0 when
35    /// the file carries no calibration.
36    pub probability: f32,
37}
38
39/// The full routing decision for one prompt.
40#[derive(Debug, Clone)]
41pub struct Routing {
42    /// best-first
43    pub scores: Vec<SkillRoute>,
44    /// winner's calibrated confidence (0 without calibration)
45    pub confidence: f32,
46    /// leader margin in `1/(1+err)` units
47    pub margin: f32,
48    /// novelty ensemble score ∈ [0,1] (NaN without calibration)
49    pub novelty: f32,
50    /// OOD verdict: calibrated θ when present, else E_min > `fallback_tau`
51    pub is_novel: bool,
52    pub calibrated: bool,
53}
54
55impl Routing {
56    pub fn winner(&self) -> Option<&SkillRoute> {
57        self.scores.first()
58    }
59}
60
61pub fn decode_f16(b64: &str) -> Option<Vec<f32>> {
62    let bytes = base64::engine::general_purpose::STANDARD.decode(b64).ok()?;
63    Some(
64        bytes
65            .chunks_exact(2)
66            .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]])))
67            .collect(),
68    )
69}
70
71fn sigmoid(x: f32) -> f32 {
72    1.0 / (1.0 + (-x).exp())
73}
74
75fn softmax(v: &mut [f32]) {
76    if v.is_empty() {
77        return;
78    }
79    let mx = v.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
80    let mut s = 0.0f32;
81    for x in v.iter_mut() {
82        *x = (*x - mx).exp();
83        s += *x;
84    }
85    for x in v.iter_mut() {
86        *x /= s.max(1e-30);
87    }
88}
89
90/// Raw squared reconstruction error of φ against a (mean, basis rows) subspace.
91pub fn recon_error(phi: &[f32], mean: &[f32], basis: &[f32], rank: usize) -> f32 {
92    let hidden = phi.len();
93    let r: Vec<f32> = phi.iter().zip(mean).map(|(p, m)| p - m).collect();
94    let rr: f32 = r.iter().map(|v| v * v).sum();
95    let mut proj = 0f32;
96    for k in 0..rank {
97        let row = &basis[k * hidden..(k + 1) * hidden];
98        let c: f32 = row.iter().zip(&r).map(|(b, v)| b * v).sum();
99        proj += c * c;
100    }
101    (rr - proj).max(0.0)
102}
103
104/// Decision from per-skill (id, raw_error, err_mean, err_std, ‖φ‖²) rows —
105/// the pure recipe, shared by `route_full` and the file-level calibration.
106pub fn decide(
107    rows: &[(String, f32, Option<f32>, Option<f32>, f32)],
108    calib: Option<&cortiq_core::format::RoutingCalibration>,
109    fallback_tau: f32,
110) -> Routing {
111    let mut idx: Vec<usize> = (0..rows.len()).collect();
112    idx.sort_by(|&a, &b| rows[a].1.total_cmp(&rows[b].1));
113    let mut scores: Vec<SkillRoute> = idx
114        .iter()
115        .map(|&i| SkillRoute {
116            id: rows[i].0.clone(),
117            error: rows[i].1 / rows[i].4.max(1e-12),
118            raw_error: rows[i].1,
119            probability: 0.0,
120        })
121        .collect();
122    if scores.is_empty() {
123        return Routing {
124            scores,
125            confidence: 0.0,
126            margin: 0.0,
127            novelty: 1.0,
128            is_novel: true,
129            calibrated: calib.is_some(),
130        };
131    }
132    let (Some(c), Some(&top)) = (calib, idx.first()) else {
133        let e_min = scores[0].error;
134        return Routing {
135            scores,
136            confidence: 0.0,
137            margin: 0.0,
138            novelty: f32::NAN,
139            is_novel: e_min > fallback_tau,
140            calibrated: false,
141        };
142    };
143    // confidence: temperature softmax over −raw_error
144    let mut logits: Vec<f32> = idx
145        .iter()
146        .map(|&i| -rows[i].1 / c.temperature.max(1e-3))
147        .collect();
148    softmax(&mut logits);
149    for (s, p) in scores.iter_mut().zip(&logits) {
150        s.probability = *p;
151    }
152    let confidence = logits[0];
153    // margin in 1/(1+err) units
154    let inv = |e: f32| 1.0 / (1.0 + e);
155    let margin = if idx.len() > 1 {
156        inv(rows[idx[0]].1) - inv(rows[idx[1]].1)
157    } else {
158        inv(rows[idx[0]].1)
159    };
160    // energy: winner z-score against its training shell
161    let (em, es) = (
162        rows[top].2.unwrap_or(0.0),
163        rows[top].3.unwrap_or(1.0).max(1e-4),
164    );
165    let z = (rows[top].1 - em) / es;
166    let novelty = NOVELTY_W_ENERGY * sigmoid(z)
167        + NOVELTY_W_MARGIN / (1.0 + margin * NOVELTY_MARGIN_K)
168        + NOVELTY_W_CONF * (1.0 - confidence);
169    Routing {
170        scores,
171        confidence,
172        margin,
173        novelty,
174        is_novel: novelty > c.novelty_theta,
175        calibrated: true,
176    }
177}
178
179/// Per-skill error rows for a φ (skills with malformed descriptors skipped).
180pub fn error_rows(
181    model: &CmfModel,
182    phi_of_layer: &mut dyn FnMut(usize) -> Vec<f32>,
183) -> Vec<(String, f32, Option<f32>, Option<f32>, f32)> {
184    let hidden = model.arch().hidden_size;
185    let mut rows = Vec::new();
186    for skill in &model.header.skills {
187        let Some(sel) = &skill.selection else {
188            continue;
189        };
190        let unit = match sel.metric.as_str() {
191            "mse" => false,
192            "mse_unit" => true,
193            m => {
194                tracing::warn!("skill '{}': unknown metric '{}'", skill.id, m);
195                continue;
196            }
197        };
198        let mut phi = phi_of_layer(sel.phi_layer);
199        if unit {
200            let n = phi.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-12);
201            for x in phi.iter_mut() {
202                *x /= n;
203            }
204        }
205        let (Some(mean), Some(basis)) = (decode_f16(&sel.mean), decode_f16(&sel.basis)) else {
206            tracing::error!("skill '{}': malformed selection payload", skill.id);
207            continue;
208        };
209        if mean.len() != hidden || basis.len() != sel.rank * hidden || phi.len() != hidden {
210            tracing::error!("skill '{}': selection dims mismatch", skill.id);
211            continue;
212        }
213        let e = recon_error(&phi, &mean, &basis, sel.rank);
214        let pp: f32 = phi.iter().map(|v| v * v).sum();
215        rows.push((skill.id.clone(), e, sel.err_mean, sel.err_std, pp));
216    }
217    rows
218}
219
220/// Full decision for a prompt.
221pub fn route_full(
222    model: &CmfModel,
223    pipeline: &mut Pipeline,
224    ids: &[u32],
225    fallback_tau: f32,
226) -> Routing {
227    let mut phi_cache: Vec<(usize, Vec<f32>)> = Vec::new();
228    let mut phi_of = |layer: usize| -> Vec<f32> {
229        if let Some((_, p)) = phi_cache.iter().find(|(l, _)| *l == layer) {
230            return p.clone();
231        }
232        let p = pipeline.probe_phi(ids, layer);
233        phi_cache.push((layer, p.clone()));
234        p
235    };
236    let rows = error_rows(model, &mut phi_of);
237    decide(&rows, model.header.routing.as_ref(), fallback_tau)
238}
239
240/// Score every routable skill; sorted best-first (compatibility API).
241pub fn route(model: &CmfModel, pipeline: &mut Pipeline, ids: &[u32]) -> Vec<SkillRoute> {
242    route_full(model, pipeline, ids, 0.30).scores
243}
244
245/// Held-out in-scope φ samples of every skill (from the descriptors), as
246/// (skill index, φ). Empty when no skill carries them.
247pub fn holdout_phis(model: &CmfModel) -> Vec<(usize, Vec<f32>)> {
248    let hidden = model.arch().hidden_size;
249    let mut out = Vec::new();
250    for (si, skill) in model.header.skills.iter().enumerate() {
251        let Some(sel) = &skill.selection else {
252            continue;
253        };
254        let (Some(h), Some(n)) = (sel.holdout.as_ref(), sel.holdout_n) else {
255            continue;
256        };
257        let Some(v) = decode_f16(h) else { continue };
258        if v.len() != n * hidden {
259            continue;
260        }
261        for i in 0..n {
262            out.push((si, v[i * hidden..(i + 1) * hidden].to_vec()));
263        }
264    }
265    out
266}
267
268/// Fit the file-level calibration from the skills' held-out φ samples: the
269/// temperature by NLL of the true skill under softmax(−err/T) over a
270/// geometric grid, then θ as the (1−fpr) quantile of in-scope novelty
271/// scores. Every skill's φ must be at the SAME phi_layer (mixed layers are
272/// scored per skill; the samples of a skill are compared against every
273/// descriptor's own layer only when equal — otherwise skipped).
274pub fn calibrate(
275    model: &CmfModel,
276    target_fpr: f32,
277) -> Option<cortiq_core::format::RoutingCalibration> {
278    let samples = holdout_phis(model);
279    if samples.is_empty() {
280        return None;
281    }
282    // per sample: rows over all skills — φ is a per-layer quantity, so only
283    // skills sharing the sample's phi_layer are comparable
284    let skills = &model.header.skills;
285    let mut per_sample: Vec<(Vec<(String, f32, Option<f32>, Option<f32>, f32)>, usize)> =
286        Vec::new();
287    for (si, phi) in &samples {
288        let layer = skills[*si]
289            .selection
290            .as_ref()
291            .map(|s| s.phi_layer)
292            .unwrap_or(0);
293        let mut phi_of =
294            |l: usize| -> Vec<f32> { if l == layer { phi.clone() } else { Vec::new() } };
295        let rows = error_rows(model, &mut phi_of);
296        let Some(pos) = rows.iter().position(|r| r.0 == skills[*si].id) else {
297            continue;
298        };
299        per_sample.push((rows, pos));
300    }
301    if per_sample.is_empty() {
302        return None;
303    }
304    // temperature: geometric grid, minimize NLL of the true skill
305    let mut best_t = 1.0f32;
306    let mut best_nll = f32::INFINITY;
307    let mut t = 1e-3f32;
308    // errors are raw squared residuals of hidden states — the scale spans
309    // orders of magnitude across models, hence the wide grid
310    while t <= 1e6 {
311        let mut nll = 0.0f32;
312        for (rows, pos) in &per_sample {
313            let mut logits: Vec<f32> = rows.iter().map(|r| -r.1 / t).collect();
314            softmax(&mut logits);
315            nll -= logits[*pos].max(1e-9).ln();
316        }
317        if nll < best_nll {
318            best_nll = nll;
319            best_t = t;
320        }
321        t *= 1.15;
322    }
323    let mut cal = cortiq_core::format::RoutingCalibration {
324        temperature: best_t,
325        novelty_theta: 0.5,
326        samples: per_sample.len(),
327        target_fpr,
328    };
329    // θ: (1−fpr) quantile of the in-scope novelty scores
330    let mut nov: Vec<f32> = per_sample
331        .iter()
332        .map(|(rows, _)| decide(rows, Some(&cal), 1.0).novelty)
333        .filter(|v| v.is_finite())
334        .collect();
335    nov.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
336    if !nov.is_empty() {
337        let q = (1.0 - target_fpr).clamp(0.0, 1.0);
338        let idx = (((nov.len() - 1) as f32) * q).round() as usize;
339        cal.novelty_theta = (nov[idx.min(nov.len() - 1)] + 1e-4).min(0.999);
340    }
341    Some(cal)
342}
cortiq_engine/router.rs

cortiq_engine/
router.rs
