#![allow(missing_docs)]
use serde::{Deserialize, Serialize};
use super::super::features::TunerFeatures;
use super::super::types::KernelType;
use super::KernelRecommendation;
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct KernelClassifier {
accuracy: f32,
}
impl KernelClassifier {
pub fn new() -> Self {
Self { accuracy: 0.85 }
}
pub fn predict(&self, features: &TunerFeatures) -> KernelRecommendation {
let batch_size = (features.batch_size_norm * 64.0).round() as u32;
let seq_len = (2.0_f32.powf(features.seq_len_log * 15.0)).round() as u32;
let (top_kernel, confidence) = if batch_size >= 4 {
(KernelType::BatchedQ4K, 0.90)
} else if batch_size >= 2 {
(KernelType::VectorizedQ4K, 0.85)
} else {
if features.cuda_graphs > 0.5 {
(KernelType::VectorizedQ4K, 0.88)
} else {
(KernelType::CoalescedQ4K, 0.82)
}
};
let attention_kernel = if seq_len > 128 {
KernelType::MultiWarpAttention
} else {
KernelType::IncrementalAttention
};
let alternatives = vec![
(KernelType::VectorizedQ4K, 0.85),
(KernelType::CoalescedQ4K, 0.75),
(attention_kernel, 0.70),
]
.into_iter()
.filter(|(k, _)| *k != top_kernel)
.take(2)
.collect();
KernelRecommendation { top_kernel, confidence, alternatives }
}
}