/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 | | } |