## Coverage Report

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

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

    /home/noah/src/realizar/src/inference_trace/mod.rs

    Line
    Count
    Source
    1
    //! Inference Tracing for debugging LLM pipelines (AWS Step Functions Parity)
    2
    //!
    3
    //! Per spec: APR-TRACE-001 v3.0.0
    4
    //! Toyota Way: Genchi Genbutsu (Go and See) + Jidoka (Built-in Quality)
    5
    //!
    6
    //! This module models inference as a deterministic **State Machine**:
    7
    //! 1. TOKENIZE: Text -> Token IDs
    8
    //! 2. EMBED: Token IDs -> Vectors
    9
    //! 3. TRANSFORMER_BLOCK: Vectors -> Vectors (×N layers)
    10
    //! 4. LM_HEAD: Vectors -> Logits
    11
    //! 5. SAMPLE: Logits -> Token ID
    12
    //! 6. DECODE: Token ID -> Text
    13
    //!
    14
    //! Each state transition emits `TaskStateEntered` and `TaskStateExited` events
    15
    //! with verified Input/Output payloads (AWS Step Functions Execution History format).
    16
    //!
    17
    //! Example:
    18
    //! ```bash
    19
    //! apr run model.gguf --prompt "Hello" --trace
    20
    //! apr run model.gguf --prompt "Hi" --trace=tokenize,sample,decode
    21
    //! apr run model.gguf --prompt "Hi" --trace --trace-output trace.json
    22
    //! ```
    23
    24
    use std::collections::HashSet;
    25
    use std::path::PathBuf;
    26
    use std::time::Instant;
    27
    28
    /// Trace configuration
    29
    #[derive(Debug, Clone, Default)]
    30
    pub struct TraceConfig {
    31
        /// Whether tracing is enabled
    32
        pub enabled: bool,
    33
        /// Which steps to trace (empty = all)
    34
        pub steps: HashSet<TraceStep>,
    35
        /// Verbose output (show tensor values)
    36
        pub verbose: bool,
    37
        /// Output file path for JSON trace (None = stderr)
    38
        pub output: Option<PathBuf>,
    39
    }
    40
    41
    impl TraceConfig {
    42
        /// Create a new trace config with tracing enabled
    43
        #[must_use]
    44
    36
        pub fn enabled() -> Self {
    45
    36
            Self {
    46
    36
                enabled: true,
    47
    36
                ..Default::default()
    48
    36
            }
    49
    36
        }
    50
    51
        /// Check if a specific step should be traced
    52
        #[must_use]
    53
    89
        pub fn should_trace(&self, step: TraceStep) -> bool {
    54
    89
            self.enabled && (self.steps86.is_empty86() || self.steps2.contains2(&step2))
    55
    89
        }
    56
    57
        /// Parse trace steps from comma-separated string
    58
        #[must_use]
    59
    1
        pub fn parse_steps(s: &str) -> HashSet<TraceStep> {
    60
    1
            s.split(',')
    61
    3
                .filter_map1(|part| TraceStep::parse(part.trim()))
    62
    1
                .collect()
    63
    1
        }
    64
    }
    65
    66
    /// Inference pipeline steps (State Machine states per AWS Step Functions model)
    67
    #[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
    68
    pub enum TraceStep {
    69
        /// Tokenization (text -> token IDs)
    70
        Tokenize,
    71
        /// Token embedding lookup
    72
        Embed,
    73
        /// Layer normalization
    74
        LayerNorm,
    75
        /// Attention computation
    76
        Attention,
    77
        /// Feed-forward network
    78
        FFN,
    79
        /// Transformer block (combines attention + FFN)
    80
        TransformerBlock,
    81
        /// LM head projection (hidden -> logits)
    82
        LmHead,
    83
        /// Token sampling
    84
        Sample,
    85
        /// Token decoding (token ID -> text)
    86
        Decode,
    87
    }
    88
    89
    impl TraceStep {
    90
        /// Parse step from string
    91
        #[must_use]
    92
    24
        pub fn parse(s: &str) -> Option<Self> {
    93
    24
            match s.to_lowercase().as_str() {
    94
    24
                "tokenize" | "encode"23 => Some(Self::Tokenize)3,
    95
    21
                "embed" | "embedding"20 => Some(Self::Embed)2,
    96
    19
                "layernorm" | "ln"18 | "norm"17 => Some(Self::LayerNorm)3,
    97
    16
                "attention" | "attn"15 => Some(Self::Attention)2,
    98
    14
                "ffn" | "mlp"13 => Some(Self::FFN)2,
    99
    12
                "transformer" | "transformer_block"11 | "layer"11 => Some(Self::TransformerBlock)2,
    100
    10
                "lmhead" | "lm_head"9 | "head"8 => Some(Self::LmHead)3,
    101
    7
                "sample" | "sampling"5 => Some(Self::Sample)3,
    102
    4
                "decode" | "detokenize"2 => Some(Self::Decode)3,
    103
    1
                _ => None,
    104
            }
    105
    24
        }
    106
    107
        /// Get display name for step (AWS Step Functions state name)
    108
        #[must_use]
    109
    29
        pub fn name(&self) -> &'static str {
    110
    29
            match self {
    111
    9
                Self::Tokenize => "TOKENIZE",
    112
    5
                Self::Embed => "EMBED",
    113
    1
                Self::LayerNorm => "LAYER_NORM",
    114
    1
                Self::Attention => "ATTENTION",
    115
    1
                Self::FFN => "FFN",
    116
    2
                Self::TransformerBlock => "TRANSFORMER_BLOCK",
    117
    3
                Self::LmHead => "LM_HEAD",
    118
    2
                Self::Sample => "SAMPLE",
    119
    5
                Self::Decode => "DECODE",
    120
            }
    121
    29
        }
    122
    123
        /// Get legacy name for backwards compatibility (deprecated)
    124
        #[deprecated(since = "3.0.0", note = "Use name() instead")]
    125
        #[must_use]
    126
    6
        pub fn legacy_name(&self) -> &'static str {
    127
    6
            match self {
    128
    1
                Self::Tokenize => "ENCODE",
    129
    1
                Self::Embed => "EMBED",
    130
    0
                Self::LayerNorm => "LAYER_NORM",
    131
    0
                Self::Attention => "ATTENTION",
    132
    0
                Self::FFN => "FFN",
    133
    1
                Self::TransformerBlock => "TRANSFORMER",
    134
    1
                Self::LmHead => "LM_HEAD",
    135
    1
                Self::Sample => "SAMPLE",
    136
    1
                Self::Decode => "DECODE",
    137
            }
    138
    6
        }
    139
    140
        /// Get step number for 7-step pipeline
    141
        #[must_use]
    142
    13
        pub fn step_number(&self) -> usize {
    143
    13
            match self {
    144
    3
                Self::Tokenize => 1,
    145
    3
                Self::Embed => 2,
    146
    1
                Self::LayerNorm | Self::Attention | Self::FFN | Self::TransformerBlock => 3,
    147
    2
                Self::LmHead => 4,
    148
    1
                Self::Sample => 5,
    149
    3
                Self::Decode => 6,
    150
            }
    151
    13
        }
    152
    }
    153
    154
    /// Tensor statistics for tracing
    155
    #[derive(Debug, Clone, Default)]
    156
    pub struct TensorStats {
    157
        /// Minimum value
    158
        pub min: f32,
    159
        /// Maximum value
    160
        pub max: f32,
    161
        /// Mean value
    162
        pub mean: f32,
    163
        /// Standard deviation
    164
        pub std: f32,
    165
        /// Whether NaN values were detected
    166
        pub has_nan: bool,
    167
        /// Whether Inf values were detected
    168
        pub has_inf: bool,
    169
    }
    170
    171
    impl TensorStats {
    172
        /// Compute stats from tensor data
    173
        #[must_use]
    174
    33
        pub fn from_slice(data: &[f32]) -> Self {
    175
    33
            if data.is_empty() {
    176
    1
                return Self::default();
    177
    32
            }
    178
    179
    32
            let mut min = f32::INFINITY;
    180
    32
            let mut max = f32::NEG_INFINITY;
    181
    32
            let mut sum = 0.0f64;
    182
    32
            let mut has_nan = false;
    183
    32
            let mut has_inf = false;
    184
    185
    155
            for &v123 in data {
    186
    123
                if v.is_nan() {
    187
    6
                    has_nan = true;
    188
    117
                } else if v.is_infinite() {
    189
    7
                    has_inf = true;
    190
    110
                } else {
    191
    110
                    min = min.min(v);
    192
    110
                    max = max.max(v);
    193
    110
                    sum += f64::from(v);
    194
    110
                }
    195
            }
    196
    197
    32
            let mean = (sum / data.len() as f64) as f32;
    198
    199
            // Compute std dev
    200
    32
            let mut var_sum = 0.0f64;
    201
    155
            for &v123 in data {
    202
    123
                if !v.is_nan() && !v.is_infinite()117 {
    203
    110
                    let diff = f64::from(v) - f64::from(mean);
    204
    110
                    var_sum += diff * diff;
    205
    110
                }13
    206
            }
    207
    32
            let std = ((var_sum / data.len() as f64).sqrt()) as f32;
    208
    209
    32
            Self {
    210
    32
                min,
    211
    32
                max,
    212
    32
                mean,
    213
    32
                std,
    214
    32
                has_nan,
    215
    32
                has_inf,
    216
    32
            }
    217
    33
        }
    218
    219
        /// Check if stats indicate an error (Jidoka)
    220
        #[must_use]
    221
    3
        pub fn has_error(&self) -> bool {
    222
    3
            self.has_nan || self.has_inf2
    223
    3
        }
    224
    }
    225
    226
    /// Trace error types (Jidoka: stop-the-line errors)
    227
    #[derive(Debug, Clone)]
    228
    pub enum TraceError {
    229
        /// Token ID exceeds vocabulary size
    230
        VocabOverflow {
    231
            /// The offending token ID
    232
            token_id: u32,
    233
            /// Size of the vocabulary
    234
            vocab_size: usize,
    235
        },
    236
        /// NaN values detected in tensor
    237
        NaNDetected {
    238
            /// Layer index where NaN was detected (None if embedding)
    239
            layer: Option<usize>,
    240
        },
    241
        /// Inf values detected in tensor
    242
        InfDetected {
    243
            /// Layer index where Inf was detected (None if embedding)
    244
            layer: Option<usize>,
    245
        },
    246
        /// Garbage characters in decoded output (APR-TOK-001)
    247
        GarbageOutput {
    248
            /// Sample of garbage output
    249
            sample: String,
    250
        },
    251
        /// Unknown token (OOV)
    252
        UnknownToken {
    253
            /// The unknown token ID
    254
            token_id: u32,
    255
        },
    256
        /// Shape mismatch
    257
        ShapeMismatch {
    258
            /// Expected shape
    259
            expected: Vec<usize>,
    260
            /// Actual shape
    261
            actual: Vec<usize>,
    262
        },
    263
        /// Execution failed (F-JID-01: Jidoka)
    264
        ExecutionFailed {
    265
            /// Cause of failure
    266
            cause: String,
    267
        },
    268
    }
    269
    270
    impl std::fmt::Display for TraceError {
    271
    12
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    272
    12
            match self {
    273
                Self::VocabOverflow {
    274
    4
                    token_id,
    275
    4
                    vocab_size,
    276
                } => {
    277
    4
                    write!(f, "Token ID {} exceeds vocab size {}", token_id, vocab_size)
    278
                },
    279
    3
                Self::NaNDetected { layer } => {
    280
    3
                    if let Some(l1) = layer {
    281
    1
                        write!(f, "NaN values detected in layer {}", l)
    282
                    } else {
    283
    2
                        write!(f, "NaN values detected")
    284
                    }
    285
                },
    286
    2
                Self::InfDetected { layer } => {
    287
    2
                    if let Some(l1) = layer {
    288
    1
                        write!(f, "Inf values detected in layer {}", l)
    289
                    } else {
    290
    1
                        write!(f, "Inf values detected")
    291
                    }
    292
                },
    293
    1
                Self::GarbageOutput { sample } => {
    294
    1
                    write!(f, "Garbage output detected: {:?}", sample)
    295
                },
    296
    1
                Self::UnknownToken { token_id } => {
    297
    1
                    write!(f, "Unknown token ID: {}", token_id)
    298
                },
    299
    1
                Self::ShapeMismatch { expected, actual } => {
    300
    1
                    write!(
    301
    1
                        f,
    302
    1
                        "Shape mismatch: expected {:?}, got {:?}",
    303
                        expected, actual
    304
                    )
    305
                },
    306
    0
                Self::ExecutionFailed { cause } => {
    307
    0
                    write!(f, "Execution failed: {}", cause)
    308
                },
    309
            }
    310
    12
        }
    311
    }
    312
    313
    /// AWS Step Functions event type (per spec v3.1.0)
    314
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    315
    pub enum AwsEventType {
    316
        /// State machine entered a state
    317
        TaskStateEntered,
    318
        /// State machine exited a state
    319
        TaskStateExited,
    320
        /// Execution failed with error
    321
        ExecutionFailed,
    322
    }
    323
    324
    impl AwsEventType {
    325
        /// Get the event type name (AWS Step Functions format)
    326
        #[must_use]
    327
    13
        pub fn name(&self) -> &'static str {
    328
    13
            match self {
    329
    6
                Self::TaskStateEntered => "TaskStateEntered",
    330
    6
                Self::TaskStateExited => "TaskStateExited",
    331
    1
                Self::ExecutionFailed => "ExecutionFailed",
    332
            }
    333
    13
        }
    334
    }
    335
    336
    /// Trace event emitted during inference (AWS Step Functions Parity)
    337
    #[derive(Debug, Clone)]
    338
    pub struct TraceEvent {
    339
        /// Unique event ID (AWS Step Functions: monotonically increasing)
    340
        pub id: u64,
    341
        /// ISO 8601 timestamp
    342
        pub timestamp: String,
    343
        /// AWS Step Functions event type
    344
        pub event_type: AwsEventType,
    345
        /// Link to the entry event (for TaskStateExited)
    346
        pub previous_event_id: Option<u64>,
    347
        /// Pipeline step (state name)
    348
        pub step: TraceStep,
    349
        /// Generation iteration (0 for prefill)
    350
        pub iteration: usize,
    351
        /// Layer index (for transformer layers)
    352
        pub layer: Option<usize>,
    353
        /// Input shape
    354
        pub input_shape: Vec<usize>,
    355
        /// Output shape
    356
        pub output_shape: Vec<usize>,
    357
        /// Tensor statistics
    358
        pub stats: TensorStats,
    359
        /// Duration in microseconds
    360
        pub duration_us: u64,
    361
        /// Error if any (Jidoka)
    362
        pub error: Option<TraceError>,
    363
        /// Cause of failure (F-AWS-05: required for ExecutionFailed events)
    364
        pub cause: Option<String>,
    365
        /// Additional details (step-specific)
    366
        pub details: TraceDetails,
    367
    }
    368
    369
    /// Step-specific trace details
    370
    #[derive(Debug, Clone, Default)]
    371
    pub struct TraceDetails {
    372
        /// Input text (for encode step)
    373
        pub input_text: Option<String>,
    374
        /// Output tokens (for encode step)
    375
        pub output_tokens: Option<Vec<u32>>,
    376
        /// Vocabulary entries (for encode step, OOV detection)
    377
        pub vocab_entries: Option<Vec<String>>,
    378
        /// Top-k logits with token IDs (for lm_head/sample step)
    379
        pub top_k_logits: Option<Vec<(u32, f32)>>,
    380
        /// Top-k probabilities with token IDs (for sample step)
    381
        pub top_k_probs: Option<Vec<(u32, f32)>>,
    382
        /// Sampled token ID (for sample/decode step)
    383
        pub sampled_token: Option<u32>,
    384
        /// Decoded text output (for decode step)
    385
        pub decoded_text: Option<String>,
    386
        /// Token string representation (for decode step)
    387
        pub token_string: Option<String>,
    388
        /// Temperature parameter used (for sample step)
    389
        pub temperature: Option<f32>,
    390
        /// Top-k parameter used (for sample step)
    391
        pub top_k: Option<usize>,
    392
    }
    393
    394
    /// Inference tracer
    395
    #[derive(Debug)]
    396
    pub struct InferenceTracer {
    397
        /// Configuration
    398
        config: TraceConfig,
    399
        /// Collected events
    400
        events: Vec<TraceEvent>,
    401
        /// Model info
    402
        model_info: ModelInfo,
    403
        /// Current step timer
    404
        step_start: Option<Instant>,
    405
        /// Total errors count
    406
        error_count: usize,
    407
        /// Total warnings count
    408
        warning_count: usize,
    409
        /// Next event ID (monotonically increasing per AWS Step Functions)
    410
        next_event_id: u64,
    411
        /// ID of the last TaskStateEntered event (for linking TaskStateExited)
    412
        last_entered_id: Option<u64>,
    413
    }
    414
    415
    /// Model information for trace header
    416
    #[derive(Debug, Clone, Default)]
    417
    pub struct ModelInfo {
    418
        /// Model name/path
    419
        pub name: String,
    420
        /// Number of layers
    421
        pub num_layers: usize,
    422
        /// Hidden dimension
    423
        pub hidden_dim: usize,
    424
        /// Vocabulary size
    425
        pub vocab_size: usize,
    426
        /// Number of attention heads
    427
        pub num_heads: usize,
    428
        /// Quantization type (e.g., "Q4_K_M")
    429
        pub quant_type: Option<String>,
    430
    }
    431
    432
    impl InferenceTracer {
    433
        /// Create a new tracer with config
    434
        #[must_use]
    435
    37
        pub fn new(config: TraceConfig) -> Self {
    436
    37
            Self {
    437
    37
                config,
    438
    37
                events: Vec::new(),
    439
    37
                model_info: ModelInfo::default(),
    440
    37
                step_start: None,
    441
    37
                error_count: 0,
    442
    37
                warning_count: 0,
    443
    37
                next_event_id: 1, // AWS Step Functions IDs start at 1
    444
    37
                last_entered_id: None,
    445
    37
            }
    446
    37
        }
    447
    448
        /// Create a disabled tracer (no-op)
    449
        #[must_use]
    450
    1
        pub fn disabled() -> Self {
    451
    1
            Self::new(TraceConfig::default())
    452
    1
        }
    453
    454
        /// Set model info
    455
    9
        pub fn set_model_info(&mut self, info: ModelInfo) {
    456
    9
            self.model_info = info;
    457
    9
        }
    458
    459
        /// Check if tracing is enabled
    460
        #[must_use]
    461
    1
        pub fn is_enabled(&self) -> bool {
    462
    1
            self.config.enabled
    463
    1
        }
    464
    465
        /// Check if verbose tracing is enabled (requires D2H sync for stats)
    466
        #[must_use]
    467
    2
        pub fn is_verbose(&self) -> bool {
    468
    2
            self.config.enabled && self.config.verbose
    469
    2
        }
    470
    471
        /// Get next event ID and increment (AWS Step Functions: monotonically increasing)
    472
    77
        fn next_id(&mut self) -> u64 {
    473
    77
            let id = self.next_event_id;
    474
    77
            self.next_event_id += 1;
    475
    77
            id
    476
    77
        }
    477
    478
        /// Generate ISO 8601 timestamp
    479
    77
        fn timestamp() -> String {
    480
    77
            chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
    481
    77
        }
    482
    483
        /// Start timing a step and emit TaskStateEntered event (AWS Step Functions F-AWS-01)
    484
    39
        pub fn start_step(&mut self, step: TraceStep) {
    485
    39
            if self.config.should_trace(step) {
    486
    38
                self.step_start = Some(Instant::now());
    487
    38
    488
    38
                // Emit TaskStateEntered event (F-AWS-01: Entry/Exit pairing)
    489
    38
                let entry_id = self.next_id();
    490
    38
                let event = TraceEvent {
    491
    38
                    id: entry_id,
    492
    38
                    timestamp: Self::timestamp(),
    493
    38
                    event_type: AwsEventType::TaskStateEntered,
    494
    38
                    previous_event_id: None, // Entry events have no predecessor
    495
    38
                    step,
    496
    38
                    iteration: 0,
    497
    38
                    layer: None,
    498
    38
                    input_shape: vec![],
    499
    38
                    output_shape: vec![],
    500
    38
                    stats: TensorStats::default(),
    501
    38
                    duration_us: 0,
    502
    38
                    error: None,
    503
    38
                    cause: None,
    504
    38
                    details: TraceDetails::default(),
    505
    38
                };
    506
    38
                self.events.push(event);
    507
    38
                // Store entry ID for the corresponding Exit event (F-AWS-02)
    508
    38
                self.last_entered_id = Some(entry_id);
    509
    38
            }1
    510
    39
        }
    511
    512
        /// Trace encode step (tokenization)
    513
    9
        pub fn trace_encode(&mut self, input_text: &str, output_tokens: &[u32], vocab_size: usize) {
    514
    9
            if !self.config.should_trace(TraceStep::Tokenize) {
    515
    1
                return;
    516
    8
            }
    517
    518
    8
            let duration = self
    519
    8
                .step_start
    520
    8
                .map_or(0, |s| s.elapsed().as_micros() as u64);
    521
    522
            // Check for OOV tokens (Jidoka)
    523
    8
            let mut error = None;
    524
    32
            for &token_id24 in output_tokens {
    525
    24
                if token_id as usize >= vocab_size {
    526
    0
                    error = Some(TraceError::VocabOverflow {
    527
    0
                        token_id,
    528
    0
                        vocab_size,
    529
    0
                    });
    530
    0
                    self.error_count += 1;
    531
    0
                    break;
    532
    24
                }
    533
            }
    534
    535
    8
            let event = TraceEvent {
    536
    8
                id: self.next_id(),
    537
    8
                timestamp: Self::timestamp(),
    538
    8
                event_type: AwsEventType::TaskStateExited,
    539
    8
                previous_event_id: self.last_entered_id.take(),
    540
    8
                step: TraceStep::Tokenize,
    541
    8
                iteration: 0,
    542
    8
                layer: None,
    543
    8
                input_shape: vec![input_text.len()],
    544
    8
                output_shape: vec![output_tokens.len()],
    545
    8
                stats: TensorStats::default(),
    546
    8
                duration_us: duration,
    547
    8
                error,
    548
    8
                cause: None,
    549
    8
                details: TraceDetails {
    550
    8
                    input_text: Some(input_text.to_string()),
    551
    8
                    output_tokens: Some(output_tokens.to_vec()),
    552
    8
                    ..Default::default()
    553
    8
                },
    554
    8
            };
    555
    556
    8
            self.events.push(event);
    557
    8
            self.step_start = None;
    558
    9
        }
    559
    560
        /// Trace embed step
    561
    7
        pub fn trace_embed(
    562
    7
            &mut self,
    563
    7
            token_count: usize,
    564
    7
            hidden_dim: usize,
    565
    7
            embeddings: Option<&[f32]>,
    566
    7
        ) {
    567
    7
            if !self.config.should_trace(TraceStep::Embed) {
    568
    0
                return;
    569
    7
            }
    570
    571
    7
            let duration = self
    572
    7
                .step_start
    573
    7
                .map_or(0, |s| s.elapsed().as_micros() as u64);
    574
    7
            let stats = embeddings.map(TensorStats::from_slice).unwrap_or_default();
    575
    576
    7
            let mut error = None;
    577
    7
            if stats.has_nan {
    578
    2
                error = Some(TraceError::NaNDetected { layer: None });
    579
    2
                self.error_count += 1;
    580
    5
            } else if stats.has_inf {
    581
    1
                error = Some(TraceError::InfDetected { layer: None });
    582
    1
                self.error_count += 1;
    583
    4
            }
    584
    585
    7
            let event = TraceEvent {
    586
    7
                id: self.next_id(),
    587
    7
                timestamp: Self::timestamp(),
    588
    7
                event_type: AwsEventType::TaskStateExited,
    589
    7
                previous_event_id: self.last_entered_id.take(),
    590
    7
                step: TraceStep::Embed,
    591
    7
                iteration: 0,
    592
    7
                layer: None,
    593
    7
                input_shape: vec![token_count],
    594
    7
                output_shape: vec![token_count, hidden_dim],
    595
    7
                stats,
    596
    7
                duration_us: duration,
    597
    7
                error,
    598
    7
                cause: None,
    599
    7
                details: TraceDetails::default(),
    600
    7
            };
    601
    602
    7
            self.events.push(event);
    603
    7
            self.step_start = None;
    604
    7
        }
    605
    606
        /// Trace transformer layer
    607
    8
        pub fn trace_layer(
    608
    8
            &mut self,
    609
    8
            layer_idx: usize,
    610
    8
            iteration: usize,
    611
    8
            hidden_state: Option<&[f32]>,
    612
    8
            seq_len: usize,
    613
    8
            hidden_dim: usize,
    614
    8
        ) {
    615
    8
            if !self.config.should_trace(TraceStep::TransformerBlock) {
    616
    0
                return;
    617
    8
            }
    618
    619
    8
            let duration = self
    620
    8
                .step_start
    621
    8
                .map_or(0, |s| s.elapsed().as_micros() as u64);
    622
    8
            let stats = hidden_state
    623
    8
                .map(TensorStats::from_slice)
    624
    8
                .unwrap_or_default();
    625
    626
    8
            let mut error = None;
    627
    8
            if stats.has_nan {
    628
    1
                error = Some(TraceError::NaNDetected {
    629
    1
                    layer: Some(layer_idx),
    630
    1
                });
    631
    1
                self.error_count += 1;
    632
    7
            } else if stats.has_inf {
    633
    1
                error = Some(TraceError::InfDetected {
    634
    1
                    layer: Some(layer_idx),
    635
    1
                });
    636
    1
                self.error_count += 1;
    637
    6
            }
    638
    639
    8
            let event = TraceEvent {
    640
    8
                id: self.next_id(),
    641
    8
                timestamp: Self::timestamp(),
    642
    8
                event_type: AwsEventType::TaskStateExited,
    643
    8
                previous_event_id: self.last_entered_id.take(),
    644
    8
                step: TraceStep::TransformerBlock,
    645
    8
                iteration,
    646
    8
                layer: Some(layer_idx),
    647
    8
                input_shape: vec![seq_len, hidden_dim],
    648
    8
                output_shape: vec![seq_len, hidden_dim],
    649
    8
                stats,
    650
    8
                duration_us: duration,
    651
    8
                error,
    652
    8
                cause: None,
    653
    8
                details: TraceDetails::default(),
    654
    8
            };
    655
    656
    8
            self.events.push(event);
    657
    8
            self.step_start = None;
    658
    8
        }
    659
    660
        /// Trace LM head projection
    661
    6
        pub fn trace_lm_head(&mut self, iteration: usize, logits: &[f32], vocab_size: usize) {
    662
    6
            if !self.config.should_trace(TraceStep::LmHead) {
    663
    0
                return;
    664
    6
            }
    665
    666
    6
            let duration = self
    667
    6
                .step_start
    668
    6
                .map_or(0, |s| s.elapsed().as_micros() as u64);
    669
    6
            let stats = TensorStats::from_slice(logits);
    670
    671
            // Get top-5 logits
    672
    6
            let top_k = get_top_k_indices(logits, 5);
    673
    674
    6
            let mut error = None;
    675
    6
            if stats.has_nan {
    676
    1
                error = Some(TraceError::NaNDetected { layer: None });
    677
    1
                self.error_count += 1;
    678
    5
            } else if stats.has_inf {
    679
    1
                error = Some(TraceError::InfDetected { layer: None });
    680
    1
                self.error_count += 1;
    681
    4
            }
    682
    683
    6
            let event = TraceEvent {
    684
    6
                id: self.next_id(),
    685
    6
                timestamp: Self::timestamp(),
    686
    6
                event_type: AwsEventType::TaskStateExited,
    687
    6
                previous_event_id: self.last_entered_id.take(),
    688
    6
                step: TraceStep::LmHead,
    689
    6
                iteration,
    690
    6
                layer: None,
    691
    6
                input_shape: vec![self.model_info.hidden_dim],
    692
    6
                output_shape: vec![vocab_size],
    693
    6
                stats,
    694
    6
                duration_us: duration,
    695
    6
                error,
    696
    6
                cause: None,
    697
    6
                details: TraceDetails {
    698
    6
                    top_k_logits: Some(top_k),
    699
    6
                    ..Default::default()
    700
    6
                },
    701
    6
            };
    702
    703
    6
            self.events.push(event);
    704
    6
            self.step_start = None;
    705
    6
        }
    706
    707
        /// Trace sampling step
    708
    3
        pub fn trace_sample(
    709
    3
            &mut self,
    710
    3
            iteration: usize,
    711
    3
            logits: &[f32],
    712
    3
            sampled_token: u32,
    713
    3
            temperature: f32,
    714
    3
            top_k: usize,
    715
    3
        ) {
    716
    3
            if !self.config.should_trace(TraceStep::Sample) {
    717
    0
                return;
    718
    3
            }
    719
    720
    3
            let duration = self
    721
    3
                .step_start
    722
    3
                .map_or(0, |s| s.elapsed().as_micros() as u64);
    723
    724
            // Compute softmax probabilities for top-k display
    725
    3
            let top_k_logits = get_top_k_indices(logits, top_k.min(10));
    726
    3
            let top_k_probs = compute_top_k_probs(logits, &top_k_logits);
    727
    728
    3
            let event = TraceEvent {
    729
    3
                id: self.next_id(),
    730
    3
                timestamp: Self::timestamp(),
    731
    3
                event_type: AwsEventType::TaskStateExited,
    732
    3
                previous_event_id: self.last_entered_id.take(),
    733
    3
                step: TraceStep::Sample,
    734
    3
                iteration,
    735
    3
                layer: None,
    736
    3
                input_shape: vec![logits.len()],
    737
    3
                output_shape: vec![1],
    738
    3
                stats: TensorStats::from_slice(logits),
    739
    3
                duration_us: duration,
    740
    3
                error: None,
    741
    3
                cause: None,
    742
    3
                details: TraceDetails {
    743
    3
                    top_k_logits: Some(top_k_logits),
    744
    3
                    top_k_probs: Some(top_k_probs),
    745
    3
                    sampled_token: Some(sampled_token),
    746
    3
                    temperature: Some(temperature),
    747
    3
                    top_k: Some(top_k),
    748
    3
                    ..Default::default()
    749
    3
                },
    750
    3
            };
    751
    752
    3
            self.events.push(event);
    753
    3
            self.step_start = None;
    754
    3
        }
    755
    756
        /// Trace decode step
    757
    6
        pub fn trace_decode(
    758
    6
            &mut self,
    759
    6
            iteration: usize,
    760
    6
            token_id: u32,
    761
    6
            decoded_text: &str,
    762
    6
            vocab_size: usize,
    763
    6
        ) {
    764
    6
            if !self.config.should_trace(TraceStep::Decode) {
    765
    0
                return;
    766
    6
            }
    767
    768
    6
            let duration = self
    769
    6
                .step_start
    770
    6
                .map_or(0, |s| s.elapsed().as_micros() as u64);
    771
    772
            // Check for garbage output (APR-TOK-001 Jidoka)
    773
    6
            let mut error = None;
    774
    6
            if token_id as usize >= vocab_size {
    775
    3
                error = Some(TraceError::VocabOverflow {
    776
    3
                    token_id,
    777
    3
                    vocab_size,
    778
    3
                });
    779
    3
                self.error_count += 1;
    780
    3
            } else if is_garbage_output(decoded_text) {
    781
    1
                error = Some(TraceError::GarbageOutput {
    782
    1
                    sample: decoded_text.chars().take(20).collect(),
    783
    1
                });
    784
    1
                self.error_count += 1;
    785
    2
            }
    786
    787
    6
            let event = TraceEvent {
    788
    6
                id: self.next_id(),
    789
    6
                timestamp: Self::timestamp(),
    790
    6
                event_type: AwsEventType::TaskStateExited,
    791
    6
                previous_event_id: self.last_entered_id.take(),
    792
    6
                step: TraceStep::Decode,
    793
    6
                iteration,
    794
    6
                layer: None,
    795
    6
                input_shape: vec![1],
    796
    6
                output_shape: vec![decoded_text.len()],
    797
    6
                stats: TensorStats::default(),
    798
    6
                duration_us: duration,
    799
    6
                error,
    800
    6
                cause: None,
    801
    6
                details: TraceDetails {
    802
    6
                    sampled_token: Some(token_id),
    803
    6
                    decoded_text: Some(decoded_text.to_string()),
    804
    6
                    ..Default::default()
    805
    6
                },
    806
    6
            };
    807
    808
    6
            self.events.push(event);
    809
    6
            self.step_start = None;
    810
    6
        }
    811
    812
        /// Get all collected events
    813
        #[must_use]
    814
    17
        pub fn events(&self) -> &[TraceEvent] {
    815
    17
            &self.events
    816
    17
        }
    817
    818
        /// Get error count
    819
        #[must_use]
    820
    10
        pub fn error_count(&self) -> usize {
    821
    10
            self.error_count
    822
    10
        }
    823
    824
        /// Record an execution failure (F-JID-01: Jidoka error handling)
    825
        ///
    826
        /// Emits an `ExecutionFailed` event per AWS Step Functions parity (F-AWS-05).
    827
        /// Use this when the inference cannot proceed due to missing config,
    828
        /// invalid model format, or other fatal errors.
    829
        ///
    830
        /// # Arguments
    831
        /// * `error` - High-level error category (e.g., "Initialization Failure")
    832
        /// * `cause` - Specific cause of failure (e.g., "Missing config.json")
    833
    1
        pub fn record_execution_failed(&mut self, error: &str, cause: &str) {
    834
    1
            if !self.config.enabled {
    835
    0
                return;
    836
    1
            }
    837
    838
    1
            let event = TraceEvent {
    839
    1
                id: self.next_id(),
    840
    1
                timestamp: Self::timestamp(),
    841
    1
                event_type: AwsEventType::ExecutionFailed,
    842
    1
                previous_event_id: None,
    843
    1
                step: TraceStep::Tokenize, // Use first step as placeholder
    844
    1
                iteration: 0,
    845
    1
                layer: None,
    846
    1
                input_shape: vec![],
    847
    1
                output_shape: vec![],
    848
    1
                stats: TensorStats::default(),
    849
    1
                duration_us: 0,
    850
    1
                error: Some(TraceError::ExecutionFailed {
    851
    1
                    cause: error.to_string(),
    852
    1
                }),
    853
    1
                cause: Some(cause.to_string()),
    854
    1
                details: TraceDetails::default(),
    855
    1
            };
    856
    1
            self.events.push(event);
    857
    1
            self.error_count += 1;
    858
    1
        }
    859
    860
        /// Format trace output as text
    861
        #[must_use]
    862
    9
        pub fn format_text(&self) -> String {
    863
    9
            let mut output = String::new();
    864
    865
            // Header
    866
    9
            output.push_str("=== APR Inference Trace ===\n");
    867
    9
            if !self.model_info.name.is_empty() {
    868
    4
                output.push_str(&format!(
    869
    4
                    "Model: {} ({} layers, hidden={})\n",
    870
    4
                    self.model_info.name, self.model_info.num_layers, self.model_info.hidden_dim
    871
    4
                ));
    872
    5
            }
    873
    9
            output.push('\n');
    874
    875
            // Group events by step type for cleaner output
    876
    9
            let mut current_step = None;
    877
    9
            let mut layer_count = 0;
    878
    879
    37
            for event28 in &self.events {
    880
                // Step header
    881
    28
                if current_step != Some(event.step) {
    882
    10
                    if current_step == Some(TraceStep::TransformerBlock) && layer_count > 01 {
    883
    1
                        output.push_str(&format!("  ... ({} layers total)\n", layer_count));
    884
    9
                    }
    885
    10
                    current_step = Some(event.step);
    886
    10
                    layer_count = 0;
    887
    888
    10
                    output.push_str(&format!(
    889
    10
                        "[{}/7] {}\n",
    890
    10
                        event.step.step_number(),
    891
    10
                        event.step.name()
    892
    10
                    ));
    893
    18
                }
    894
    895
                // Step content
    896
    28
                match event.step {
    897
                    TraceStep::Tokenize => {
    898
    4
                        if let Some(ref text2) = event.details.input_text {
    899
    2
                            let display_text = if text.len() > 50 {
    900
    1
                                format!("{}...", &text[..50])
    901
                            } else {
    902
    1
                                text.clone()
    903
                            };
    904
    2
                            output.push_str(&format!("  Input:  {:?}\n", display_text));
    905
    2
                        }
    906
    4
                        if let Some(ref tokens2) = event.details.output_tokens {
    907
    2
                            let display_tokens: Vec<_> = tokens.iter().take(10).collect();
    908
    2
                            if tokens.len() > 10 {
    909
    1
                                output.push_str(&format!(
    910
    1
                                    "  Output: {:?}...  ({} tokens)\n",
    911
    1
                                    display_tokens,
    912
    1
                                    tokens.len()
    913
    1
                                ));
    914
    1
                            } else {
    915
    1
                                output.push_str(&format!(
    916
    1
                                    "  Output: {:?}  ({} tokens)\n",
    917
    1
                                    display_tokens,
    918
    1
                                    tokens.len()
    919
    1
                                ));
    920
    1
                            }
    921
    2
                        }
    922
                    },
    923
    4
                    TraceStep::Embed => {
    924
    4
                        output.push_str(&format!(
    925
    4
                            "  Input:  [{} token IDs]\n",
    926
    4
                            event.input_shape.first().unwrap_or(&0)
    927
    4
                        ));
    928
    4
                        output.push_str(&format!("  Output: {:?} float32\n", event.output_shape));
    929
    4
                        output.push_str(&format!(
    930
    4
                            "  Range:  min={:.2}, max={:.2}, mean={:.3}\n",
    931
    4
                            event.stats.min, event.stats.max, event.stats.mean
    932
    4
                        ));
    933
    4
                    },
    934
                    TraceStep::TransformerBlock => {
    935
    10
                        layer_count += 1;
    936
    10
                        if layer_count <= 3 || self.config.verbose7 {
    937
    3
                            output.push_str(&format!(
    938
    3
                                "  Layer {:2}: attn {} ffn {}  {:?} range=[{:.1}, {:.1}]\n",
    939
    3
                                event.layer.unwrap_or(0),
    940
    3
                                if event.error.is_none() { "OK" } else { "ERR"0 },
    941
    3
                                if event.error.is_none() { "OK" } else { "ERR"0 },
    942
                                event.output_shape,
    943
                                event.stats.min,
    944
                                event.stats.max
    945
                            ));
    946
    7
                        }
    947
                    },
    948
                    TraceStep::LmHead => {
    949
    4
                        output.push_str(&format!(
    950
    4
                            "  Input:  [{}] (last token hidden state)\n",
    951
    4
                            event.input_shape.first().unwrap_or(&0)
    952
    4
                        ));
    953
    4
                        output.push_str(&format!(
    954
    4
                            "  Output: [{}] logits\n",
    955
    4
                            event.output_shape.first().unwrap_or(&0)
    956
    4
                        ));
    957
    4
                        if let Some(ref top_k2) = event.details.top_k_logits {
    958
    2
                            output.push_str("  Top 5:  ");
    959
    9
                            for (i, (tok, logit)) in top_k.iter()2.take2(5).enumerate2() {
    960
    9
                                if i > 0 {
    961
    7
                                    output.push_str(", ");
    962
    7
                                }2
    963
    9
                                output.push_str(&format!("{}={:.2}", tok, logit));
    964
                            }
    965
    2
                            output.push('\n');
    966
    2
                        }
    967
                    },
    968
                    TraceStep::Sample => {
    969
    2
                        output.push_str(&format!(
    970
    2
                            "  Logits:  [{}] -> scaled -> filtered\n",
    971
    2
                            event.input_shape.first().unwrap_or(&0)
    972
    2
                        ));
    973
    2
                        if let Some(ref probs1) = event.details.top_k_probs {
    974
    1
                            output.push_str("  Probs:   ");
    975
    3
                            for (i, (tok, prob)) in probs.iter()1.take1(5).enumerate1() {
    976
    3
                                if i > 0 {
    977
    2
                                    output.push_str(", ");
    978
    2
                                }1
    979
    3
                                output.push_str(&format!("{}={:.2}", tok, prob));
    980
                            }
    981
    1
                            output.push('\n');
    982
    1
                        }
    983
    2
                        if let Some(token1) = event.details.sampled_token {
    984
    1
                            output.push_str(&format!("  Sampled: token_id={}\n", token));
    985
    1
                        }
    986
                    },
    987
                    TraceStep::Decode => {
    988
    4
                        if let Some(token2) = event.details.sampled_token {
    989
    2
                            output.push_str(&format!("  Token ID:  {}\n", token));
    990
    2
                        }
    991
    4
                        if let Some(ref text2) = event.details.decoded_text {
    992
    2
                            output.push_str(&format!("  Decoded:   {:?}\n", text));
    993
    2
                        }
    994
                    },
    995
    0
                    _ => {},
    996
                }
    997
    998
                // Error output (Jidoka)
    999
    28
                if let Some(ref err1) = event.error {
    1000
    1
                    output.push_str(&format!("  ERROR: {}\n", err));
    1001
    1
                    output.push_str(&format!("  Hint: {}\n", get_error_hint(err)));
    1002
    27
                } else {
    1003
    27
                    output.push_str("  OK\n");
    1004
    27
                }
    1005
    28
                output.push('\n');
    1006
            }
    1007
    1008
            // Summary
    1009
    9
            if self.error_count > 0 {
    1010
    1
                output.push_str(&format!(
    1011
    1
                    "\n=== TRACE SUMMARY: {} errors, {} warnings ===\n",
    1012
    1
                    self.error_count, self.warning_count
    1013
    1
                ));
    1014
    8
            } else {
    1015
    8
                output.push_str("\n=== TRACE COMPLETE: No errors ===\n");
    1016
    8
            }
    1017
    1018
    9
            output
    1019
    9
        }
    1020
    1021
        /// Format trace as JSON
    1022
        #[must_use]
    1023
    6
        pub fn to_json(&self) -> String {
    1024
    6
            let mut json = String::from("{\n");
    1025
    6
            json.push_str("  \"version\": \"1.0\",\n");
    1026
    6
            json.push_str(&format!(
    1027
    6
                "  \"timestamp\": \"{}\",\n",
    1028
    6
                chrono::Utc::now().to_rfc3339()
    1029
    6
            ));
    1030
    1031
            // Model info
    1032
    6
            json.push_str("  \"model\": {\n");
    1033
    6
            json.push_str(&format!("    \"name\": {:?},\n", self.model_info.name));
    1034
    6
            json.push_str(&format!(
    1035
    6
                "    \"num_layers\": {},\n",
    1036
    6
                self.model_info.num_layers
    1037
    6
            ));
    1038
    6
            json.push_str(&format!(
    1039
    6
                "    \"hidden_dim\": {},\n",
    1040
    6
                self.model_info.hidden_dim
    1041
    6
            ));
    1042
    6
            json.push_str(&format!(
    1043
    6
                "    \"vocab_size\": {},\n",
    1044
    6
                self.model_info.vocab_size
    1045
    6
            ));
    1046
    6
            json.push_str(&format!(
    1047
    6
                "    \"num_heads\": {}\n",
    1048
    6
                self.model_info.num_heads
    1049
    6
            ));
    1050
    6
            json.push_str("  },\n");
    1051
    1052
            // Events
    1053
    6
            json.push_str("  \"events\": [\n");
    1054
    10
            for (i, event) in self.events.iter()6.enumerate6() {
    1055
    10
                if i > 0 {
    1056
    5
                    json.push_str(",\n");
    1057
    5
                }
    1058
    10
                json.push_str("    {\n");
    1059
                // AWS Step Functions parity fields (F-AWS-01, F-AWS-02)
    1060
    10
                json.push_str(&format!("      \"id\": {},\n", event.id));
    1061
    10
                json.push_str(&format!("      \"timestamp\": {:?},\n", event.timestamp));
    1062
    10
                json.push_str(&format!("      \"type\": {:?},\n", event.event_type.name()));
    1063
    10
                json.push_str(&format!(
    1064
    10
                    "      \"previous_event_id\": {},\n",
    1065
    10
                    event
    1066
    10
                        .previous_event_id
    1067
    10
                        .map_or("null".to_string(), |id| id5.to_string5())
    1068
                ));
    1069
                // State details
    1070
    10
                json.push_str(&format!("      \"step\": {:?},\n", event.step.name()));
    1071
    10
                json.push_str(&format!("      \"iteration\": {},\n", event.iteration));
    1072
    10
                json.push_str(&format!(
    1073
    10
                    "      \"layer\": {},\n",
    1074
    10
                    event.layer.map_or("null".to_string(), |l| l0.to_string0())
    1075
                ));
    1076
    10
                json.push_str(&format!(
    1077
    10
                    "      \"input_shape\": {:?},\n",
    1078
    10
                    event.input_shape
    1079
    10
                ));
    1080
    10
                json.push_str(&format!(
    1081
    10
                    "      \"output_shape\": {:?},\n",
    1082
    10
                    event.output_shape
    1083
    10
                ));
    1084
    10
                json.push_str(&format!("      \"duration_us\": {},\n", event.duration_us));
    1085
    10
                json.push_str("      \"stats\": {\n");
    1086
    10
                json.push_str(&format!(
    1087
    10
                    "        \"min\": {},\n",
    1088
    10
                    format_json_float(event.stats.min)
    1089
    10
                ));
    1090
    10
                json.push_str(&format!(
    1091
    10
                    "        \"max\": {},\n",
    1092
    10
                    format_json_float(event.stats.max)
    1093
    10
                ));
    1094
    10
                json.push_str(&format!(
    1095
    10
                    "        \"mean\": {},\n",
    1096
    10
                    format_json_float(event.stats.mean)
    1097
    10
                ));
    1098
    10
                json.push_str(&format!(
    1099
    10
                    "        \"std\": {},\n",
    1100
    10
                    format_json_float(event.stats.std)
    1101
    10
                ));
    1102
    10
                json.push_str(&format!("        \"has_nan\": {},\n", event.stats.has_nan));
    1103
    10
                json.push_str(&format!("        \"has_inf\": {}\n", event.stats.has_inf));
    1104
    10
                json.push_str("      },\n");
    1105
    10
                json.push_str(&format!(
    1106
    10
                    "      \"error\": {},\n",
    1107
    10
                    event
    1108
    10
                        .error
    1109
    10
                        .as_ref()
    1110
    10
                        .map_or("null".to_string(), |e| format!2("{:?}"2, e2.to_string2()))
    1111
                ));
    1112
                // F-AWS-05: cause field required for ExecutionFailed events
    1113
    10
                json.push_str(&format!(
    1114
    10
                    "      \"cause\": {}\n",
    1115
    10
                    event
    1116
    10
                        .cause
    1117
    10
                        .as_ref()
    1118
    10
                        .map_or("null".to_string(), |c| format!0("{:?}"0, c))
    1119
                ));
    1120
    10
                json.push_str("    }");
    1121
            }
    1122
    6
            json.push_str("\n  ],\n");
    1123
    1124
            // Summary
    1125
    6
            json.push_str(&format!("  \"error_count\": {},\n", self.error_count));
    1126
    6
            json.push_str(&format!("  \"warning_count\": {}\n", self.warning_count));
    1127
    6
            json.push_str("}\n");
    1128
    1129
    6
            json
    1130
    6
        }
    1131
    1132
        /// Write trace to configured output
    1133
    2
        pub fn write_output(&self) -> std::io::Result<()> {
    1134
    2
            let output = if self.config.output.is_some() {
    1135
    1
                self.to_json()
    1136
            } else {
    1137
    1
                self.format_text()
    1138
            };
    1139
    1140
    2
            if let Some(ref path1) = self.config.output {
    1141
    1
                std::fs::write(path, output)?0;
    1142
    1
            } else {
    1143
    1
                eprint!("{}", output);
    1144
    1
            }
    1145
    1146
    2
            Ok(())
    1147
    2
        }
    1148
    }
    1149
    1150
    /// Get top-k indices with values from logits
    1151
    11
    fn get_top_k_indices(logits: &[f32], k: usize) -> Vec<(u32, f32)> {
    1152
    11
        let mut indexed: Vec<(u32, f32)> = logits
    1153
    11
            .iter()
    1154
    11
            .enumerate()
    1155
    42
            .map11(|(i, &v)| (i as u32, v))
    1156
    11
            .collect();
    1157
    61
        indexed11.sort_by11(|(_, a), (_, b)| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
    1158
    11
        indexed.into_iter().take(k).collect()
    1159
    11
    }
    1160
    1161
    /// Compute softmax probabilities for top-k tokens
    1162
    4
    fn compute_top_k_probs(logits: &[f32], top_k: &[(u32, f32)]) -> Vec<(u32, f32)> {
    1163
        // Find max for numerical stability
    1164
    4
        let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    1165
    1166
        // Compute exp sum for softmax
    1167
    14
        let exp_sum4: f324 = logits4.iter4().map4(|&l| (l - max_logit).exp()).sum4();
    1168
    1169
        // Compute probs for top-k
    1170
    4
        top_k
    1171
    4
            .iter()
    1172
    13
            .map4(|&(idx, logit)| {
    1173
    13
                let prob = (logit - max_logit).exp() / exp_sum;
    1174
    13
                (idx, prob)
    1175
    13
            })
    1176
    4
            .collect()
    1177
    4
    }
    1178
    1179
    /// Check if decoded output contains garbage characters (APR-TOK-001)
    1180
    13
    fn is_garbage_output(text: &str) -> bool {
    1181
    13
        if text.is_empty() {
    1182
    1
            return false;
    1183
    12
        }
    1184
    1185
        // Count suspicious characters (CJK private use, replacement chars, etc.)
    1186
    12
        let suspicious_count = text
    1187
    12
            .chars()
    1188
    71
            .filter12(|&c| {
    1189
                // Unicode replacement character
    1190
    71
                c == '\u{FFFD}'
    1191
                    // Private use area (often indicates bad decoding)
    1192
    63
                    || ('\u{E000}'..='\u{F8FF}').contains(&c)
    1193
                    // CJK Extension B/C/D (rarely used, often garbage)
    1194
    60
                    || ('\u{20000}'..='\u{2FFFF}').contains(&c)
    1195
    71
            })
    1196
    12
            .count();
    1197
    1198
        // If more than 30% suspicious, likely garbage
    1199
    12
        suspicious_count * 3 > text.chars().count()
    1200
    13
    }
    1201
    1202
    /// Get hint for error (Jidoka: actionable feedback)
    1203
    7
    fn get_error_hint(error: &TraceError) -> &'static str {
    1204
    7
        match error {
    1205
            TraceError::VocabOverflow { .. } => {
    1206
    2
                "Check GGUF vocab loading or tokenizer.json compatibility"
    1207
            },
    1208
    1
            TraceError::NaNDetected { .. } => "Check for numerical overflow in matmul or softmax",
    1209
    1
            TraceError::InfDetected { .. } => "Check for division by zero or very large values",
    1210
            TraceError::GarbageOutput { .. } => {
    1211
    1
                "Token ID may not match tokenizer vocab. Check tokenizer.json vs GGUF vocab"
    1212
            },
    1213
    1
            TraceError::UnknownToken { .. } => "Token not in vocabulary. Check tokenizer configuration",
    1214
            TraceError::ShapeMismatch { .. } => {
    1215
    1
                "Tensor dimensions don't match. Check model architecture"
    1216
            },
    1217
            TraceError::ExecutionFailed { .. } => {
    1218
    0
                "Execution failed. Check model config and dependencies"
    1219
            },
    1220
        }
    1221
    7
    }
    1222
    1223
    /// Format float for JSON (handle NaN/Inf)
    1224
    44
    fn format_json_float(v: f32) -> String {
    1225
    44
        if v.is_nan() {
    1226
    1
            "null".to_string()
    1227
    43
        } else if v.is_infinite() {
    1228
    2
            if v.is_sign_positive() {
    1229
    1
                "\"Infinity\"".to_string()
    1230
            } else {
    1231
    1
                "\"-Infinity\"".to_string()
    1232
            }
    1233
        } else {
    1234
    41
            format!("{:.6}", v)
    1235
        }
    1236
    44
    }
    1237
    1238
    // Tests extracted to tests.rs (PMAT-802)
    1239
    #[cfg(test)]
    1240
    #[path = "tests.rs"]
    1241
    mod inference_trace_tests;