Skip to main content

ferrox_core/
kernel_registry.rs

1//! Sealed kernel-lookup registry: makes a *missing* kernel loud instead
2//! of silently slow.
3//!
4//! ## Why this exists
5//!
6//! The worst bug class this engine has shipped is not a wrong kernel but
7//! an absent one, silently replaced by a correct-but-slow fallback. Two
8//! real examples:
9//!
10//! - IQ4_XS batched prefill ran on the **CPU** because
11//!   [`crate::weight_matrix::WeightMatrix`]'s `metal_kind_supported`
12//!   predicate and `apply_gpu_batch`'s per-kind dispatch table disagreed
13//!   about one kind. The only symptom was a benchmark row 13.7x behind.
14//! - Gemma-4-E2B is slower on Metal than on CPU. Output is correct, so
15//!   nothing fails; the model simply never reaches a batched path.
16//!
17//! Both are *lookups that missed*. Neither produced a diagnostic.
18//!
19//! ## The shape
20//!
21//! Every dispatch decision that asks "is there a kernel for this
22//! (backend, op, quant kind)?" records the answer here.
23//!
24//! - **Build phase.** While the model is constructed, each weight matrix
25//!   is probed eagerly ([`crate::weight_matrix::WeightMatrix::probe_kernels`]):
26//!   the same predicates the hot path will consult are evaluated once per
27//!   weight and recorded. This is what turns "why is this row slow" into
28//!   a startup line.
29//! - **[`seal`].** Called once the model is loaded. It summarises the
30//!   build picture, warns about quantized weights that will run off the
31//!   selected accelerator, and switches on post-seal checking.
32//! - **Run phase.** After sealing, a dispatch-site lookup that misses a
33//!   *combination the build probe never saw* is by definition an
34//!   unpredicted slow path. It warns once, loudly, naming the call site
35//!   ([`#[track_caller]`](std::panic::Location)) and the quant kind, and
36//!   is a hard error under `FERROX_STRICT_KERNELS=1` so CI and benchmarks
37//!   can run closed.
38//!
39//! ## Cost
40//!
41//! This sits next to dispatch, so it follows the `OnceLock` discipline
42//! used by `metal_dense_enabled` / `min_task_macs`: the environment is
43//! read exactly once per process, never per dispatch.
44//!
45//! On the hot path the added cost is **zero instructions on a hit** —
46//! hits are only ever recorded by the build probe, never by a dispatch
47//! site. A dispatch site records only when it is *about to take a
48//! fallback*, i.e. only when it is already paying orders of magnitude
49//! more than the bookkeeping costs.
50//!
51//! And that bookkeeping never takes an exclusive lock in steady state: a
52//! repeat lookup is a shared `RwLock` read plus one relaxed
53//! `fetch_add`. The write lock is taken once per distinct call site, to
54//! create the row and decide whether to warn. A dispatch path must not
55//! be able to serialize rayon workers behind a mutex.
56//!
57//! `FERROX_KERNEL_REGISTRY=0` reduces every entry point to one relaxed
58//! atomic load and a return.
59//!
60//! ## What it must not do
61//!
62//! Observe only. Nothing in this module may change a dispatch decision;
63//! the predicates it calls are the same ones the dispatch takes, read a
64//! second time, and their results are recorded rather than acted on.
65
66use crate::weight_matrix::QuantKind;
67use std::collections::{HashMap, HashSet};
68use std::panic::Location;
69use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
70use std::sync::{OnceLock, RwLock};
71
72/// Which execution backend a lookup was resolved against.
73///
74/// **The GPU variants are generated**, from
75/// [`crate::weight_matrix::gpu_backend::gpu_backend_table`] — the same
76/// ordered list `WeightMatrix::apply_gpu` dispatches over and
77/// `active_backend` reads. That is the beachhead verdict's point 4:
78/// "a backend cannot be dispatched to without being reportable". It
79/// used to be a hand-kept list beside the dispatch order, i.e. two
80/// structures that had to agree about one thing with nothing enforcing
81/// it, which is the failure this file exists to catch in *kernels* and
82/// was quietly repeating in *backends*.
83///
84/// `Cpu` is written out rather than generated: it is the fallthrough
85/// every accelerator miss lands on, not a backend anything dispatches
86/// to, and it has no row in the table for the same reason it has no
87/// cargo feature.
88///
89/// This works only because the registry and the seam are one crate. A
90/// `macro_rules!` cannot add a variant to an enum in another crate, so
91/// an out-of-tree backend could not own its own variant; if one is ever
92/// wanted, this becomes a trait object or a `&'static str`, not a
93/// second list.
94macro_rules! define_backend_enum {
95    ([] $(($variant:ident, $name:literal, $ty:ident)),* $(,)?) => {
96        #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
97        pub enum Backend {
98            Cpu,
99            $($variant,)*
100        }
101
102        impl Backend {
103            pub fn as_str(self) -> &'static str {
104                match self {
105                    Backend::Cpu => "cpu",
106                    $(Backend::$variant => $name,)*
107                }
108            }
109
110            /// Every backend this build knows about, `Cpu` first. Lets a
111            /// test iterate the set instead of restating it.
112            pub const ALL: &'static [Backend] = &[Backend::Cpu, $(Backend::$variant,)*];
113        }
114    };
115}
116crate::weight_matrix::gpu_backend::gpu_backend_table!(define_backend_enum);
117
118impl Backend {
119    /// True for the accelerator backends. A miss here means work the
120    /// user asked to run on a GPU is running somewhere else.
121    pub fn is_accelerator(self) -> bool {
122        !matches!(self, Backend::Cpu)
123    }
124}
125
126impl std::fmt::Display for Backend {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        f.write_str(self.as_str())
129    }
130}
131
132/// Canonical op names. Free-form `&'static str` is accepted everywhere,
133/// but sticking to these keeps the report groupable.
134pub mod op {
135    /// One activation against a whole weight matrix (decode).
136    pub const MATVEC: &str = "matvec";
137    /// Several independent matvecs sharing one activation (fused q/k/v).
138    pub const MATVEC_MULTI: &str = "matvec_multi";
139    /// A real batched GEMM over `batch` activations (prefill).
140    pub const GEMM_PREFILL: &str = "gemm_prefill";
141    /// Whole gate/up/down SwiGLU fused into one dispatch.
142    pub const FFN_SWIGLU: &str = "ffn_swiglu";
143    /// An engine-level capability rather than a per-tensor kernel: does
144    /// this engine have a batched prefill path at all?
145    pub const ENGINE_PREFILL_BATCH: &str = "engine.prefill_batch";
146}
147
148/// The identity of one lookup. Deduplicated on this; a repeated lookup
149/// bumps [`Entry::count`] instead of adding a row.
150#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
151pub struct Key {
152    pub backend: Backend,
153    /// See [`op`].
154    pub op: &'static str,
155    /// Tensor role (`"attn_q"`, `"ffn_down"`, `"moe_router"`, ...) or
156    /// `"(dispatch)"` when recorded from the hot path, where the role is
157    /// not known.
158    pub role: &'static str,
159    /// `None` for a non-quantized (F32 / MXFP4-pair) matrix.
160    pub kind: Option<QuantKind>,
161    pub file: &'static str,
162    pub line: u32,
163}
164
165impl Key {
166    /// The part of the identity that decides whether a kernel exists.
167    /// Call site and tensor role are diagnostics, not part of the
168    /// question being asked.
169    fn shape(&self) -> (Backend, &'static str, Option<QuantKind>) {
170        (self.backend, self.op, self.kind)
171    }
172
173    fn kind_name(&self) -> &'static str {
174        self.kind.map_or("f32", QuantKind::name)
175    }
176}
177
178impl std::fmt::Display for Key {
179    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
180        write!(
181            f,
182            "{} {} {} ({}) at {}:{}",
183            self.backend,
184            self.op,
185            self.kind_name(),
186            self.role,
187            self.file,
188            self.line
189        )
190    }
191}
192
193/// How bad a miss is. Recorded at the call site, which is the only
194/// place that knows — inferring it at report time from the kind or the
195/// backend is how a signal turns into noise nobody reads.
196#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
197pub enum Severity {
198    /// The fallback *is* the intended path. An MoE router is a lone
199    /// small F32 matvec that costs more to ship to a GPU than to compute
200    /// on the host; that is a decision, not an omission.
201    ByDesign,
202    /// No kernel exists, so the work runs somewhere slower than the
203    /// backend the user selected, with correct output and no symptom.
204    /// This is the class the registry exists to surface.
205    SlowPath,
206}
207
208/// What the lookup resolved to.
209#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
210pub enum Outcome {
211    /// A real kernel exists and will be used.
212    Hit,
213    /// No kernel; the caller takes `fallback` instead. The string is the
214    /// fallback's name, so the report reads as a sentence.
215    Miss {
216        fallback: &'static str,
217        severity: Severity,
218    },
219}
220
221impl Outcome {
222    /// A miss that is a silent slow path.
223    pub fn slow_path(fallback: &'static str) -> Self {
224        Outcome::Miss {
225            fallback,
226            severity: Severity::SlowPath,
227        }
228    }
229
230    /// A miss whose fallback is the deliberate, documented choice.
231    pub fn by_design(fallback: &'static str) -> Self {
232        Outcome::Miss {
233            fallback,
234            severity: Severity::ByDesign,
235        }
236    }
237
238    pub fn is_miss(self) -> bool {
239        matches!(self, Outcome::Miss { .. })
240    }
241
242    pub fn is_slow_path(self) -> bool {
243        matches!(
244            self,
245            Outcome::Miss {
246                severity: Severity::SlowPath,
247                ..
248            }
249        )
250    }
251
252    pub fn fallback(self) -> Option<&'static str> {
253        match self {
254            Outcome::Miss { fallback, .. } => Some(fallback),
255            Outcome::Hit => None,
256        }
257    }
258}
259
260/// Whether the lookup happened while the model was being built or after
261/// [`seal`].
262#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
263pub enum Phase {
264    Build,
265    Run,
266}
267
268/// A lookup being reported, minus the call site (which `#[track_caller]`
269/// supplies).
270#[derive(Clone, Copy, Debug)]
271pub struct Lookup {
272    pub backend: Backend,
273    pub op: &'static str,
274    pub role: &'static str,
275    pub kind: Option<QuantKind>,
276}
277
278impl Lookup {
279    pub fn new(backend: Backend, op: &'static str, kind: Option<QuantKind>) -> Self {
280        Lookup {
281            backend,
282            op,
283            role: "(dispatch)",
284            kind,
285        }
286    }
287
288    pub fn with_role(mut self, role: &'static str) -> Self {
289        self.role = role;
290        self
291    }
292}
293
294/// One deduplicated row of the registry.
295#[derive(Clone, Copy, Debug)]
296pub struct Entry {
297    pub key: Key,
298    pub outcome: Outcome,
299    pub phase: Phase,
300    /// How many lookups collapsed into this row.
301    pub count: u64,
302}
303
304impl std::fmt::Display for Entry {
305    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
306        match self.outcome {
307            Outcome::Hit => write!(f, "hit   {}  x{}", self.key, self.count),
308            Outcome::Miss {
309                fallback,
310                severity: Severity::ByDesign,
311            } => write!(f, "host  {} -> {}  x{}", self.key, fallback, self.count),
312            Outcome::Miss {
313                fallback,
314                severity: Severity::SlowPath,
315            } => write!(f, "MISS  {} -> {}  x{}", self.key, fallback, self.count),
316        }
317    }
318}
319
320/// The picture at [`seal`] time.
321#[derive(Clone, Debug, Default)]
322pub struct SealReport {
323    /// Every build-phase row, sorted for stable printing.
324    pub entries: Vec<Entry>,
325    /// Build-phase rows whose outcome was a miss, of either severity.
326    pub misses: Vec<Entry>,
327    /// The subset of `misses` that is a real silent slow path
328    /// ([`Severity::SlowPath`]) on an *accelerator* backend the process
329    /// actually selected — work the user asked to run on a GPU that will
330    /// not. A CPU-backend miss is informational: there is nothing to
331    /// fall off.
332    pub violations: Vec<Entry>,
333}
334
335impl SealReport {
336    /// One line per row, for `FERROX_KERNEL_REGISTRY=1`.
337    pub fn render(&self) -> String {
338        let mut s = String::new();
339        for e in &self.entries {
340            s.push_str("ferrox kernels: ");
341            s.push_str(&e.to_string());
342            s.push('\n');
343        }
344        s
345    }
346
347    /// The loud one-paragraph summary printed whenever a violation
348    /// exists, registry verbose or not.
349    pub fn render_violations(&self) -> String {
350        let mut s = String::new();
351        for e in &self.violations {
352            // Engine-level capabilities are counted once, not per
353            // weight, so do not label their count "weights".
354            let unit = if e.key.op.starts_with("engine.") {
355                String::new()
356            } else {
357                format!(", {} weights", e.count)
358            };
359            let kind = match e.key.kind {
360                Some(k) => format!(" {}", k.name()),
361                None => String::new(),
362            };
363            s.push_str(&format!(
364                "ferrox: NO KERNEL for {} {}{} ({}) -> falls back to {} [{}:{}{}]\n",
365                e.key.backend,
366                e.key.op,
367                kind,
368                e.key.role,
369                e.outcome.fallback().unwrap_or("(hit)"),
370                e.key.file,
371                e.key.line,
372                unit,
373            ));
374        }
375        s
376    }
377}
378
379/// Raised by [`seal_or_error`] under `FERROX_STRICT_KERNELS=1`.
380#[derive(Clone, Debug)]
381pub struct StrictKernelError {
382    pub report: SealReport,
383}
384
385impl std::fmt::Display for StrictKernelError {
386    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
387        write!(
388            f,
389            "FERROX_STRICT_KERNELS=1 and {} kernel lookup(s) missed:\n{}",
390            self.report.violations.len(),
391            self.report.render_violations()
392        )
393    }
394}
395
396impl std::error::Error for StrictKernelError {}
397
398/// One stored row. `count` is atomic so a repeat lookup -- the only
399/// thing a dispatch site ever does after the first call -- needs a
400/// *shared* read lock and a relaxed increment, never exclusive access.
401/// A dispatch path must not be able to serialize rayon workers on a
402/// mutex, whatever else it does.
403struct Row {
404    outcome: Outcome,
405    phase: Phase,
406    count: AtomicU64,
407}
408
409#[derive(Default)]
410struct State {
411    rows: HashMap<Key, Row>,
412    /// Lookup *shapes* seen during the build phase. A post-seal miss on
413    /// a shape in here was predicted at load time and already reported;
414    /// one that is not is the thing this registry exists to find.
415    known: HashSet<(Backend, &'static str, Option<QuantKind>)>,
416    /// Post-seal unpredicted misses, in discovery order.
417    surprises: Vec<Entry>,
418}
419
420impl State {
421    fn snapshot(&self, build_only: bool) -> Vec<Entry> {
422        let mut entries: Vec<Entry> = self
423            .rows
424            .iter()
425            .filter(|(_, row)| !build_only || row.phase == Phase::Build)
426            .map(|(key, row)| Entry {
427                key: *key,
428                outcome: row.outcome,
429                phase: row.phase,
430                count: row.count.load(Ordering::Relaxed),
431            })
432            .collect();
433        entries.sort_by_key(|e| {
434            (
435                e.key.backend,
436                e.key.op,
437                e.key.kind.map(|k| k.name()).unwrap_or("f32"),
438                e.key.role,
439                e.key.line,
440            )
441        });
442        entries
443    }
444}
445
446/// A registry instance. There is one [`global`] instance; tests build
447/// their own so they never race the process-wide one.
448pub struct Registry {
449    inner: RwLock<State>,
450    sealed: AtomicBool,
451}
452
453impl Default for Registry {
454    fn default() -> Self {
455        Self::new()
456    }
457}
458
459impl Registry {
460    pub fn new() -> Self {
461        Registry {
462            inner: RwLock::new(State::default()),
463            sealed: AtomicBool::new(false),
464        }
465    }
466
467    fn read(&self) -> std::sync::RwLockReadGuard<'_, State> {
468        self.inner.read().unwrap_or_else(|e| e.into_inner())
469    }
470
471    fn write(&self) -> std::sync::RwLockWriteGuard<'_, State> {
472        self.inner.write().unwrap_or_else(|e| e.into_inner())
473    }
474
475    /// Shared-lock increment for a row that already exists. Returns
476    /// false when the key is new and the caller must take the write
477    /// lock. This is the whole steady-state cost of the registry on a
478    /// dispatch path.
479    fn bump_existing(&self, key: &Key) -> bool {
480        match self.read().rows.get(key) {
481            Some(row) => {
482                row.count.fetch_add(1, Ordering::Relaxed);
483                true
484            }
485            None => false,
486        }
487    }
488
489    pub fn is_sealed(&self) -> bool {
490        self.sealed.load(Ordering::Relaxed)
491    }
492
493    /// Records a build-phase lookup. Always accepted, even after
494    /// [`Self::seal`] — a process that loads a second model (draft +
495    /// target for speculative decoding) probes it too, and those
496    /// lookups are predictions, not surprises.
497    pub fn record_build_at(&self, loc: &'static Location<'static>, l: Lookup, outcome: Outcome) {
498        let key = Key {
499            backend: l.backend,
500            op: l.op,
501            role: l.role,
502            kind: l.kind,
503            file: loc.file(),
504            line: loc.line(),
505        };
506        if self.bump_existing(&key) {
507            return;
508        }
509        let mut st = self.write();
510        st.known.insert(key.shape());
511        st.rows
512            .entry(key)
513            .or_insert_with(|| Row {
514                outcome,
515                phase: Phase::Build,
516                count: AtomicU64::new(0),
517            })
518            .count
519            .fetch_add(1, Ordering::Relaxed);
520    }
521
522    /// Records a dispatch-site lookup. Before sealing this is just
523    /// bookkeeping; after sealing, a miss on a shape the build probe
524    /// never predicted warns once and is collected into
525    /// [`Self::surprises`].
526    pub fn record_at(&self, loc: &'static Location<'static>, l: Lookup, outcome: Outcome) {
527        let key = Key {
528            backend: l.backend,
529            op: l.op,
530            role: l.role,
531            kind: l.kind,
532            file: loc.file(),
533            line: loc.line(),
534        };
535        // Steady state: the key already exists, so the warn-or-not
536        // decision was settled when it was created. A shared read lock
537        // and one relaxed add, and nothing here blocks anything.
538        if self.bump_existing(&key) {
539            return;
540        }
541        let sealed = self.is_sealed();
542        let mut st = self.write();
543        // Re-check under the write lock: another thread may have raced
544        // us to this key, in which case it already made the decision.
545        if let Some(row) = st.rows.get(&key) {
546            row.count.fetch_add(1, Ordering::Relaxed);
547            return;
548        }
549        let phase = if sealed { Phase::Run } else { Phase::Build };
550        st.rows.insert(
551            key,
552            Row {
553                outcome,
554                phase,
555                count: AtomicU64::new(1),
556            },
557        );
558        if !sealed {
559            st.known.insert(key.shape());
560            return;
561        }
562        if !outcome.is_slow_path() || st.known.contains(&key.shape()) {
563            return;
564        }
565        st.surprises.push(Entry {
566            key,
567            outcome,
568            phase: Phase::Run,
569            count: 1,
570        });
571        drop(st);
572        let fallback = outcome
573            .fallback()
574            .unwrap_or("(unknown)" /* unreachable: guarded by is_slow_path */);
575        eprintln!(
576            "ferrox: SILENT SLOW PATH — kernel lookup missed after the model was sealed.\n\
577             ferrox:   {} {} for {} has no kernel; falling back to {}.\n\
578             ferrox:   call site {}:{} (role {}).\n\
579             ferrox:   this was not predicted at load time, so no startup diagnostic covered it.\n\
580             ferrox:   set FERROX_STRICT_KERNELS=1 to make this a hard error.",
581            key.backend,
582            key.op,
583            key.kind_name(),
584            fallback,
585            key.file,
586            key.line,
587            key.role,
588        );
589    }
590
591    /// Freezes the build picture and switches on post-seal checking.
592    /// Idempotent: re-sealing recomputes the report over everything
593    /// recorded so far.
594    pub fn seal(&self) -> SealReport {
595        self.sealed.store(true, Ordering::Relaxed);
596        let entries = self.read().snapshot(true);
597        let misses: Vec<Entry> = entries
598            .iter()
599            .copied()
600            .filter(|e| e.outcome.is_miss())
601            .collect();
602        let violations: Vec<Entry> = misses
603            .iter()
604            .copied()
605            .filter(|e| e.key.backend.is_accelerator() && e.outcome.is_slow_path())
606            .collect();
607        SealReport {
608            entries,
609            misses,
610            violations,
611        }
612    }
613
614    /// Post-seal misses that the build probe did not predict.
615    pub fn surprises(&self) -> Vec<Entry> {
616        self.read().surprises.clone()
617    }
618
619    /// Every row, build and run, sorted like [`SealReport::entries`].
620    pub fn entries(&self) -> Vec<Entry> {
621        self.read().snapshot(false)
622    }
623}
624
625static GLOBAL: OnceLock<Registry> = OnceLock::new();
626
627/// The process-wide registry.
628pub fn global() -> &'static Registry {
629    GLOBAL.get_or_init(Registry::new)
630}
631
632/// Whether the registry records at all.
633///
634/// Default **on**: a warning nobody enabled is the entire point. Read
635/// once per process (`FERROX_KERNEL_REGISTRY=0|false|off` disables), so
636/// no dispatch ever touches the environment.
637pub fn enabled() -> bool {
638    static V: OnceLock<bool> = OnceLock::new();
639    *V.get_or_init(|| {
640        !matches!(
641            std::env::var("FERROX_KERNEL_REGISTRY").ok().as_deref(),
642            Some("0") | Some("false") | Some("off")
643        )
644    })
645}
646
647/// Whether the full build-phase table is printed at [`seal`].
648pub fn verbose() -> bool {
649    static V: OnceLock<bool> = OnceLock::new();
650    *V.get_or_init(|| {
651        matches!(
652            std::env::var("FERROX_KERNEL_REGISTRY").ok().as_deref(),
653            Some("1") | Some("true") | Some("on") | Some("verbose")
654        )
655    })
656}
657
658/// Whether a missing kernel is a hard error. Set this in CI and in
659/// benchmark harnesses so a slow path cannot be published as a number.
660pub fn strict() -> bool {
661    static V: OnceLock<bool> = OnceLock::new();
662    *V.get_or_init(|| {
663        matches!(
664            std::env::var("FERROX_STRICT_KERNELS").ok().as_deref(),
665            Some("1") | Some("true") | Some("on")
666        )
667    })
668}
669
670/// Record an eager, load-time lookup against the global registry.
671#[track_caller]
672pub fn record_build(l: Lookup, outcome: Outcome) {
673    if !enabled() {
674        return;
675    }
676    global().record_build_at(Location::caller(), l, outcome);
677}
678
679/// Record a dispatch-site lookup that found a kernel. Dispatch sites do
680/// not normally call this — hits cost nothing precisely because they are
681/// not recorded on the hot path — but it exists for completeness.
682#[track_caller]
683pub fn hit(l: Lookup) {
684    if !enabled() {
685        return;
686    }
687    global().record_at(Location::caller(), l, Outcome::Hit);
688}
689
690/// Record a dispatch-site lookup that missed and is taking a slower
691/// `fallback` than the selected backend.
692///
693/// Call this on the fallback branch only. By construction the branch is
694/// already about to do far more work than a hash and a lock.
695#[track_caller]
696pub fn miss(l: Lookup, fallback: &'static str) {
697    if !enabled() {
698        return;
699    }
700    global().record_at(Location::caller(), l, Outcome::slow_path(fallback));
701}
702
703/// Record a dispatch-site lookup that missed, where the fallback is the
704/// deliberate choice rather than a gap. Never warns; recorded so the
705/// report is a complete picture instead of a filtered one.
706#[track_caller]
707pub fn miss_by_design(l: Lookup, fallback: &'static str) {
708    if !enabled() {
709        return;
710    }
711    global().record_at(Location::caller(), l, Outcome::by_design(fallback));
712}
713
714/// Seal the global registry: print what the build probe found, warn
715/// loudly about quantized weights that will run off the selected
716/// accelerator, and switch on post-seal checking.
717pub fn seal() -> SealReport {
718    let report = global().seal();
719    if !enabled() {
720        return report;
721    }
722    if verbose() {
723        eprint!("{}", report.render());
724    }
725    if !report.violations.is_empty() && !strict() {
726        eprint!("{}", report.render_violations());
727        eprintln!(
728            "ferrox: {} kernel lookup(s) above will run on a slower path than the \
729             selected backend. Set FERROX_STRICT_KERNELS=1 to refuse to run instead.",
730            report.violations.len()
731        );
732    }
733    report
734}
735
736/// [`seal`], but returns `Err` under `FERROX_STRICT_KERNELS=1` when a
737/// quantized weight has no kernel on the selected accelerator.
738pub fn seal_or_error() -> Result<SealReport, StrictKernelError> {
739    let report = seal();
740    if strict() && !report.violations.is_empty() {
741        return Err(StrictKernelError { report });
742    }
743    Ok(report)
744}
745
746#[cfg(test)]
747mod tests {
748    use super::*;
749
750    fn lookup(backend: Backend, kind: Option<QuantKind>) -> Lookup {
751        Lookup::new(backend, op::GEMM_PREFILL, kind).with_role("ffn_down")
752    }
753
754    #[test]
755    fn a_build_hit_is_recorded_once_per_shape_with_a_count() {
756        let r = Registry::new();
757        let loc = Location::caller();
758        for _ in 0..5 {
759            r.record_build_at(
760                loc,
761                lookup(Backend::Metal, Some(QuantKind::Q4K)),
762                Outcome::Hit,
763            );
764        }
765        let report = r.seal();
766        assert_eq!(report.entries.len(), 1);
767        assert_eq!(report.entries[0].count, 5);
768        assert!(report.misses.is_empty());
769        assert!(report.violations.is_empty());
770    }
771
772    /// A quantized weight with no accelerator kernel is exactly the
773    /// IQ4_XS bug: correct output, silent CPU fallback, no diagnostic.
774    /// Seal must call it a violation.
775    #[test]
776    fn a_quantized_weight_with_no_accelerator_kernel_is_a_violation() {
777        let r = Registry::new();
778        r.record_build_at(
779            Location::caller(),
780            lookup(Backend::Metal, Some(QuantKind::IQ4XS)),
781            Outcome::slow_path("CPU apply_batch"),
782        );
783        let report = r.seal();
784        assert_eq!(report.violations.len(), 1);
785        assert_eq!(report.violations[0].key.kind, Some(QuantKind::IQ4XS));
786        assert!(report.render_violations().contains("IQ4_XS"));
787    }
788
789    /// An F32 matrix has no quantized kernel by construction and its
790    /// host GEMV is a deliberate decision. It is a miss, but not a
791    /// violation — otherwise every MoE router would fail the gate and
792    /// the signal would be worthless.
793    #[test]
794    fn an_f32_miss_is_reported_but_is_not_a_violation() {
795        let r = Registry::new();
796        r.record_build_at(
797            Location::caller(),
798            lookup(Backend::Metal, None),
799            Outcome::by_design("host GEMV"),
800        );
801        let report = r.seal();
802        assert_eq!(report.misses.len(), 1);
803        assert!(report.violations.is_empty());
804    }
805
806    /// A CPU-only process has no accelerator to fall off, so its misses
807    /// are informational.
808    #[test]
809    fn a_cpu_backend_miss_is_not_a_violation() {
810        let r = Registry::new();
811        r.record_build_at(
812            Location::caller(),
813            lookup(Backend::Cpu, Some(QuantKind::IQ2XXS)),
814            Outcome::slow_path("f32 dequant-dot"),
815        );
816        assert!(r.seal().violations.is_empty());
817    }
818
819    #[test]
820    fn a_post_seal_miss_on_a_predicted_shape_is_not_a_surprise() {
821        let r = Registry::new();
822        let loc = Location::caller();
823        r.record_build_at(
824            loc,
825            lookup(Backend::Metal, Some(QuantKind::IQ4XS)),
826            Outcome::slow_path("CPU apply_batch"),
827        );
828        r.seal();
829        r.record_at(
830            loc,
831            lookup(Backend::Metal, Some(QuantKind::IQ4XS)),
832            Outcome::slow_path("CPU apply_batch"),
833        );
834        assert!(r.surprises().is_empty(), "seal already reported this shape");
835    }
836
837    /// The registry's whole purpose: a lookup that misses on a shape the
838    /// build probe never saw is an unpredicted slow path.
839    #[test]
840    fn a_post_seal_miss_on_an_unpredicted_shape_is_a_surprise_reported_once() {
841        let r = Registry::new();
842        let loc = Location::caller();
843        r.record_build_at(
844            loc,
845            lookup(Backend::Metal, Some(QuantKind::Q4K)),
846            Outcome::Hit,
847        );
848        r.seal();
849        for _ in 0..3 {
850            r.record_at(
851                loc,
852                lookup(Backend::Metal, Some(QuantKind::Q2K)),
853                Outcome::slow_path("CPU apply_batch"),
854            );
855        }
856        let surprises = r.surprises();
857        assert_eq!(surprises.len(), 1, "warned once, not once per dispatch");
858        assert_eq!(surprises[0].key.kind, Some(QuantKind::Q2K));
859        assert_eq!(surprises[0].phase, Phase::Run);
860    }
861
862    #[test]
863    fn a_post_seal_hit_is_never_a_surprise() {
864        let r = Registry::new();
865        r.seal();
866        r.record_at(
867            Location::caller(),
868            lookup(Backend::Metal, Some(QuantKind::Q4K)),
869            Outcome::Hit,
870        );
871        assert!(r.surprises().is_empty());
872    }
873
874    /// Probing a second model after the first was sealed must not
875    /// manufacture surprises (speculative decoding loads two).
876    #[test]
877    fn build_records_after_seal_extend_the_predicted_set() {
878        let r = Registry::new();
879        let loc = Location::caller();
880        r.seal();
881        r.record_build_at(
882            loc,
883            lookup(Backend::Metal, Some(QuantKind::Q6K)),
884            Outcome::slow_path("CPU apply_batch"),
885        );
886        r.record_at(
887            loc,
888            lookup(Backend::Metal, Some(QuantKind::Q6K)),
889            Outcome::slow_path("CPU apply_batch"),
890        );
891        assert!(r.surprises().is_empty());
892    }
893
894    /// The dispatch path runs on rayon workers, so the first-sighting
895    /// write must be raced-into by many threads and still warn exactly
896    /// once, with every lookup counted.
897    #[test]
898    fn concurrent_dispatch_misses_warn_once_and_count_all() {
899        let r = std::sync::Arc::new(Registry::new());
900        let loc = Location::caller();
901        r.seal();
902        std::thread::scope(|s| {
903            for _ in 0..8 {
904                let r = std::sync::Arc::clone(&r);
905                s.spawn(move || {
906                    for _ in 0..250 {
907                        r.record_at(
908                            loc,
909                            lookup(Backend::Metal, Some(QuantKind::IQ1S)),
910                            Outcome::slow_path("CPU apply_batch"),
911                        );
912                    }
913                });
914            }
915        });
916        assert_eq!(r.surprises().len(), 1, "warned once across 8 threads");
917        let counted: u64 = r
918            .entries()
919            .iter()
920            .filter(|e| e.key.kind == Some(QuantKind::IQ1S))
921            .map(|e| e.count)
922            .sum();
923        assert_eq!(counted, 2000, "every lookup counted exactly once");
924    }
925
926    #[test]
927    fn the_report_names_the_call_site_and_the_quant_kind() {
928        let r = Registry::new();
929        r.record_build_at(
930            Location::caller(),
931            lookup(Backend::Metal, Some(QuantKind::Q5K)),
932            Outcome::slow_path("CPU apply_batch"),
933        );
934        let rendered = r.seal().render();
935        assert!(rendered.contains("Q5_K"), "{rendered}");
936        assert!(rendered.contains("kernel_registry.rs"), "{rendered}");
937        assert!(rendered.contains("ffn_down"), "{rendered}");
938        assert!(rendered.contains("CPU apply_batch"), "{rendered}");
939    }
940}