metrale_model_layers/layers/ops/
model_stats.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: [`ModelStats`]: diagnostic counters and one-shot latches owned by one model.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - A new `ModelStats` starts with zero counters and every latch unfired, so a second model
8//!   neither inherits the first one's counts nor finds its dumps already spent.
9//!
10//! `TransformerModel` owns one next to its [`ModelLevers`](super::ModelLevers) and lends it to
11//! each `ForwardContext` as `stats`. The counters are atomics because they are updated through that
12//! shared reference.
13
14use std::sync::atomic::AtomicU64;
15
16/// 2026-09-25: Per-model diagnostic state.
17#[derive(Debug, Default)]
18pub struct ModelStats {
19    /// 2026-09-25: MoE expert-union sampling under `ModelLevers::moe_union_stats`
20    /// (`moe/union_stats.rs`).
21    pub moe_union: MoeUnionStats,
22    /// 2026-09-25: One-shot latches for dumps and once-per-model log lines.
23    pub dumped: DumpLatches,
24}
25
26/// 2026-09-25: Expert-union sampling counters for one model: calls seen, calls sampled, and the
27/// running unique-expert and slot totals.
28#[derive(Debug, Default)]
29pub struct MoeUnionStats {
30    pub calls: AtomicU64,
31    pub samples: AtomicU64,
32    pub unique_sum: AtomicU64,
33    pub slots_sum: AtomicU64,
34}
35
36/// 2026-09-25: One-shot latches for one model, keyed by a `&'static str`
37/// ([`keyed`](DumpLatches::keyed)).
38#[derive(Debug, Default)]
39pub struct DumpLatches {
40    /// 2026-09-25: The keys that have fired.
41    fired: std::sync::Mutex<std::collections::BTreeSet<&'static str>>,
42}
43
44impl ModelStats {
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// 2026-09-25: `true` the first time this model reaches `key`, `false` after; the same latch
50    /// set as [`DumpLatches::keyed`]. Prefix keys by purpose (`"log:..."`, `"dump:..."`) so two
51    /// unrelated sites cannot share one.
52    pub fn once(&self, key: &'static str) -> bool {
53        self.dumped.keyed(key)
54    }
55}
56
57impl DumpLatches {
58    /// 2026-09-25: `true` exactly once per model for `key`. Calling it consumes the shot, so test
59    /// the dump's lever first; a disabled dump then does not spend it.
60    pub fn keyed(&self, key: &'static str) -> bool {
61        self.fired
62            .lock()
63            .expect("dump latches poisoned")
64            .insert(key)
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use std::sync::atomic::Ordering;
72
73    #[test]
74    fn a_keyed_latch_fires_once_per_key_per_model() {
75        let a = ModelStats::new();
76        assert!(a.dumped.keyed("dflash_block"));
77        assert!(!a.dumped.keyed("dflash_block"), "and only once");
78        assert!(a.dumped.keyed("dflash_ctx"), "a different dump is separate");
79        assert!(
80            ModelStats::new().dumped.keyed("dflash_block"),
81            "and a new model re-arms every key"
82        );
83    }
84
85    #[test]
86    fn two_models_count_expert_unions_independently() {
87        let a = ModelStats::new();
88        let b = ModelStats::new();
89        a.moe_union.calls.fetch_add(9, Ordering::Relaxed);
90        assert_eq!(a.moe_union.calls.load(Ordering::Relaxed), 9);
91        assert_eq!(
92            b.moe_union.calls.load(Ordering::Relaxed),
93            0,
94            "a second model starts clean rather than inheriting a mean"
95        );
96    }
97}
metrale_model_layers/layers/ops/model_stats.rs

metrale_model_layers/layers/ops/
model_stats.rs
