## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/trueno/src/brick/tracing.rs

    Line
    Count
    Source
    1
    //! Model-Level Inference Tracing (Phase 13, E.11)
    2
    //!
    3
    //! Comprehensive tracing system for transformer model inference:
    4
    //! - MLT-01: LayerActivationTrace - anomaly detection per layer
    5
    //! - MLT-02: AttentionWeightTrace - attention pattern analysis
    6
    //! - MLT-03: LogitEvolutionTrace - token probability evolution
    7
    //! - MLT-04: QuantizationErrorTrace - quantization quality metrics
    8
    //! - MLT-05: KvCacheStateTrace - KV cache efficiency tracking
    9
    //!
    10
    //! # Example
    11
    //!
    12
    //! ```rust,ignore
    13
    //! use trueno::brick::{ModelTracer, ModelTracerConfig};
    14
    //!
    15
    //! let config = ModelTracerConfig::lightweight();
    16
    //! let mut tracer = ModelTracer::new(config);
    17
    //!
    18
    //! tracer.begin_forward(position);
    19
    //! // ... forward pass with trace hooks ...
    20
    //! if let Some(anomaly) = tracer.end_forward() {
    21
    //!     log::warn!("Anomaly: {}", anomaly);
    22
    //! }
    23
    //! ```
    24
    25
    use std::fmt;
    26
    27
    use super::exec_graph::BrickId;
    28
    29
    // ============================================================================
    30
    // QuantType - Quantization type tracking
    31
    // ============================================================================
    32
    33
    /// Quantization type for tracking quantization errors (MLT-04).
    34
    ///
    35
    /// Note: Variant names follow GGML conventions (e.g., Q4_K) for interoperability.
    36
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
    37
    #[allow(non_camel_case_types)]
    38
    pub enum QuantType {
    39
        /// Full precision (FP32)
    40
        #[default]
    41
        F32,
    42
        /// Half precision (FP16)
    43
        F16,
    44
        /// Brain floating point (BF16)
    45
        Bf16,
    46
        /// 8-bit integer quantization
    47
        Q8_0,
    48
        /// 4-bit quantization (GGML)
    49
        Q4_0,
    50
        /// 4-bit quantization with k-quants
    51
        Q4_K,
    52
        /// 5-bit quantization with k-quants
    53
        Q5_K,
    54
        /// 6-bit quantization with k-quants
    55
        Q6_K,
    56
        /// 2-bit quantization
    57
        Q2_K,
    58
        /// 3-bit quantization
    59
        Q3_K,
    60
    }
    61
    62
    impl QuantType {
    63
        /// Get bits per element for this quantization type.
    64
    0
        pub fn bits_per_element(self) -> f32 {
    65
    0
            match self {
    66
    0
                Self::F32 => 32.0,
    67
    0
                Self::F16 | Self::Bf16 => 16.0,
    68
    0
                Self::Q8_0 => 8.0,
    69
    0
                Self::Q6_K => 6.5,
    70
    0
                Self::Q5_K => 5.5,
    71
    0
                Self::Q4_0 | Self::Q4_K => 4.5,
    72
    0
                Self::Q3_K => 3.5,
    73
    0
                Self::Q2_K => 2.5,
    74
            }
    75
    0
        }
    76
    77
        /// Get compression ratio vs FP32.
    78
    0
        pub fn compression_ratio(self) -> f32 {
    79
    0
            32.0 / self.bits_per_element()
    80
    0
        }
    81
    }
    82
    83
    // ============================================================================
    84
    // E.11.2: LayerActivationTrace (MLT-01)
    85
    // ============================================================================
    86
    87
    /// Statistics for a tensor without storing the tensor itself.
    88
    ///
    89
    /// Computes min, max, mean, std, L2 norm, NaN/Inf counts in a single pass.
    90
    /// Used for anomaly detection (explosion, vanishing gradients, NaN propagation).
    91
    ///
    92
    /// # Example
    93
    /// ```rust,ignore
    94
    /// let stats = TensorStats::from_slice(&tensor_data);
    95
    /// if stats.has_anomaly() {
    96
    ///     log::warn!("Anomaly detected: {}", stats.anomaly_description());
    97
    /// }
    98
    /// ```
    99
    #[derive(Debug, Clone, Default, PartialEq)]
    100
    pub struct TensorStats {
    101
        /// Number of elements analyzed
    102
        pub count: usize,
    103
        /// Minimum value (ignoring NaN/Inf)
    104
        pub min: f32,
    105
        /// Maximum value (ignoring NaN/Inf)
    106
        pub max: f32,
    107
        /// Mean value (ignoring NaN/Inf)
    108
        pub mean: f32,
    109
        /// Standard deviation (ignoring NaN/Inf)
    110
        pub std: f32,
    111
        /// Count of NaN values
    112
        pub nan_count: usize,
    113
        /// Count of Inf values
    114
        pub inf_count: usize,
    115
        /// L2 norm (sqrt of sum of squares)
    116
        pub l2_norm: f32,
    117
    }
    118
    119
    impl TensorStats {
    120
        /// Compute statistics from a slice in a single pass.
    121
        ///
    122
        /// Uses Welford's algorithm for numerically stable mean/variance.
    123
    0
        pub fn from_slice(data: &[f32]) -> Self {
    124
    0
            if data.is_empty() {
    125
    0
                return Self::default();
    126
    0
            }
    127
    128
    0
            let mut count = 0usize;
    129
    0
            let mut nan_count = 0usize;
    130
    0
            let mut inf_count = 0usize;
    131
    0
            let mut min = f32::MAX;
    132
    0
            let mut max = f32::MIN;
    133
    0
            let mut sum_sq = 0.0f64;
    134
    135
            // Welford's algorithm for online mean/variance
    136
    0
            let mut mean = 0.0f64;
    137
    0
            let mut m2 = 0.0f64;
    138
    139
    0
            for &val in data {
    140
    0
                if val.is_nan() {
    141
    0
                    nan_count += 1;
    142
    0
                    continue;
    143
    0
                }
    144
    0
                if val.is_infinite() {
    145
    0
                    inf_count += 1;
    146
    0
                    continue;
    147
    0
                }
    148
    149
    0
                count += 1;
    150
    0
                min = min.min(val);
    151
    0
                max = max.max(val);
    152
    0
                sum_sq += (val as f64) * (val as f64);
    153
    154
                // Welford's update
    155
    0
                let delta = val as f64 - mean;
    156
    0
                mean += delta / count as f64;
    157
    0
                let delta2 = val as f64 - mean;
    158
    0
                m2 += delta * delta2;
    159
            }
    160
    161
    0
            let std = if count > 1 {
    162
    0
                (m2 / (count - 1) as f64).sqrt() as f32
    163
            } else {
    164
    0
                0.0
    165
            };
    166
    167
    0
            let l2_norm = sum_sq.sqrt() as f32;
    168
    169
            Self {
    170
    0
                count: data.len(),
    171
    0
                min: if count > 0 { min } else { 0.0 },
    172
    0
                max: if count > 0 { max } else { 0.0 },
    173
    0
                mean: mean as f32,
    174
    0
                std,
    175
    0
                nan_count,
    176
    0
                inf_count,
    177
    0
                l2_norm,
    178
            }
    179
    0
        }
    180
    181
        /// Check if this tensor has any anomalies.
    182
        ///
    183
        /// Anomaly detection rules (from E.11.2):
    184
        /// - NaN detected: `nan_count > 0`
    185
        /// - Explosion: `max.abs() > 1e6` or `std > 1e4`
    186
        /// - Vanishing: `std < 1e-6` (should check after first few layers)
    187
    0
        pub fn has_anomaly(&self) -> bool {
    188
    0
            self.nan_count > 0
    189
    0
                || self.inf_count > 0
    190
    0
                || self.max.abs() > 1e6
    191
    0
                || self.min.abs() > 1e6
    192
    0
                || self.std > 1e4
    193
    0
        }
    194
    195
        /// Check if values are vanishing (for layers past warmup).
    196
    0
        pub fn is_vanishing(&self) -> bool {
    197
    0
            self.std < 1e-6 && self.count > 0
    198
    0
        }
    199
    200
        /// Get a description of any anomaly detected.
    201
    0
        pub fn anomaly_description(&self) -> Option<String> {
    202
    0
            if self.nan_count > 0 {
    203
    0
                return Some(format!("NaN detected: {} values", self.nan_count));
    204
    0
            }
    205
    0
            if self.inf_count > 0 {
    206
    0
                return Some(format!("Inf detected: {} values", self.inf_count));
    207
    0
            }
    208
    0
            if self.max.abs() > 1e6 || self.min.abs() > 1e6 {
    209
    0
                return Some(format!(
    210
    0
                    "Explosion: min={:.2e}, max={:.2e}",
    211
    0
                    self.min, self.max
    212
    0
                ));
    213
    0
            }
    214
    0
            if self.std > 1e4 {
    215
    0
                return Some(format!("High variance: std={:.2e}", self.std));
    216
    0
            }
    217
    0
            None
    218
    0
        }
    219
    }
    220
    221
    /// Activation trace for a single transformer layer.
    222
    ///
    223
    /// Records tensor statistics at each stage of a transformer layer:
    224
    /// input → norm → attention → residual → ffn → output
    225
    #[derive(Debug, Clone, Default)]
    226
    pub struct LayerActivationTrace {
    227
        /// Layer index (0-indexed)
    228
        pub layer_idx: usize,
    229
        /// Input hidden state statistics
    230
        pub input_stats: TensorStats,
    231
        /// After RMSNorm/LayerNorm statistics
    232
        pub post_norm_stats: TensorStats,
    233
        /// After attention statistics
    234
        pub post_attn_stats: TensorStats,
    235
        /// After FFN statistics
    236
        pub post_ffn_stats: TensorStats,
    237
        /// Output hidden state statistics
    238
        pub output_stats: TensorStats,
    239
        /// Residual connection magnitude ratio (output_norm / (output_norm + attn_norm))
    240
        pub residual_ratio: f32,
    241
    }
    242
    243
    impl LayerActivationTrace {
    244
        /// Create a new layer activation trace.
    245
    0
        pub fn new(layer_idx: usize) -> Self {
    246
    0
            Self {
    247
    0
                layer_idx,
    248
    0
                ..Default::default()
    249
    0
            }
    250
    0
        }
    251
    252
        /// Check if this layer has any anomalies.
    253
    0
        pub fn has_anomaly(&self) -> bool {
    254
    0
            self.input_stats.has_anomaly()
    255
    0
                || self.post_norm_stats.has_anomaly()
    256
    0
                || self.post_attn_stats.has_anomaly()
    257
    0
                || self.post_ffn_stats.has_anomaly()
    258
    0
                || self.output_stats.has_anomaly()
    259
    0
                || self.residual_ratio > 0.99 // Skip connection bypass
    260
    0
        }
    261
    262
        /// Get anomaly description for this layer.
    263
    0
        pub fn anomaly_description(&self) -> Option<String> {
    264
    0
            if let Some(desc) = self.input_stats.anomaly_description() {
    265
    0
                return Some(format!("Layer {} input: {}", self.layer_idx, desc));
    266
    0
            }
    267
    0
            if let Some(desc) = self.post_norm_stats.anomaly_description() {
    268
    0
                return Some(format!("Layer {} post_norm: {}", self.layer_idx, desc));
    269
    0
            }
    270
    0
            if let Some(desc) = self.post_attn_stats.anomaly_description() {
    271
    0
                return Some(format!("Layer {} post_attn: {}", self.layer_idx, desc));
    272
    0
            }
    273
    0
            if let Some(desc) = self.post_ffn_stats.anomaly_description() {
    274
    0
                return Some(format!("Layer {} post_ffn: {}", self.layer_idx, desc));
    275
    0
            }
    276
    0
            if let Some(desc) = self.output_stats.anomaly_description() {
    277
    0
                return Some(format!("Layer {} output: {}", self.layer_idx, desc));
    278
    0
            }
    279
    0
            if self.residual_ratio > 0.99 {
    280
    0
                return Some(format!(
    281
    0
                    "Layer {} residual dominance: ratio={:.4}",
    282
    0
                    self.layer_idx, self.residual_ratio
    283
    0
                ));
    284
    0
            }
    285
    0
            None
    286
    0
        }
    287
    }
    288
    289
    /// Full model activation trace for one forward pass.
    290
    #[derive(Debug, Clone, Default)]
    291
    pub struct ModelActivationTrace {
    292
        /// Per-layer activation traces
    293
        pub layers: Vec<LayerActivationTrace>,
    294
        /// Embedding output statistics
    295
        pub embedding_stats: TensorStats,
    296
        /// Final logits statistics
    297
        pub logits_stats: TensorStats,
    298
        /// Whether any anomaly was detected
    299
        pub has_anomaly: bool,
    300
        /// Description of first anomaly found
    301
        pub anomaly_desc: Option<String>,
    302
    }
    303
    304
    impl ModelActivationTrace {
    305
        /// Create a new model activation trace with expected layer count.
    306
    0
        pub fn with_capacity(num_layers: usize) -> Self {
    307
    0
            Self {
    308
    0
                layers: Vec::with_capacity(num_layers),
    309
    0
                ..Default::default()
    310
    0
            }
    311
    0
        }
    312
    313
        /// Add a layer trace.
    314
    0
        pub fn add_layer(&mut self, trace: LayerActivationTrace) {
    315
    0
            if !self.has_anomaly {
    316
    0
                if let Some(desc) = trace.anomaly_description() {
    317
    0
                    self.has_anomaly = true;
    318
    0
                    self.anomaly_desc = Some(desc);
    319
    0
                }
    320
    0
            }
    321
    0
            self.layers.push(trace);
    322
    0
        }
    323
    324
        /// Finalize the trace and check embedding/logits.
    325
    0
        pub fn finalize(&mut self) {
    326
    0
            if !self.has_anomaly {
    327
    0
                if let Some(desc) = self.embedding_stats.anomaly_description() {
    328
    0
                    self.has_anomaly = true;
    329
    0
                    self.anomaly_desc = Some(format!("Embedding: {}", desc));
    330
    0
                }
    331
    0
            }
    332
    0
            if !self.has_anomaly {
    333
    0
                if let Some(desc) = self.logits_stats.anomaly_description() {
    334
    0
                    self.has_anomaly = true;
    335
    0
                    self.anomaly_desc = Some(format!("Logits: {}", desc));
    336
    0
                }
    337
    0
            }
    338
    0
        }
    339
    }
    340
    341
    // ============================================================================
    342
    // E.11.3: AttentionWeightTrace (MLT-02)
    343
    // ============================================================================
    344
    345
    /// Sparse attention weight storage for a single head.
    346
    ///
    347
    /// Records top-k attended positions to avoid storing the full attention matrix.
    348
    /// Useful for debugging repetition, context loss, and attention sinks.
    349
    #[derive(Debug, Clone, Default)]
    350
    pub struct AttentionWeightTrace {
    351
        /// Layer index
    352
        pub layer_idx: usize,
    353
        /// Head index within the layer
    354
        pub head_idx: usize,
    355
        /// Query position (current token being generated)
    356
        pub query_pos: usize,
    357
        /// Top-k attended positions (sorted by weight descending)
    358
        pub top_k_positions: Vec<usize>,
    359
        /// Corresponding attention weights
    360
        pub top_k_weights: Vec<f32>,
    361
        /// Sum of weights outside top-k (attention mass lost to tail)
    362
        pub tail_mass: f32,
    363
        /// Entropy of attention distribution (higher = more uniform)
    364
        pub entropy: f32,
    365
    }
    366
    367
    impl AttentionWeightTrace {
    368
        /// Create from full attention weights, extracting top-k.
    369
    0
        pub fn from_weights(
    370
    0
            layer_idx: usize,
    371
    0
            head_idx: usize,
    372
    0
            query_pos: usize,
    373
    0
            weights: &[f32],
    374
    0
            k: usize,
    375
    0
        ) -> Self {
    376
    0
            let k = k.min(weights.len());
    377
    378
            // Create position-weight pairs and sort by weight descending
    379
    0
            let mut pairs: Vec<(usize, f32)> = weights.iter().copied().enumerate().collect();
    380
    0
            pairs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    381
    382
    0
            let top_k_positions: Vec<usize> = pairs.iter().take(k).map(|(pos, _)| *pos).collect();
    383
    0
            let top_k_weights: Vec<f32> = pairs.iter().take(k).map(|(_, w)| *w).collect();
    384
    385
    0
            let top_k_mass: f32 = top_k_weights.iter().sum();
    386
    0
            let total_mass: f32 = weights.iter().sum();
    387
    0
            let tail_mass = (total_mass - top_k_mass).max(0.0);
    388
    389
            // Compute entropy: H = -sum(p * log(p)) for non-zero probabilities
    390
    0
            let entropy = weights
    391
    0
                .iter()
    392
    0
                .filter(|&&w| w > 1e-10)
    393
    0
                .map(|&w| -w * w.ln())
    394
    0
                .sum();
    395
    396
    0
            Self {
    397
    0
                layer_idx,
    398
    0
                head_idx,
    399
    0
                query_pos,
    400
    0
                top_k_positions,
    401
    0
                top_k_weights,
    402
    0
                tail_mass,
    403
    0
                entropy,
    404
    0
            }
    405
    0
        }
    406
    407
        /// Check if attention is concentrated on first position (attention sink).
    408
    0
        pub fn is_attention_sink(&self, threshold: f32) -> bool {
    409
    0
            self.top_k_positions.first() == Some(&0)
    410
    0
                && self.top_k_weights.first().copied().unwrap_or(0.0) > threshold
    411
    0
        }
    412
    413
        /// Check if attention is too uniform (confused model).
    414
    0
        pub fn is_uniform(&self, entropy_threshold: f32) -> bool {
    415
    0
            self.entropy > entropy_threshold
    416
    0
        }
    417
    418
        /// Check for repetition pattern (high weight on recent positions).
    419
    0
        pub fn has_recency_bias(&self, recency_window: usize, threshold: f32) -> bool {
    420
    0
            if self.query_pos == 0 {
    421
    0
                return false;
    422
    0
            }
    423
    0
            let recency_start = self.query_pos.saturating_sub(recency_window);
    424
    0
            let recent_mass: f32 = self
    425
    0
                .top_k_positions
    426
    0
                .iter()
    427
    0
                .zip(self.top_k_weights.iter())
    428
    0
                .filter(|(pos, _)| **pos >= recency_start)
    429
    0
                .map(|(_, w)| w)
    430
    0
                .sum();
    431
    0
            recent_mass > threshold
    432
    0
        }
    433
    }
    434
    435
    /// Configuration for attention weight tracing.
    436
    #[derive(Debug, Clone)]
    437
    pub struct AttentionTraceConfig {
    438
        /// Number of top positions to record per head
    439
        pub top_k: usize,
    440
        /// Layers to trace (None = all)
    441
        pub layers: Option<Vec<usize>>,
    442
        /// Heads to trace (None = all)
    443
        pub heads: Option<Vec<usize>>,
    444
        /// Minimum weight to consider (positions with weight below this are ignored)
    445
        pub weight_threshold: f32,
    446
    }
    447
    448
    impl Default for AttentionTraceConfig {
    449
    0
        fn default() -> Self {
    450
    0
            Self {
    451
    0
                top_k: 10,
    452
    0
                layers: None,
    453
    0
                heads: None,
    454
    0
                weight_threshold: 0.01,
    455
    0
            }
    456
    0
        }
    457
    }
    458
    459
    impl AttentionTraceConfig {
    460
        /// Check if a layer should be traced.
    461
    0
        pub fn should_trace_layer(&self, layer_idx: usize) -> bool {
    462
    0
            self.layers
    463
    0
                .as_ref()
    464
    0
                .is_none_or(|layers| layers.contains(&layer_idx))
    465
    0
        }
    466
    467
        /// Check if a head should be traced.
    468
    0
        pub fn should_trace_head(&self, head_idx: usize) -> bool {
    469
    0
            self.heads
    470
    0
                .as_ref()
    471
    0
                .is_none_or(|heads| heads.contains(&head_idx))
    472
    0
        }
    473
    }
    474
    475
    // ============================================================================
    476
    // E.11.4: LogitEvolutionTrace (MLT-03)
    477
    // ============================================================================
    478
    479
    /// Logit evolution for a single token through layers.
    480
    ///
    481
    /// Tracks how a token's logit value and rank change as hidden states
    482
    /// pass through transformer layers.
    483
    #[derive(Debug, Clone, Default)]
    484
    pub struct TokenLogitEvolution {
    485
        /// Token ID being tracked
    486
        pub token_id: u32,
    487
        /// Token string representation (for display)
    488
        pub token_str: String,
    489
        /// Logit value after each layer's contribution
    490
        pub per_layer_logit: Vec<f32>,
    491
        /// Rank among vocabulary at each layer (0 = highest probability)
    492
        pub per_layer_rank: Vec<usize>,
    493
        /// Final probability after softmax
    494
        pub final_probability: f32,
    495
        /// Final rank (0 = selected token)
    496
        pub final_rank: usize,
    497
    }
    498
    499
    impl TokenLogitEvolution {
    500
        /// Create a new token evolution tracker.
    501
    0
        pub fn new(token_id: u32, token_str: String) -> Self {
    502
    0
            Self {
    503
    0
                token_id,
    504
    0
                token_str,
    505
    0
                ..Default::default()
    506
    0
            }
    507
    0
        }
    508
    509
        /// Record logit value at a layer.
    510
    0
        pub fn record_layer(&mut self, logit: f32, rank: usize) {
    511
    0
            self.per_layer_logit.push(logit);
    512
    0
            self.per_layer_rank.push(rank);
    513
    0
        }
    514
    515
        /// Get the layer where this token's rank changed most dramatically.
    516
    0
        pub fn decisive_layer(&self) -> Option<usize> {
    517
    0
            if self.per_layer_rank.len() < 2 {
    518
    0
                return None;
    519
    0
            }
    520
    521
    0
            let mut max_change = 0i64;
    522
    0
            let mut decisive = 0;
    523
    524
    0
            for i in 1..self.per_layer_rank.len() {
    525
    0
                let change =
    526
    0
                    (self.per_layer_rank[i] as i64 - self.per_layer_rank[i - 1] as i64).abs();
    527
    0
                if change > max_change {
    528
    0
                    max_change = change;
    529
    0
                    decisive = i;
    530
    0
                }
    531
            }
    532
    533
    0
            Some(decisive)
    534
    0
        }
    535
    }
    536
    537
    /// Full logit trace for one generation step.
    538
    #[derive(Debug, Clone, Default)]
    539
    pub struct LogitEvolutionTrace {
    540
        /// Position being generated
    541
        pub position: usize,
    542
        /// Tokens being tracked (typically top-k candidates + ground truth)
    543
        pub tracked_tokens: Vec<TokenLogitEvolution>,
    544
        /// Which layer had the largest impact on the selected token
    545
        pub decisive_layer: usize,
    546
        /// Temperature used for sampling
    547
        pub temperature: f32,
    548
        /// Top-p (nucleus) value used
    549
        pub top_p: f32,
    550
    }
    551
    552
    impl LogitEvolutionTrace {
    553
        /// Create a new logit evolution trace.
    554
    0
        pub fn new(position: usize, temperature: f32, top_p: f32) -> Self {
    555
    0
            Self {
    556
    0
                position,
    557
    0
                temperature,
    558
    0
                top_p,
    559
    0
                ..Default::default()
    560
    0
            }
    561
    0
        }
    562
    563
        /// Add a token to track.
    564
    0
        pub fn track_token(&mut self, token_id: u32, token_str: String) -> &mut TokenLogitEvolution {
    565
    0
            self.tracked_tokens
    566
    0
                .push(TokenLogitEvolution::new(token_id, token_str));
    567
    0
            self.tracked_tokens.last_mut().expect("invariant: just pushed")
    568
    0
        }
    569
    570
        /// Compute rank of a token in a logit distribution.
    571
    0
        pub fn compute_rank(logits: &[f32], token_id: u32) -> usize {
    572
    0
            let target_logit = logits.get(token_id as usize).copied().unwrap_or(f32::MIN);
    573
    574
    0
            logits.iter().filter(|&&l| l > target_logit).count()
    575
    0
        }
    576
    577
        /// Finalize the trace after generation completes.
    578
    0
        pub fn finalize(&mut self, selected_token_id: u32) {
    579
            // Find the decisive layer for the selected token
    580
    0
            for token in &self.tracked_tokens {
    581
    0
                if token.token_id == selected_token_id {
    582
    0
                    if let Some(layer) = token.decisive_layer() {
    583
    0
                        self.decisive_layer = layer;
    584
    0
                    }
    585
    0
                    break;
    586
    0
                }
    587
            }
    588
    0
        }
    589
    }
    590
    591
    // ============================================================================
    592
    // E.11.5: QuantizationErrorTrace (MLT-04)
    593
    // ============================================================================
    594
    595
    /// Quantization error measurement for a single operation.
    596
    ///
    597
    /// Compares quantized computation against FP32 reference using multiple metrics.
    598
    #[derive(Debug, Clone)]
    599
    pub struct QuantizationErrorTrace {
    600
        /// Brick type being measured
    601
        pub brick_id: BrickId,
    602
        /// Layer index
    603
        pub layer_idx: usize,
    604
        /// Mean squared error vs FP32 reference
    605
        pub mse: f32,
    606
        /// Maximum absolute error
    607
        pub max_abs_error: f32,
    608
        /// Cosine similarity (1.0 = perfect match)
    609
        pub cosine_similarity: f32,
    610
        /// Signal-to-noise ratio in dB
    611
        pub snr_db: f32,
    612
        /// Quantization type used
    613
        pub quant_type: QuantType,
    614
    }
    615
    616
    impl QuantizationErrorTrace {
    617
        /// Compute error metrics between quantized and reference outputs.
    618
    0
        pub fn compute(
    619
    0
            brick_id: BrickId,
    620
    0
            layer_idx: usize,
    621
    0
            quantized: &[f32],
    622
    0
            reference: &[f32],
    623
    0
            quant_type: QuantType,
    624
    0
        ) -> Self {
    625
    0
            assert_eq!(quantized.len(), reference.len(), "Length mismatch");
    626
    0
            let n = quantized.len();
    627
    0
            if n == 0 {
    628
    0
                return Self {
    629
    0
                    brick_id,
    630
    0
                    layer_idx,
    631
    0
                    mse: 0.0,
    632
    0
                    max_abs_error: 0.0,
    633
    0
                    cosine_similarity: 1.0, // Perfect match when both empty
    634
    0
                    snr_db: f32::INFINITY,
    635
    0
                    quant_type,
    636
    0
                };
    637
    0
            }
    638
    639
            // MSE and max abs error
    640
    0
            let mut sum_sq_error = 0.0f64;
    641
    0
            let mut max_abs_error = 0.0f32;
    642
    0
            for (q, r) in quantized.iter().zip(reference.iter()) {
    643
    0
                let error = q - r;
    644
    0
                sum_sq_error += (error as f64) * (error as f64);
    645
    0
                max_abs_error = max_abs_error.max(error.abs());
    646
    0
            }
    647
    0
            let mse = (sum_sq_error / n as f64) as f32;
    648
    649
            // Cosine similarity
    650
    0
            let mut dot = 0.0f64;
    651
    0
            let mut norm_q = 0.0f64;
    652
    0
            let mut norm_r = 0.0f64;
    653
    0
            for (q, r) in quantized.iter().zip(reference.iter()) {
    654
    0
                dot += (*q as f64) * (*r as f64);
    655
    0
                norm_q += (*q as f64) * (*q as f64);
    656
    0
                norm_r += (*r as f64) * (*r as f64);
    657
    0
            }
    658
    0
            let cosine_similarity = if norm_q > 0.0 && norm_r > 0.0 {
    659
    0
                (dot / (norm_q.sqrt() * norm_r.sqrt())) as f32
    660
            } else {
    661
    0
                0.0
    662
            };
    663
    664
            // SNR in dB: 10 * log10(signal_power / noise_power)
    665
    0
            let signal_power = norm_r / n as f64;
    666
    0
            let noise_power = sum_sq_error / n as f64;
    667
    0
            let snr_db = if noise_power > 1e-10 {
    668
    0
                (10.0 * (signal_power / noise_power).log10()) as f32
    669
            } else {
    670
    0
                f32::INFINITY
    671
            };
    672
    673
    0
            Self {
    674
    0
                brick_id,
    675
    0
                layer_idx,
    676
    0
                mse,
    677
    0
                max_abs_error,
    678
    0
                cosine_similarity,
    679
    0
                snr_db,
    680
    0
                quant_type,
    681
    0
            }
    682
    0
        }
    683
    684
        /// Check if error is acceptable (cosine > 0.995).
    685
    0
        pub fn is_acceptable(&self) -> bool {
    686
    0
            self.cosine_similarity > 0.995
    687
    0
        }
    688
    689
        /// Check if error is in warning zone (0.99 < cosine < 0.995).
    690
    0
        pub fn is_warning(&self) -> bool {
    691
    0
            self.cosine_similarity > 0.99 && self.cosine_similarity <= 0.995
    692
    0
        }
    693
    694
        /// Check if error is critical (cosine < 0.99).
    695
    0
        pub fn is_critical(&self) -> bool {
    696
    0
            self.cosine_similarity < 0.99
    697
    0
        }
    698
    }
    699
    700
    /// Cumulative quantization error across an entire model.
    701
    #[derive(Debug, Clone, Default)]
    702
    pub struct ModelQuantizationError {
    703
        /// Per-brick error traces
    704
        pub brick_errors: Vec<QuantizationErrorTrace>,
    705
        /// Overall cosine similarity of final logits
    706
        pub logits_cosine: f32,
    707
        /// KL divergence of output probability distributions
    708
        pub output_kl_divergence: f32,
    709
        /// Perplexity difference (PPL_quant - PPL_fp32)
    710
        pub perplexity_delta: f32,
    711
    }
    712
    713
    impl ModelQuantizationError {
    714
        /// Add a brick error trace.
    715
    0
        pub fn add_error(&mut self, trace: QuantizationErrorTrace) {
    716
    0
            self.brick_errors.push(trace);
    717
    0
        }
    718
    719
        /// Get count of critical errors.
    720
    0
        pub fn critical_count(&self) -> usize {
    721
    0
            self.brick_errors.iter().filter(|e| e.is_critical()).count()
    722
    0
        }
    723
    724
        /// Get count of warning errors.
    725
    0
        pub fn warning_count(&self) -> usize {
    726
    0
            self.brick_errors.iter().filter(|e| e.is_warning()).count()
    727
    0
        }
    728
    729
        /// Get worst brick by cosine similarity.
    730
    0
        pub fn worst_brick(&self) -> Option<&QuantizationErrorTrace> {
    731
    0
            self.brick_errors
    732
    0
                .iter()
    733
    0
                .min_by(|a, b| a.cosine_similarity.partial_cmp(&b.cosine_similarity).unwrap_or(std::cmp::Ordering::Equal))
    734
    0
        }
    735
    }
    736
    737
    // ============================================================================
    738
    // E.11.6: KvCacheStateTrace (MLT-05)
    739
    // ============================================================================
    740
    741
    /// KV cache state at a single generation step.
    742
    #[derive(Debug, Clone, Default)]
    743
    pub struct KvCacheStateTrace {
    744
        /// Generation step (0-indexed)
    745
        pub step: usize,
    746
        /// Total cache size in bytes
    747
        pub cache_size_bytes: usize,
    748
        /// Number of valid (filled) positions in cache
    749
        pub valid_positions: usize,
    750
        /// Maximum positions (context window size)
    751
        pub max_positions: usize,
    752
        /// Evictions performed this step
    753
        pub evictions_this_step: usize,
    754
        /// Cache hit rate (reused positions / total lookups)
    755
        pub cache_hit_rate: f32,
    756
        /// Oldest position still in cache
    757
        pub oldest_position: usize,
    758
        /// Memory fragmentation (0.0 = compact, 1.0 = fully scattered)
    759
        pub fragmentation: f32,
    760
        /// Positions accessed this step (for locality analysis)
    761
        pub accessed_positions: Vec<usize>,
    762
    }
    763
    764
    impl KvCacheStateTrace {
    765
        /// Create a new trace for a step.
    766
    0
        pub fn new(step: usize, max_positions: usize) -> Self {
    767
    0
            Self {
    768
    0
                step,
    769
    0
                max_positions,
    770
    0
                ..Default::default()
    771
    0
            }
    772
    0
        }
    773
    774
        /// Check if context window is exhausted.
    775
    0
        pub fn is_window_exhausted(&self) -> bool {
    776
    0
            self.valid_positions >= self.max_positions
    777
    0
        }
    778
    779
        /// Get cache utilization ratio.
    780
    0
        pub fn utilization(&self) -> f32 {
    781
    0
            if self.max_positions == 0 {
    782
    0
                return 0.0;
    783
    0
            }
    784
    0
            self.valid_positions as f32 / self.max_positions as f32
    785
    0
        }
    786
    }
    787
    788
    /// Full KV cache trace for a generation session.
    789
    #[derive(Debug, Clone, Default)]
    790
    pub struct KvCacheSessionTrace {
    791
        /// Per-step traces
    792
        pub steps: Vec<KvCacheStateTrace>,
    793
        /// Total evictions across the session
    794
        pub total_evictions: usize,
    795
        /// Average cache hit rate
    796
        pub avg_hit_rate: f32,
    797
        /// Peak memory usage in bytes
    798
        pub peak_memory_bytes: usize,
    799
    }
    800
    801
    impl KvCacheSessionTrace {
    802
        /// Add a step trace.
    803
    0
        pub fn add_step(&mut self, trace: KvCacheStateTrace) {
    804
    0
            self.total_evictions += trace.evictions_this_step;
    805
    0
            self.peak_memory_bytes = self.peak_memory_bytes.max(trace.cache_size_bytes);
    806
    807
            // Update rolling average
    808
    0
            let n = self.steps.len() as f32 + 1.0;
    809
    0
            self.avg_hit_rate =
    810
    0
                (self.avg_hit_rate * (n - 1.0) + trace.cache_hit_rate) / n;
    811
    812
    0
            self.steps.push(trace);
    813
    0
        }
    814
    815
        /// Check if eviction rate is concerning (>10% of steps).
    816
    0
        pub fn has_high_eviction_rate(&self) -> bool {
    817
    0
            if self.steps.is_empty() {
    818
    0
                return false;
    819
    0
            }
    820
    0
            let eviction_steps = self.steps.iter().filter(|s| s.evictions_this_step > 0).count();
    821
    0
            eviction_steps as f32 / self.steps.len() as f32 > 0.1
    822
    0
        }
    823
    824
        /// Check if KV cache is thrashing (high evictions + low hit rate).
    825
        ///
    826
        /// Returns true if the recent window shows both high eviction rate and low hit rate.
    827
        /// Uses all available steps if fewer than `window` steps exist.
    828
        ///
    829
        /// # Arguments
    830
        /// - `window`: Number of recent steps to consider (uses available if fewer)
    831
        /// - `min_hit_rate`: Minimum acceptable hit rate (0.0-1.0)
    832
    0
        pub fn has_thrashing(&self, window: usize, min_hit_rate: f32) -> bool {
    833
    0
            if self.steps.is_empty() {
    834
    0
                return false;
    835
    0
            }
    836
    837
            // Use all steps if fewer than window
    838
    0
            let actual_window = std::cmp::min(window, self.steps.len());
    839
    0
            let recent_steps = &self.steps[self.steps.len() - actual_window..];
    840
    0
            let recent_evictions: usize = recent_steps.iter().map(|s| s.evictions_this_step).sum();
    841
    0
            let recent_hit_rate: f32 =
    842
    0
                recent_steps.iter().map(|s| s.cache_hit_rate).sum::<f32>() / actual_window as f32;
    843
    844
            // Thrashing: more than half the steps have evictions AND hit rate below threshold
    845
    0
            recent_evictions > actual_window / 2 && recent_hit_rate < min_hit_rate
    846
    0
        }
    847
    }
    848
    849
    // ============================================================================
    850
    // E.11.7: Unified ModelTracer
    851
    // ============================================================================
    852
    853
    /// Configuration for model-level tracing.
    854
    #[derive(Debug, Clone, Default)]
    855
    pub struct ModelTracerConfig {
    856
        /// Enable layer activation tracing (MLT-01)
    857
        pub trace_activations: bool,
    858
        /// Enable attention weight tracing (MLT-02)
    859
        pub trace_attention: bool,
    860
        /// Attention trace configuration
    861
        pub attention_config: AttentionTraceConfig,
    862
        /// Enable logit evolution tracing (MLT-03)
    863
        pub trace_logits: bool,
    864
        /// Specific tokens to track (None = auto-select top-k)
    865
        pub tracked_tokens: Option<Vec<u32>>,
    866
        /// Enable quantization error tracing (MLT-04) - expensive!
    867
        pub trace_quant_error: bool,
    868
        /// Enable KV cache state tracing (MLT-05)
    869
        pub trace_kv_cache: bool,
    870
    }
    871
    872
    impl ModelTracerConfig {
    873
        /// Create a config that traces everything (for debugging).
    874
    0
        pub fn full() -> Self {
    875
    0
            Self {
    876
    0
                trace_activations: true,
    877
    0
                trace_attention: true,
    878
    0
                attention_config: AttentionTraceConfig::default(),
    879
    0
                trace_logits: true,
    880
    0
                tracked_tokens: None,
    881
    0
                trace_quant_error: true,
    882
    0
                trace_kv_cache: true,
    883
    0
            }
    884
    0
        }
    885
    886
        /// Create a lightweight config (activations + KV cache only).
    887
    0
        pub fn lightweight() -> Self {
    888
    0
            Self {
    889
    0
                trace_activations: true,
    890
    0
                trace_kv_cache: true,
    891
    0
                ..Default::default()
    892
    0
            }
    893
    0
        }
    894
    895
        /// Check if any tracing is enabled.
    896
    0
        pub fn is_enabled(&self) -> bool {
    897
    0
            self.trace_activations
    898
    0
                || self.trace_attention
    899
    0
                || self.trace_logits
    900
    0
                || self.trace_quant_error
    901
    0
                || self.trace_kv_cache
    902
    0
        }
    903
    }
    904
    905
    /// Unified model tracer that coordinates all trace types.
    906
    ///
    907
    /// # Example
    908
    /// ```rust,ignore
    909
    /// let config = ModelTracerConfig::lightweight();
    910
    /// let mut tracer = ModelTracer::new(config);
    911
    ///
    912
    /// tracer.begin_forward(position);
    913
    /// // ... forward pass with trace hooks ...
    914
    /// if let Some(anomaly) = tracer.end_forward() {
    915
    ///     log::warn!("Anomaly: {}", anomaly);
    916
    /// }
    917
    /// ```
    918
    pub struct ModelTracer {
    919
        config: ModelTracerConfig,
    920
        /// Current forward pass position
    921
        current_position: usize,
    922
        /// Accumulated activation traces
    923
        activation_traces: Vec<ModelActivationTrace>,
    924
        /// Current activation trace (in progress)
    925
        current_activation_trace: Option<ModelActivationTrace>,
    926
        /// Accumulated attention traces
    927
        attention_traces: Vec<AttentionWeightTrace>,
    928
        /// Accumulated logit evolution traces
    929
        logit_traces: Vec<LogitEvolutionTrace>,
    930
        /// Current logit trace (in progress)
    931
        current_logit_trace: Option<LogitEvolutionTrace>,
    932
        /// Accumulated quantization error traces
    933
        quant_traces: Vec<ModelQuantizationError>,
    934
        /// KV cache session trace
    935
        kv_trace: KvCacheSessionTrace,
    936
    }
    937
    938
    impl ModelTracer {
    939
        /// Create a new tracer with the given configuration.
    940
    0
        pub fn new(config: ModelTracerConfig) -> Self {
    941
    0
            Self {
    942
    0
                config,
    943
    0
                current_position: 0,
    944
    0
                activation_traces: Vec::new(),
    945
    0
                current_activation_trace: None,
    946
    0
                attention_traces: Vec::new(),
    947
    0
                logit_traces: Vec::new(),
    948
    0
                current_logit_trace: None,
    949
    0
                quant_traces: Vec::new(),
    950
    0
                kv_trace: KvCacheSessionTrace::default(),
    951
    0
            }
    952
    0
        }
    953
    954
        /// Get the configuration.
    955
    0
        pub fn config(&self) -> &ModelTracerConfig {
    956
    0
            &self.config
    957
    0
        }
    958
    959
        /// Get a reference to the current logit trace (if any).
    960
    0
        pub fn current_logit_trace(&self) -> Option<&LogitEvolutionTrace> {
    961
    0
            self.current_logit_trace.as_ref()
    962
    0
        }
    963
    964
        /// Set the current logit trace (for testing purposes).
    965
    0
        pub fn set_current_logit_trace(&mut self, trace: Option<LogitEvolutionTrace>) {
    966
    0
            self.current_logit_trace = trace;
    967
    0
        }
    968
    969
        /// Begin a forward pass at the given position.
    970
    0
        pub fn begin_forward(&mut self, position: usize) {
    971
    0
            self.current_position = position;
    972
    973
    0
            if self.config.trace_activations {
    974
    0
                self.current_activation_trace = Some(ModelActivationTrace::default());
    975
    0
            }
    976
    977
    0
            if self.config.trace_logits {
    978
    0
                self.current_logit_trace = Some(LogitEvolutionTrace::new(position, 1.0, 1.0));
    979
    0
            }
    980
    0
        }
    981
    982
        /// Record layer activation (called by executor after each layer).
    983
    0
        pub fn record_layer_activation(&mut self, trace: LayerActivationTrace) {
    984
    0
            if let Some(ref mut activation) = self.current_activation_trace {
    985
    0
                activation.add_layer(trace);
    986
    0
            }
    987
    0
        }
    988
    989
        /// Record attention weights (called by attention brick).
    990
    0
        pub fn record_attention(&mut self, trace: AttentionWeightTrace) {
    991
    0
            if self.config.trace_attention {
    992
    0
                self.attention_traces.push(trace);
    993
    0
            }
    994
    0
        }
    995
    996
        /// Record logit state at a layer (called by lm_head or probe).
    997
    0
        pub fn record_logits(&mut self, layer_idx: usize, logits: &[f32]) {
    998
    0
            if let Some(ref mut logit_trace) = self.current_logit_trace {
    999
    0
                for token_evo in &mut logit_trace.tracked_tokens {
    1000
    0
                    let logit = logits.get(token_evo.token_id as usize).copied().unwrap_or(0.0);
    1001
    0
                    let rank = LogitEvolutionTrace::compute_rank(logits, token_evo.token_id);
    1002
    0
                    token_evo.record_layer(logit, rank);
    1003
    0
                }
    1004
                // Store decisive layer based on rank changes
    1005
    0
                logit_trace.decisive_layer = layer_idx;
    1006
    0
            }
    1007
    0
        }
    1008
    1009
        /// Record KV cache state (called after each generation step).
    1010
    0
        pub fn record_kv_state(&mut self, trace: KvCacheStateTrace) {
    1011
    0
            if self.config.trace_kv_cache {
    1012
    0
                self.kv_trace.add_step(trace);
    1013
    0
            }
    1014
    0
        }
    1015
    1016
        /// Record quantization error for a brick.
    1017
    0
        pub fn record_quant_error(&mut self, trace: QuantizationErrorTrace) {
    1018
    0
            if self.config.trace_quant_error {
    1019
    0
                if self.quant_traces.is_empty() {
    1020
    0
                    self.quant_traces.push(ModelQuantizationError::default());
    1021
    0
                }
    1022
    0
                if let Some(model_error) = self.quant_traces.last_mut() {
    1023
    0
                    model_error.add_error(trace);
    1024
    0
                }
    1025
    0
            }
    1026
    0
        }
    1027
    1028
        /// Complete forward pass and check for anomalies.
    1029
        ///
    1030
        /// Returns a description of the first anomaly detected, if any.
    1031
    0
        pub fn end_forward(&mut self) -> Option<String> {
    1032
    0
            let mut anomaly = None;
    1033
    1034
            // Finalize activation trace
    1035
    0
            if let Some(mut trace) = self.current_activation_trace.take() {
    1036
    0
                trace.finalize();
    1037
    0
                if trace.has_anomaly {
    1038
    0
                    anomaly = trace.anomaly_desc.clone();
    1039
    0
                }
    1040
    0
                self.activation_traces.push(trace);
    1041
    0
            }
    1042
    1043
            // Finalize logit trace
    1044
    0
            if let Some(trace) = self.current_logit_trace.take() {
    1045
    0
                self.logit_traces.push(trace);
    1046
    0
            }
    1047
    1048
    0
            anomaly
    1049
    0
        }
    1050
    1051
        /// Get summary statistics.
    1052
    0
        pub fn summary(&self) -> ModelTracerSummary {
    1053
            ModelTracerSummary {
    1054
    0
                total_forwards: self.activation_traces.len(),
    1055
    0
                anomalies_detected: self.activation_traces.iter().filter(|t| t.has_anomaly).count(),
    1056
    0
                attention_traces: self.attention_traces.len(),
    1057
    0
                logit_traces: self.logit_traces.len(),
    1058
    0
                kv_steps: self.kv_trace.steps.len(),
    1059
    0
                total_evictions: self.kv_trace.total_evictions,
    1060
    0
                avg_hit_rate: self.kv_trace.avg_hit_rate,
    1061
    0
                quant_warnings: self.quant_traces.iter().map(|t| t.warning_count()).sum(),
    1062
    0
                quant_criticals: self.quant_traces.iter().map(|t| t.critical_count()).sum(),
    1063
            }
    1064
    0
        }
    1065
    1066
        /// Export summary as JSON for artifact validation.
    1067
    0
        pub fn summary_to_json(&self) -> String {
    1068
    0
            let summary = self.summary();
    1069
    0
            format!(
    1070
    0
                r#"{{"total_forwards":{},"anomalies_detected":{},"attention_traces":{},"logit_traces":{},"kv_steps":{},"total_evictions":{},"avg_hit_rate":{:.4},"quant_warnings":{},"quant_criticals":{}}}"#,
    1071
                summary.total_forwards,
    1072
                summary.anomalies_detected,
    1073
                summary.attention_traces,
    1074
                summary.logit_traces,
    1075
                summary.kv_steps,
    1076
                summary.total_evictions,
    1077
                summary.avg_hit_rate,
    1078
                summary.quant_warnings,
    1079
                summary.quant_criticals
    1080
            )
    1081
    0
        }
    1082
    1083
        /// Clear all accumulated traces (free memory).
    1084
    0
        pub fn clear(&mut self) {
    1085
    0
            self.activation_traces.clear();
    1086
    0
            self.attention_traces.clear();
    1087
    0
            self.logit_traces.clear();
    1088
    0
            self.quant_traces.clear();
    1089
    0
            self.kv_trace = KvCacheSessionTrace::default();
    1090
    0
        }
    1091
    }
    1092
    1093
    /// Summary of model tracer state.
    1094
    #[derive(Debug, Clone, Default)]
    1095
    pub struct ModelTracerSummary {
    1096
        /// Total forward passes traced
    1097
        pub total_forwards: usize,
    1098
        /// Number of forward passes with anomalies
    1099
        pub anomalies_detected: usize,
    1100
        /// Total attention traces collected
    1101
        pub attention_traces: usize,
    1102
        /// Total logit evolution traces
    1103
        pub logit_traces: usize,
    1104
        /// Total KV cache steps traced
    1105
        pub kv_steps: usize,
    1106
        /// Total KV cache evictions
    1107
        pub total_evictions: usize,
    1108
        /// Average KV cache hit rate
    1109
        pub avg_hit_rate: f32,
    1110
        /// Quantization warning count
    1111
        pub quant_warnings: usize,
    1112
        /// Quantization critical count
    1113
        pub quant_criticals: usize,
    1114
    }
    1115
    1116
    impl fmt::Display for ModelTracerSummary {
    1117
    0
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    1118
    0
            writeln!(f, "ModelTracer Summary:")?;
    1119
    0
            writeln!(f, "  Forward passes: {}", self.total_forwards)?;
    1120
    0
            writeln!(f, "  Anomalies: {}", self.anomalies_detected)?;
    1121
    0
            writeln!(f, "  Attention traces: {}", self.attention_traces)?;
    1122
    0
            writeln!(f, "  Logit traces: {}", self.logit_traces)?;
    1123
    0
            writeln!(f, "  KV cache steps: {}", self.kv_steps)?;
    1124
    0
            writeln!(f, "  KV evictions: {}", self.total_evictions)?;
    1125
    0
            writeln!(f, "  Avg hit rate: {:.2}%", self.avg_hit_rate * 100.0)?;
    1126
    0
            writeln!(f, "  Quant warnings: {}", self.quant_warnings)?;
    1127
    0
            write!(f, "  Quant criticals: {}", self.quant_criticals)
    1128
    0
        }
    1129
    }
    1130
    1131
    #[cfg(test)]
    1132
    mod tests {
    1133
        use super::*;
    1134
    1135
        // ========================================================================
    1136
        // QuantType Tests
    1137
        // ========================================================================
    1138
    1139
        #[test]
    1140
        fn test_quant_type_bits() {
    1141
            assert_eq!(QuantType::F32.bits_per_element(), 32.0);
    1142
            assert_eq!(QuantType::F16.bits_per_element(), 16.0);
    1143
            assert_eq!(QuantType::Q8_0.bits_per_element(), 8.0);
    1144
            assert_eq!(QuantType::Q4_K.bits_per_element(), 4.5);
    1145
        }
    1146
    1147
        #[test]
    1148
        fn test_quant_type_compression_ratio() {
    1149
            // F32 -> F32 = 1x
    1150
            assert!((QuantType::F32.compression_ratio() - 1.0).abs() < 0.01);
    1151
            // F32 -> F16 = 2x
    1152
            assert!((QuantType::F16.compression_ratio() - 2.0).abs() < 0.01);
    1153
            // F32 -> Q4_K = ~7.1x
    1154
            assert!(QuantType::Q4_K.compression_ratio() > 7.0);
    1155
        }
    1156
    1157
        // ========================================================================
    1158
        // TensorStats Tests
    1159
        // ========================================================================
    1160
    1161
        #[test]
    1162
        fn test_tensor_stats_empty() {
    1163
            let stats = TensorStats::from_slice(&[]);
    1164
            assert_eq!(stats.count, 0);
    1165
            assert_eq!(stats.nan_count, 0);
    1166
            assert!(!stats.has_anomaly());
    1167
        }
    1168
    1169
        #[test]
    1170
        fn test_tensor_stats_basic() {
    1171
            let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
    1172
            let stats = TensorStats::from_slice(&data);
    1173
            assert_eq!(stats.count, 5);
    1174
            assert_eq!(stats.min, 1.0);
    1175
            assert_eq!(stats.max, 5.0);
    1176
            assert!((stats.mean - 3.0).abs() < 0.01);
    1177
            assert!(!stats.has_anomaly());
    1178
        }
    1179
    1180
        #[test]
    1181
        fn test_tensor_stats_nan_detection() {
    1182
            let data = vec![1.0, f32::NAN, 3.0];
    1183
            let stats = TensorStats::from_slice(&data);
    1184
            assert_eq!(stats.nan_count, 1);
    1185
            assert!(stats.has_anomaly());
    1186
            assert!(stats.anomaly_description().unwrap().contains("NaN"));
    1187
        }
    1188
    1189
        #[test]
    1190
        fn test_tensor_stats_inf_detection() {
    1191
            let data = vec![1.0, f32::INFINITY, 3.0];
    1192
            let stats = TensorStats::from_slice(&data);
    1193
            assert_eq!(stats.inf_count, 1);
    1194
            assert!(stats.has_anomaly());
    1195
        }
    1196
    1197
        #[test]
    1198
        fn test_tensor_stats_explosion() {
    1199
            let data = vec![1e7, 2e7];
    1200
            let stats = TensorStats::from_slice(&data);
    1201
            assert!(stats.has_anomaly());
    1202
            assert!(stats.anomaly_description().unwrap().contains("Explosion"));
    1203
        }
    1204
    1205
        #[test]
    1206
        fn test_tensor_stats_vanishing() {
    1207
            let data = vec![1e-8, 1e-8, 1e-8];
    1208
            let stats = TensorStats::from_slice(&data);
    1209
            assert!(stats.is_vanishing());
    1210
        }
    1211
    1212
        // ========================================================================
    1213
        // LayerActivationTrace Tests
    1214
        // ========================================================================
    1215
    1216
        #[test]
    1217
        fn test_layer_activation_trace_new() {
    1218
            let trace = LayerActivationTrace::new(5);
    1219
            assert_eq!(trace.layer_idx, 5);
    1220
            assert!(!trace.has_anomaly());
    1221
        }
    1222
    1223
        #[test]
    1224
        fn test_layer_activation_trace_anomaly() {
    1225
            let mut trace = LayerActivationTrace::new(0);
    1226
            trace.input_stats = TensorStats::from_slice(&[f32::NAN]);
    1227
            assert!(trace.has_anomaly());
    1228
            assert!(trace.anomaly_description().is_some());
    1229
        }
    1230
    1231
        #[test]
    1232
        fn test_layer_activation_trace_residual_dominance() {
    1233
            let mut trace = LayerActivationTrace::new(0);
    1234
            trace.residual_ratio = 0.999;
    1235
            assert!(trace.has_anomaly());
    1236
            assert!(trace.anomaly_description().unwrap().contains("residual"));
    1237
        }
    1238
    1239
        // ========================================================================
    1240
        // ModelActivationTrace Tests
    1241
        // ========================================================================
    1242
    1243
        #[test]
    1244
        fn test_model_activation_trace_add_layer() {
    1245
            let mut trace = ModelActivationTrace::with_capacity(32);
    1246
            trace.add_layer(LayerActivationTrace::new(0));
    1247
            trace.add_layer(LayerActivationTrace::new(1));
    1248
            assert_eq!(trace.layers.len(), 2);
    1249
            assert!(!trace.has_anomaly);
    1250
        }
    1251
    1252
        #[test]
    1253
        fn test_model_activation_trace_anomaly_propagation() {
    1254
            let mut trace = ModelActivationTrace::default();
    1255
            let mut bad_layer = LayerActivationTrace::new(0);
    1256
            bad_layer.input_stats = TensorStats::from_slice(&[f32::NAN]);
    1257
            trace.add_layer(bad_layer);
    1258
            assert!(trace.has_anomaly);
    1259
        }
    1260
    1261
        // ========================================================================
    1262
        // AttentionWeightTrace Tests
    1263
        // ========================================================================
    1264
    1265
        #[test]
    1266
        fn test_attention_weight_trace_from_weights() {
    1267
            let weights = vec![0.1, 0.3, 0.4, 0.2];
    1268
            let trace = AttentionWeightTrace::from_weights(0, 0, 3, &weights, 2);
    1269
    1270
            assert_eq!(trace.layer_idx, 0);
    1271
            assert_eq!(trace.head_idx, 0);
    1272
            assert_eq!(trace.query_pos, 3);
    1273
            assert_eq!(trace.top_k_positions.len(), 2);
    1274
            // Position 2 has highest weight (0.4), then position 1 (0.3)
    1275
            assert_eq!(trace.top_k_positions[0], 2);
    1276
            assert_eq!(trace.top_k_positions[1], 1);
    1277
        }
    1278
    1279
        #[test]
    1280
        fn test_attention_sink_detection() {
    1281
            let weights = vec![0.8, 0.1, 0.05, 0.05];
    1282
            let trace = AttentionWeightTrace::from_weights(0, 0, 3, &weights, 4);
    1283
            assert!(trace.is_attention_sink(0.5));
    1284
        }
    1285
    1286
        #[test]
    1287
        fn test_recency_bias_detection() {
    1288
            // Position 3 attending mostly to positions 1 and 2
    1289
            let weights = vec![0.05, 0.4, 0.5, 0.05];
    1290
            let trace = AttentionWeightTrace::from_weights(0, 0, 3, &weights, 4);
    1291
            assert!(trace.has_recency_bias(2, 0.5));
    1292
        }
    1293
    1294
        // ========================================================================
    1295
        // TokenLogitEvolution Tests
    1296
        // ========================================================================
    1297
    1298
        #[test]
    1299
        fn test_token_logit_evolution() {
    1300
            let mut evo = TokenLogitEvolution::new(42, "test".to_string());
    1301
            evo.record_layer(1.0, 100);
    1302
            evo.record_layer(2.0, 50);
    1303
            evo.record_layer(3.0, 10);
    1304
    1305
            assert_eq!(evo.per_layer_logit.len(), 3);
    1306
            assert_eq!(evo.per_layer_rank.len(), 3);
    1307
            assert_eq!(evo.decisive_layer(), Some(1)); // 100->50 is biggest jump
    1308
        }
    1309
    1310
        #[test]
    1311
        fn test_logit_evolution_trace_compute_rank() {
    1312
            let logits = vec![1.0, 5.0, 3.0, 2.0]; // sorted: 5, 3, 2, 1
    1313
            // Token 0 has logit 1.0, rank 3 (3 values above it)
    1314
            assert_eq!(LogitEvolutionTrace::compute_rank(&logits, 0), 3);
    1315
            // Token 1 has logit 5.0, rank 0 (nothing above it)
    1316
            assert_eq!(LogitEvolutionTrace::compute_rank(&logits, 1), 0);
    1317
        }
    1318
    1319
        // ========================================================================
    1320
        // QuantizationErrorTrace Tests
    1321
        // ========================================================================
    1322
    1323
        #[test]
    1324
        fn test_quant_error_perfect_match() {
    1325
            let reference = vec![1.0, 2.0, 3.0];
    1326
            let quantized = vec![1.0, 2.0, 3.0];
    1327
            let trace = QuantizationErrorTrace::compute(
    1328
                BrickId::RmsNorm,
    1329
                0,
    1330
                &quantized,
    1331
                &reference,
    1332
                QuantType::Q4_K,
    1333
            );
    1334
    1335
            assert!((trace.mse - 0.0).abs() < 1e-6);
    1336
            assert!((trace.cosine_similarity - 1.0).abs() < 1e-6);
    1337
            assert!(trace.is_acceptable());
    1338
        }
    1339
    1340
        #[test]
    1341
        fn test_quant_error_significant_difference() {
    1342
            let reference = vec![1.0, 2.0, 3.0];
    1343
            // Non-proportional: adds different offsets, changing direction
    1344
            let quantized = vec![1.5, 2.1, 2.9];
    1345
            let trace = QuantizationErrorTrace::compute(
    1346
                BrickId::RmsNorm,
    1347
                0,
    1348
                &quantized,
    1349
                &reference,
    1350
                QuantType::Q4_K,
    1351
            );
    1352
    1353
            assert!(trace.mse > 0.0);
    1354
            assert!(trace.cosine_similarity < 1.0);
    1355
            // Cosine should still be high since vectors are close
    1356
            assert!(trace.cosine_similarity > 0.99);
    1357
        }
    1358
    1359
        // ========================================================================
    1360
        // KvCacheStateTrace Tests
    1361
        // ========================================================================
    1362
    1363
        #[test]
    1364
        fn test_kv_cache_state_trace() {
    1365
            let trace = KvCacheStateTrace::new(0, 2048);
    1366
            assert_eq!(trace.step, 0);
    1367
            assert_eq!(trace.max_positions, 2048);
    1368
            assert!(!trace.is_window_exhausted());
    1369
        }
    1370
    1371
        #[test]
    1372
        fn test_kv_cache_state_utilization() {
    1373
            let mut trace = KvCacheStateTrace::new(0, 1000);
    1374
            trace.valid_positions = 500;
    1375
            assert!((trace.utilization() - 0.5).abs() < 0.01);
    1376
        }
    1377
    1378
        #[test]
    1379
        fn test_kv_cache_session_trace() {
    1380
            let mut session = KvCacheSessionTrace::default();
    1381
            session.add_step(KvCacheStateTrace {
    1382
                step: 0,
    1383
                cache_hit_rate: 0.9,
    1384
                evictions_this_step: 0,
    1385
                cache_size_bytes: 1000,
    1386
                ..Default::default()
    1387
            });
    1388
            session.add_step(KvCacheStateTrace {
    1389
                step: 1,
    1390
                cache_hit_rate: 0.8,
    1391
                evictions_this_step: 5,
    1392
                cache_size_bytes: 2000,
    1393
                ..Default::default()
    1394
            });
    1395
    1396
            assert_eq!(session.steps.len(), 2);
    1397
            assert_eq!(session.total_evictions, 5);
    1398
            assert_eq!(session.peak_memory_bytes, 2000);
    1399
            assert!((session.avg_hit_rate - 0.85).abs() < 0.01);
    1400
        }
    1401
    1402
        // ========================================================================
    1403
        // ModelTracer Tests
    1404
        // ========================================================================
    1405
    1406
        #[test]
    1407
        fn test_model_tracer_lightweight() {
    1408
            let config = ModelTracerConfig::lightweight();
    1409
            assert!(config.trace_activations);
    1410
            assert!(config.trace_kv_cache);
    1411
            assert!(!config.trace_attention);
    1412
            assert!(!config.trace_quant_error);
    1413
        }
    1414
    1415
        #[test]
    1416
        fn test_model_tracer_full() {
    1417
            let config = ModelTracerConfig::full();
    1418
            assert!(config.trace_activations);
    1419
            assert!(config.trace_attention);
    1420
            assert!(config.trace_logits);
    1421
            assert!(config.trace_quant_error);
    1422
            assert!(config.trace_kv_cache);
    1423
        }
    1424
    1425
        #[test]
    1426
        fn test_model_tracer_forward_pass() {
    1427
            let config = ModelTracerConfig::lightweight();
    1428
            let mut tracer = ModelTracer::new(config);
    1429
    1430
            tracer.begin_forward(0);
    1431
            tracer.record_layer_activation(LayerActivationTrace::new(0));
    1432
            tracer.record_layer_activation(LayerActivationTrace::new(1));
    1433
            let anomaly = tracer.end_forward();
    1434
    1435
            assert!(anomaly.is_none());
    1436
            let summary = tracer.summary();
    1437
            assert_eq!(summary.total_forwards, 1);
    1438
            assert_eq!(summary.anomalies_detected, 0);
    1439
        }
    1440
    1441
        #[test]
    1442
        fn test_model_tracer_detects_anomaly() {
    1443
            let config = ModelTracerConfig::lightweight();
    1444
            let mut tracer = ModelTracer::new(config);
    1445
    1446
            tracer.begin_forward(0);
    1447
            let mut bad_layer = LayerActivationTrace::new(0);
    1448
            bad_layer.input_stats = TensorStats::from_slice(&[f32::NAN]);
    1449
            tracer.record_layer_activation(bad_layer);
    1450
            let anomaly = tracer.end_forward();
    1451
    1452
            assert!(anomaly.is_some());
    1453
            assert!(anomaly.unwrap().contains("NaN"));
    1454
            assert_eq!(tracer.summary().anomalies_detected, 1);
    1455
        }
    1456
    1457
        #[test]
    1458
        fn test_model_tracer_json_output() {
    1459
            let config = ModelTracerConfig::lightweight();
    1460
            let mut tracer = ModelTracer::new(config);
    1461
    1462
            tracer.begin_forward(0);
    1463
            tracer.end_forward();
    1464
    1465
            let json = tracer.summary_to_json();
    1466
            assert!(json.contains("\"total_forwards\":1"));
    1467
            assert!(json.contains("\"anomalies_detected\":0"));
    1468
        }
    1469
    1470
        // ========================================================================
    1471
        // Falsification Tests
    1472
        // ========================================================================
    1473
    1474
        /// FALSIFICATION TEST: TensorStats Welford algorithm numerical stability
    1475
        ///
    1476
        /// Welford's algorithm must produce correct mean/std even for large values.
    1477
        #[test]
    1478
        fn test_falsify_tensor_stats_welford_stability() {
    1479
            // Test with large offset - naive algorithm would lose precision
    1480
            let large_offset = 1e9;
    1481
            let data: Vec<f32> = (0..1000).map(|i| large_offset + i as f32).collect();
    1482
            let stats = TensorStats::from_slice(&data);
    1483
    1484
            // Mean should be large_offset + 499.5
    1485
            let expected_mean = large_offset + 499.5;
    1486
            assert!(
    1487
                (stats.mean - expected_mean as f32).abs() < 1.0,
    1488
                "FALSIFICATION FAILED: Welford mean {} != expected {} (relative error too high)",
    1489
                stats.mean,
    1490
                expected_mean
    1491
            );
    1492
    1493
            // Std should be ~288.7 (uniform distribution 0-999)
    1494
            assert!(
    1495
                stats.std > 280.0 && stats.std < 300.0,
    1496
                "FALSIFICATION FAILED: Welford std {} outside expected range [280, 300]",
    1497
                stats.std
    1498
            );
    1499
        }
    1500
    1501
        /// FALSIFICATION TEST: Cosine similarity must be 1.0 for identical vectors
    1502
        #[test]
    1503
        fn test_falsify_cosine_identical_vectors() {
    1504
            let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
    1505
            let trace = QuantizationErrorTrace::compute(
    1506
                BrickId::RmsNorm,
    1507
                0,
    1508
                &data,
    1509
                &data,
    1510
                QuantType::F32,
    1511
            );
    1512
    1513
            assert!(
    1514
                (trace.cosine_similarity - 1.0).abs() < 1e-6,
    1515
                "FALSIFICATION FAILED: identical vectors have cosine {} != 1.0",
    1516
                trace.cosine_similarity
    1517
            );
    1518
        }
    1519
    1520
        /// FALSIFICATION TEST: Cosine similarity must be symmetric
    1521
        #[test]
    1522
        fn test_falsify_cosine_symmetry() {
    1523
            let a = vec![1.0, 2.0, 3.0];
    1524
            let b = vec![4.0, 5.0, 6.0];
    1525
    1526
            let trace_ab = QuantizationErrorTrace::compute(
    1527
                BrickId::RmsNorm, 0, &a, &b, QuantType::F32,
    1528
            );
    1529
            let trace_ba = QuantizationErrorTrace::compute(
    1530
                BrickId::RmsNorm, 0, &b, &a, QuantType::F32,
    1531
            );
    1532
    1533
            assert!(
    1534
                (trace_ab.cosine_similarity - trace_ba.cosine_similarity).abs() < 1e-6,
    1535
                "FALSIFICATION FAILED: cosine(a,b) {} != cosine(b,a) {}",
    1536
                trace_ab.cosine_similarity,
    1537
                trace_ba.cosine_similarity
    1538
            );
    1539
        }
    1540
    1541
        /// FALSIFICATION TEST: ModelTracer layer count must match recorded layers
    1542
        #[test]
    1543
        fn test_falsify_tracer_layer_count() {
    1544
            let config = ModelTracerConfig::lightweight();
    1545
            let mut tracer = ModelTracer::new(config);
    1546
    1547
            tracer.begin_forward(0);
    1548
            let num_layers = 32;
    1549
            for i in 0..num_layers {
    1550
                tracer.record_layer_activation(LayerActivationTrace::new(i));
    1551
            }
    1552
            tracer.end_forward();
    1553
    1554
            // The activation trace should have exactly num_layers entries
    1555
            assert_eq!(
    1556
                tracer.activation_traces[0].layers.len(),
    1557
                num_layers,
    1558
                "FALSIFICATION FAILED: recorded {} layers but expected {}",
    1559
                tracer.activation_traces[0].layers.len(),
    1560
                num_layers
    1561
            );
    1562
        }
    1563
    }