Skip to main content

trueno/tuner/models/
kernel.rs

1#![allow(missing_docs)]
2//! Kernel classifier model for ML tuner.
3
4use serde::{Deserialize, Serialize};
5
6use super::super::features::TunerFeatures;
7use super::super::types::KernelType;
8use super::KernelRecommendation;
9
10/// Kernel classifier using simple rule-based logic.
11///
12/// The `ml-tuner` showcase (SHOWCASE-BRICK-001) trains an
13/// `aprender::RandomForestClassifier` on these features in
14/// `examples/ml_tuner_demo.rs` — aprender is a dev-dependency so the SIMD foundation
15/// stays free of the ML layer (no compute→core layer inversion).
16#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub struct KernelClassifier {
18    /// Kernel accuracy on validation (for confidence)
19    accuracy: f32,
20}
21
22impl KernelClassifier {
23    pub fn new() -> Self {
24        Self { accuracy: 0.85 }
25    }
26
27    /// Predict best kernel based on features
28    pub fn predict(&self, features: &TunerFeatures) -> KernelRecommendation {
29        // Rule-based kernel selection from SHOWCASE-BRICK-001 learnings
30        let batch_size = (features.batch_size_norm * 64.0).round() as u32;
31        let seq_len = (2.0_f32.powf(features.seq_len_log * 15.0)).round() as u32;
32
33        // Determine best Q4K variant based on batch size
34        let (top_kernel, confidence) = if batch_size >= 4 {
35            // M >= 4: Use batched kernels
36            (KernelType::BatchedQ4K, 0.90)
37        } else if batch_size >= 2 {
38            // M = 2-3: Vectorized is good
39            (KernelType::VectorizedQ4K, 0.85)
40        } else {
41            // M = 1: Coalesced or Vectorized
42            if features.cuda_graphs > 0.5 {
43                (KernelType::VectorizedQ4K, 0.88)
44            } else {
45                (KernelType::CoalescedQ4K, 0.82)
46            }
47        };
48
49        // Check for attention-bound cases
50        let attention_kernel = if seq_len > 128 {
51            KernelType::MultiWarpAttention
52        } else {
53            KernelType::IncrementalAttention
54        };
55
56        // Build alternatives
57        let alternatives = vec![
58            (KernelType::VectorizedQ4K, 0.85),
59            (KernelType::CoalescedQ4K, 0.75),
60            (attention_kernel, 0.70),
61        ]
62        .into_iter()
63        .filter(|(k, _)| *k != top_kernel)
64        .take(2)
65        .collect();
66
67        KernelRecommendation { top_kernel, confidence, alternatives }
68    }
69}