trueno/tuner/models/
kernel.rs1#![allow(missing_docs)]
2use serde::{Deserialize, Serialize};
5
6use super::super::features::TunerFeatures;
7use super::super::types::KernelType;
8use super::KernelRecommendation;
9
10#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub struct KernelClassifier {
18 accuracy: f32,
20}
21
22impl KernelClassifier {
23 pub fn new() -> Self {
24 Self { accuracy: 0.85 }
25 }
26
27 pub fn predict(&self, features: &TunerFeatures) -> KernelRecommendation {
29 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 let (top_kernel, confidence) = if batch_size >= 4 {
35 (KernelType::BatchedQ4K, 0.90)
37 } else if batch_size >= 2 {
38 (KernelType::VectorizedQ4K, 0.85)
40 } else {
41 if features.cuda_graphs > 0.5 {
43 (KernelType::VectorizedQ4K, 0.88)
44 } else {
45 (KernelType::CoalescedQ4K, 0.82)
46 }
47 };
48
49 let attention_kernel = if seq_len > 128 {
51 KernelType::MultiWarpAttention
52 } else {
53 KernelType::IncrementalAttention
54 };
55
56 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}