## Coverage Report

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

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

    /home/noah/src/trueno/src/tuner/models.rs

    Line
    Count
    Source
    1
    //! ML Models for Tuner
    2
    //!
    3
    //! Throughput regressor, kernel classifier, and bottleneck classifier implementations.
    4
    5
    use serde::{Deserialize, Serialize};
    6
    7
    #[cfg(feature = "ml-tuner")]
    8
    use aprender::{
    9
        tree::{RandomForestClassifier, RandomForestRegressor},
    10
        Matrix, Vector,
    11
    };
    12
    13
    use super::error::TunerError;
    14
    use super::features::TunerFeatures;
    15
    use super::types::{BottleneckClass, KernelType};
    16
    17
    // ============================================================================
    18
    // Prediction Results
    19
    // ============================================================================
    20
    21
    /// Throughput prediction result
    22
    #[derive(Debug, Clone, Serialize, Deserialize)]
    23
    pub struct ThroughputPrediction {
    24
        /// Predicted tokens per second
    25
        pub predicted_tps: f32,
    26
        /// Confidence (0-1)
    27
        pub confidence: f32,
    28
        /// Top contributing features
    29
        pub top_features: Vec<(String, f32)>,
    30
    }
    31
    32
    /// Kernel recommendation result
    33
    #[derive(Debug, Clone, Serialize, Deserialize)]
    34
    pub struct KernelRecommendation {
    35
        /// Top recommended kernel
    36
        pub top_kernel: KernelType,
    37
        /// Confidence (0-1)
    38
        pub confidence: f32,
    39
        /// Alternative kernels with probabilities
    40
        pub alternatives: Vec<(KernelType, f32)>,
    41
    }
    42
    43
    /// Bottleneck prediction result
    44
    #[derive(Debug, Clone, Serialize, Deserialize)]
    45
    pub struct BottleneckPrediction {
    46
        /// Predicted bottleneck class
    47
        pub class: BottleneckClass,
    48
        /// Confidence (0-1)
    49
        pub confidence: f32,
    50
        /// Human-readable explanation
    51
        pub explanation: String,
    52
        /// Recommended action
    53
        pub recommended_action: String,
    54
    }
    55
    56
    // ============================================================================
    57
    // ThroughputRegressor
    58
    // ============================================================================
    59
    60
    /// Simple linear regression model for throughput prediction.
    61
    ///
    62
    /// Uses closed-form solution: w = (X^T X)^-1 X^T y
    63
    /// With `ml-tuner` feature: uses aprender::RandomForestRegressor (SHOWCASE-BRICK-001)
    64
    #[derive(Debug, Clone, Serialize, Deserialize)]
    65
    pub struct ThroughputRegressor {
    66
        /// Model weights (one per feature + bias) - fallback when ml-tuner disabled
    67
        pub(crate) weights: Vec<f32>,
    68
        /// Feature importance scores
    69
        pub(crate) feature_importance: Vec<(String, f32)>,
    70
        /// Training sample count
    71
        pub(crate) sample_count: usize,
    72
        /// Mean absolute percentage error on validation
    73
        pub(crate) mape: f32,
    74
        /// Whether the RandomForest model is trained (ml-tuner feature)
    75
        #[cfg(feature = "ml-tuner")]
    76
        #[serde(skip)]
    77
        rf_model: Option<RandomForestRegressor>,
    78
    }
    79
    80
    impl Default for ThroughputRegressor {
    81
    0
        fn default() -> Self {
    82
    0
            Self::new()
    83
    0
        }
    84
    }
    85
    86
    impl ThroughputRegressor {
    87
        /// Create a new regressor with default weights
    88
    0
        pub fn new() -> Self {
    89
            // Initialize with heuristic-based weights
    90
            // These encode domain knowledge from SHOWCASE-BRICK-001
    91
    0
            let mut weights = vec![0.0; TunerFeatures::DIM + 1]; // +1 for bias
    92
    93
            // Bias: baseline throughput ~200 tok/s normalized
    94
    0
            weights[0] = 0.4;
    95
    96
            // Batch size has largest positive impact (index 6)
    97
    0
            weights[7] = 0.3; // batch_size_norm
    98
    99
            // CUDA graphs help (index 8)
    100
    0
            weights[9] = 0.1; // cuda_graphs
    101
    102
            // GPU memory bandwidth matters (index 35)
    103
    0
            weights[36] = 0.15; // gpu_mem_bw_norm
    104
    105
            // GPU SM count matters (index 37)
    106
    0
            weights[38] = 0.1; // gpu_sm_norm
    107
    108
            // Larger models are slower (negative impact)
    109
    0
            weights[1] = -0.15; // model_params_b
    110
    111
            // Longer sequences slower for decode
    112
    0
            weights[8] = -0.05; // seq_len_log
    113
    114
    0
            Self {
    115
    0
                weights,
    116
    0
                feature_importance: Self::default_feature_importance(),
    117
    0
                sample_count: 0,
    118
    0
                mape: 0.15, // 15% default MAPE
    119
    0
                #[cfg(feature = "ml-tuner")]
    120
    0
                rf_model: None,
    121
    0
            }
    122
    0
        }
    123
    124
        /// Create a new regressor using aprender RandomForest (ml-tuner feature)
    125
        #[cfg(feature = "ml-tuner")]
    126
        pub fn with_random_forest(n_estimators: usize) -> Self {
    127
            let mut instance = Self::new();
    128
            instance.rf_model = Some(RandomForestRegressor::new(n_estimators));
    129
            instance
    130
        }
    131
    132
    0
        fn default_feature_importance() -> Vec<(String, f32)> {
    133
    0
            vec![
    134
    0
                ("batch_size".into(), 0.25),
    135
    0
                ("gpu_mem_bw".into(), 0.20),
    136
    0
                ("model_params".into(), 0.15),
    137
    0
                ("cuda_graphs".into(), 0.10),
    138
    0
                ("gpu_sm_count".into(), 0.10),
    139
    0
                ("hidden_dim".into(), 0.08),
    140
    0
                ("quant_type".into(), 0.07),
    141
    0
                ("seq_len".into(), 0.05),
    142
            ]
    143
    0
        }
    144
    145
        /// Train the model on labeled data
    146
    0
        pub fn train(&mut self, data: &[(TunerFeatures, f32)]) -> Result<(), TunerError> {
    147
    0
            if data.len() < 10 {
    148
    0
                return Err(TunerError::InsufficientData(data.len()));
    149
    0
            }
    150
    151
            // Simple gradient descent (in production: aprender's GBDT)
    152
    0
            let learning_rate = 0.01;
    153
    0
            let epochs = 100;
    154
    155
    0
            for _ in 0..epochs {
    156
    0
                let mut gradients = vec![0.0; self.weights.len()];
    157
    158
    0
                for (features, target) in data {
    159
    0
                    let x = features.to_vector();
    160
    0
                    let predicted = self.predict_raw(&x);
    161
    0
                    let error = predicted - target;
    162
    163
                    // Gradient for bias
    164
    0
                    gradients[0] += error;
    165
    166
                    // Gradient for features
    167
    0
                    for (i, xi) in x.iter().enumerate() {
    168
    0
                        gradients[i + 1] += error * xi;
    169
    0
                    }
    170
                }
    171
    172
                // Update weights
    173
    0
                let n = data.len() as f32;
    174
    0
                for (i, g) in gradients.iter().enumerate() {
    175
    0
                    self.weights[i] -= learning_rate * g / n;
    176
    0
                }
    177
            }
    178
    179
            // Calculate MAPE on training data
    180
    0
            let mut total_ape = 0.0;
    181
    0
            for (features, target) in data {
    182
    0
                let predicted = self.predict_raw(&features.to_vector());
    183
    0
                total_ape += ((predicted - target) / target.max(1.0)).abs();
    184
    0
            }
    185
    0
            self.mape = total_ape / data.len() as f32;
    186
    0
            self.sample_count = data.len();
    187
    188
    0
            Ok(())
    189
    0
        }
    190
    191
        /// Train using aprender RandomForest (ml-tuner feature)
    192
        ///
    193
        /// Provides superior throughput prediction via ensemble learning.
    194
        /// See: SHOWCASE-BRICK-001 Section 12.3
    195
        #[cfg(feature = "ml-tuner")]
    196
        pub fn train_random_forest(&mut self, data: &[(TunerFeatures, f32)]) -> Result<(), TunerError> {
    197
            if data.len() < 10 {
    198
                return Err(TunerError::InsufficientData(data.len()));
    199
            }
    200
    201
            // Convert to aprender matrix format (f32 for RandomForestRegressor)
    202
            let n_samples = data.len();
    203
            let n_features = TunerFeatures::DIM;
    204
            let mut x_data = Vec::with_capacity(n_samples * n_features);
    205
            let mut y_data = Vec::with_capacity(n_samples);
    206
    207
            for (features, target) in data {
    208
                x_data.extend(features.to_vector());
    209
                y_data.push(*target);
    210
            }
    211
    212
            let x_matrix = Matrix::from_vec(n_samples, n_features, x_data)
    213
                .map_err(|e| TunerError::TrainingFailed(e.to_string()))?;
    214
            let y_vector = Vector::from_vec(y_data);
    215
    216
            // Train RandomForest
    217
            let rf = self
    218
                .rf_model
    219
                .get_or_insert_with(|| RandomForestRegressor::new(100));
    220
            rf.fit(&x_matrix, &y_vector)
    221
                .map_err(|e| TunerError::TrainingFailed(e.to_string()))?;
    222
    223
            // Calculate MAPE on training data
    224
            let predictions = rf.predict(&x_matrix);
    225
            let mut total_ape = 0.0;
    226
            for (i, (_, target)) in data.iter().enumerate() {
    227
                let pred = predictions.as_slice()[i];
    228
                total_ape += ((pred - target) / target.max(1.0)).abs();
    229
            }
    230
            self.mape = total_ape / data.len() as f32;
    231
            self.sample_count = data.len();
    232
    233
            Ok(())
    234
        }
    235
    236
    0
        pub(crate) fn predict_raw(&self, x: &[f32]) -> f32 {
    237
    0
            let mut result = self.weights[0]; // bias
    238
    0
            for (i, xi) in x.iter().enumerate() {
    239
    0
                if i + 1 < self.weights.len() {
    240
    0
                    result += self.weights[i + 1] * xi;
    241
    0
                }
    242
            }
    243
            // Convert from normalized to tok/s (scale ~1000)
    244
    0
            (result * 1000.0).max(1.0)
    245
    0
        }
    246
    247
        /// Predict throughput for features
    248
        ///
    249
        /// With `ml-tuner` feature: uses trained RandomForest if available.
    250
        /// Falls back to linear model otherwise.
    251
    0
        pub fn predict(&self, features: &TunerFeatures) -> ThroughputPrediction {
    252
    0
            let x = features.to_vector();
    253
    254
            // Use RandomForest if trained (ml-tuner feature)
    255
            #[cfg(feature = "ml-tuner")]
    256
            let raw_predicted_tps = if let Some(ref rf) = self.rf_model {
    257
                // Use f32 matrix for RandomForestRegressor
    258
                if let Ok(x_matrix) = Matrix::from_vec(1, TunerFeatures::DIM, x.to_vec()) {
    259
                    let predictions = rf.predict(&x_matrix);
    260
                    predictions.as_slice().first().copied().unwrap_or(0.0)
    261
                } else {
    262
                    self.predict_raw(&x)
    263
                }
    264
            } else {
    265
                self.predict_raw(&x)
    266
            };
    267
    268
            #[cfg(not(feature = "ml-tuner"))]
    269
    0
            let raw_predicted_tps = self.predict_raw(&x);
    270
    271
            // v1.1.0: Roofline clamping - predictions must not exceed theoretical maximum
    272
    0
            let theoretical_max_tps = Self::compute_roofline_bound(features);
    273
    0
            let predicted_tps = raw_predicted_tps.min(theoretical_max_tps);
    274
    275
            // Confidence based on training MAPE and feature validity
    276
            // Lower confidence if we hit the roofline cap
    277
    0
            let roofline_penalty = if raw_predicted_tps > theoretical_max_tps {
    278
    0
                0.9 // 10% confidence penalty for capped predictions
    279
            } else {
    280
    0
                1.0
    281
            };
    282
    0
            let confidence = (1.0 - self.mape).max(0.5) * roofline_penalty;
    283
    284
    0
            ThroughputPrediction {
    285
    0
                predicted_tps,
    286
    0
                confidence,
    287
    0
                top_features: self.feature_importance.iter().take(5).cloned().collect(),
    288
    0
            }
    289
    0
        }
    290
    291
        /// Compute theoretical maximum throughput based on roofline model (v1.1.0)
    292
        ///
    293
        /// For memory-bound LLM inference (decode phase):
    294
        /// max_tps = memory_bw_bytes_per_sec / bytes_per_token
    295
        /// bytes_per_token = model_params × bytes_per_param / batch_size
    296
    0
        pub fn compute_roofline_bound(features: &TunerFeatures) -> f32 {
    297
            // Denormalize model params: normalized = (log10(b) + 1) / 3
    298
            // log10(b) = normalized * 3 - 1
    299
            // b = 10^(normalized * 3 - 1)
    300
    0
            let model_params_b = 10.0_f32.powf(features.model_params_b * 3.0 - 1.0);
    301
    302
            // Get bytes per param from quant type one-hot encoding
    303
    0
            let bytes_per_param = Self::bytes_per_param_from_onehot(&features.quant_type_onehot);
    304
    305
            // Denormalize memory bandwidth: normalized = bw / 3000 GB/s
    306
    0
            let gpu_mem_bw_gbs = features.gpu_mem_bw_norm * 3000.0;
    307
    308
            // Denormalize batch size: normalized = batch_size / 64
    309
    0
            let batch_size = (features.batch_size_norm * 64.0).max(1.0);
    310
    311
            // Roofline calculation:
    312
            // model_bytes = model_params_b * bytes_per_param * 1e9
    313
            // bytes_per_token = model_bytes / batch_size
    314
            // max_tps = (gpu_mem_bw_gbs * 1e9) / bytes_per_token
    315
            //         = (gpu_mem_bw_gbs * 1e9 * batch_size) / (model_params_b * bytes_per_param * 1e9)
    316
            //         = (gpu_mem_bw_gbs * batch_size) / (model_params_b * bytes_per_param)
    317
    0
            let theoretical_max = (gpu_mem_bw_gbs * batch_size) / (model_params_b * bytes_per_param);
    318
    319
            // Clamp to reasonable range (1 tok/s to 10000 tok/s)
    320
    0
            theoretical_max.clamp(1.0, 10000.0)
    321
    0
        }
    322
    323
        /// Extract bytes per param from quant type one-hot encoding
    324
    0
        pub fn bytes_per_param_from_onehot(onehot: &[f32; 8]) -> f32 {
    325
            // One-hot indices map to QuantType variants
    326
            // 0: Q4_0, 1: Q4_1, 2: Q4K, 3: Q5K, 4: Q6K, 5: Q8_0, 6: F16, 7: F32
    327
    0
            let bytes_per_param = [0.5625, 0.5625, 0.5625, 0.6875, 0.8125, 1.0, 2.0, 4.0];
    328
    329
            // Find the active index (max value in one-hot)
    330
    0
            let idx = onehot
    331
    0
                .iter()
    332
    0
                .enumerate()
    333
    0
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    334
    0
                .map(|(i, _)| i)
    335
    0
                .unwrap_or(2); // Default to Q4K if ambiguous
    336
    337
    0
            bytes_per_param[idx]
    338
    0
        }
    339
    }
    340
    341
    // ============================================================================
    342
    // KernelClassifier
    343
    // ============================================================================
    344
    345
    /// Kernel classifier using simple rule-based logic.
    346
    ///
    347
    /// With `ml-tuner` feature: uses aprender::RandomForestClassifier (SHOWCASE-BRICK-001)
    348
    #[derive(Debug, Clone, Serialize, Deserialize, Default)]
    349
    pub struct KernelClassifier {
    350
        /// Kernel accuracy on validation (for confidence)
    351
        accuracy: f32,
    352
        /// RandomForest classifier when ml-tuner feature is enabled
    353
        #[cfg(feature = "ml-tuner")]
    354
        #[serde(skip)]
    355
        rf_classifier: Option<RandomForestClassifier>,
    356
    }
    357
    358
    impl KernelClassifier {
    359
    0
        pub fn new() -> Self {
    360
    0
            Self {
    361
    0
                accuracy: 0.85,
    362
    0
                #[cfg(feature = "ml-tuner")]
    363
    0
                rf_classifier: None,
    364
    0
            }
    365
    0
        }
    366
    367
        /// Create a classifier with aprender RandomForest (ml-tuner feature)
    368
        #[cfg(feature = "ml-tuner")]
    369
        pub fn with_random_forest(n_estimators: usize) -> Self {
    370
            Self {
    371
                accuracy: 0.85,
    372
                rf_classifier: Some(RandomForestClassifier::new(n_estimators)),
    373
            }
    374
        }
    375
    376
        /// Train the classifier using aprender RandomForest (ml-tuner feature)
    377
        ///
    378
        /// Labels should be kernel type indices (0=TiledQ4K, 1=CoalescedQ4K, etc.)
    379
        #[cfg(feature = "ml-tuner")]
    380
        pub fn train(&mut self, data: &[(TunerFeatures, u32)]) -> Result<(), TunerError> {
    381
            if data.len() < 10 {
    382
                return Err(TunerError::InsufficientData(data.len()));
    383
            }
    384
    385
            // Convert to aprender format (Matrix<f32> for features, &[usize] for labels)
    386
            let n_samples = data.len();
    387
            let n_features = TunerFeatures::DIM;
    388
            let mut x_data = Vec::with_capacity(n_samples * n_features);
    389
            let mut y_data: Vec<usize> = Vec::with_capacity(n_samples);
    390
    391
            for (features, label) in data {
    392
                x_data.extend(features.to_vector());
    393
                y_data.push(*label as usize);
    394
            }
    395
    396
            let x_matrix = Matrix::from_vec(n_samples, n_features, x_data)
    397
                .map_err(|e| TunerError::TrainingFailed(e.to_string()))?;
    398
    399
            let rf = self
    400
                .rf_classifier
    401
                .get_or_insert_with(|| RandomForestClassifier::new(50));
    402
            rf.fit(&x_matrix, &y_data)
    403
                .map_err(|e| TunerError::TrainingFailed(e.to_string()))?;
    404
    405
            // Calculate accuracy on training data
    406
            let predictions = rf.predict(&x_matrix);
    407
            let mut correct = 0;
    408
            for (i, (_, label)) in data.iter().enumerate() {
    409
                if predictions[i] as u32 == *label {
    410
                    correct += 1;
    411
                }
    412
            }
    413
            self.accuracy = correct as f32 / data.len() as f32;
    414
    415
            Ok(())
    416
        }
    417
    418
        /// Predict best kernel based on features
    419
    0
        pub fn predict(&self, features: &TunerFeatures) -> KernelRecommendation {
    420
            // Rule-based kernel selection from SHOWCASE-BRICK-001 learnings
    421
    0
            let batch_size = (features.batch_size_norm * 64.0).round() as u32;
    422
    0
            let seq_len = (2.0_f32.powf(features.seq_len_log * 15.0)).round() as u32;
    423
    424
            // Determine best Q4K variant based on batch size
    425
    0
            let (top_kernel, confidence) = if batch_size >= 4 {
    426
                // M >= 4: Use batched kernels
    427
    0
                (KernelType::BatchedQ4K, 0.90)
    428
    0
            } else if batch_size >= 2 {
    429
                // M = 2-3: Vectorized is good
    430
    0
                (KernelType::VectorizedQ4K, 0.85)
    431
            } else {
    432
                // M = 1: Coalesced or Vectorized
    433
    0
                if features.cuda_graphs > 0.5 {
    434
    0
                    (KernelType::VectorizedQ4K, 0.88)
    435
                } else {
    436
    0
                    (KernelType::CoalescedQ4K, 0.82)
    437
                }
    438
            };
    439
    440
            // Check for attention-bound cases
    441
    0
            let attention_kernel = if seq_len > 128 {
    442
    0
                KernelType::MultiWarpAttention
    443
            } else {
    444
    0
                KernelType::IncrementalAttention
    445
            };
    446
    447
            // Build alternatives
    448
    0
            let alternatives = vec![
    449
    0
                (KernelType::VectorizedQ4K, 0.85),
    450
    0
                (KernelType::CoalescedQ4K, 0.75),
    451
    0
                (attention_kernel, 0.70),
    452
            ]
    453
    0
            .into_iter()
    454
    0
            .filter(|(k, _)| *k != top_kernel)
    455
    0
            .take(2)
    456
    0
            .collect();
    457
    458
    0
            KernelRecommendation {
    459
    0
                top_kernel,
    460
    0
                confidence,
    461
    0
                alternatives,
    462
    0
            }
    463
    0
        }
    464
    }
    465
    466
    // ============================================================================
    467
    // BottleneckClassifier
    468
    // ============================================================================
    469
    470
    /// Bottleneck classifier using heuristics from profiler data.
    471
    #[derive(Debug, Clone, Serialize, Deserialize, Default)]
    472
    pub struct BottleneckClassifier {
    473
        /// Classification accuracy
    474
        accuracy: f32,
    475
    }
    476
    477
    impl BottleneckClassifier {
    478
    0
        pub fn new() -> Self {
    479
    0
            Self { accuracy: 0.90 }
    480
    0
        }
    481
    482
        /// Predict bottleneck from features
    483
    0
        pub fn predict(&self, features: &TunerFeatures) -> BottleneckPrediction {
    484
            // Use already-computed bottleneck if available
    485
    0
            if let Some(class) = features.bottleneck_class {
    486
    0
                return BottleneckPrediction {
    487
    0
                    class,
    488
    0
                    confidence: 0.95,
    489
    0
                    explanation: format!("Bottleneck classified from profiler data: {}", class),
    490
    0
                    recommended_action: class.recommended_action().to_string(),
    491
    0
                };
    492
    0
            }
    493
    494
            // Heuristic classification based on features
    495
    0
            let batch_size = (features.batch_size_norm * 64.0).round() as u32;
    496
    0
            let seq_len = (2.0_f32.powf(features.seq_len_log * 15.0)).round() as u32;
    497
    498
    0
            let (class, confidence, explanation) = if batch_size == 1 && features.cuda_graphs < 0.5 {
    499
    0
                (
    500
    0
                    BottleneckClass::LaunchBound,
    501
    0
                    0.75,
    502
    0
                    "Single sequence without CUDA graphs: kernel launch overhead may dominate".into(),
    503
    0
                )
    504
    0
            } else if seq_len > 512 {
    505
    0
                (
    506
    0
                    BottleneckClass::AttentionBound,
    507
    0
                    0.80,
    508
    0
                    format!(
    509
    0
                        "Long sequence (len={}) likely makes attention the bottleneck",
    510
    0
                        seq_len
    511
    0
                    ),
    512
    0
                )
    513
            } else {
    514
    0
                (
    515
    0
                    BottleneckClass::MemoryBound,
    516
    0
                    0.85,
    517
    0
                    "Q4K GEMV is typically memory-bound for LLM inference".into(),
    518
    0
                )
    519
            };
    520
    521
    0
            BottleneckPrediction {
    522
    0
                class,
    523
    0
                confidence,
    524
    0
                explanation,
    525
    0
                recommended_action: class.recommended_action().to_string(),
    526
    0
            }
    527
    0
        }
    528
    }