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

Source
//! Feature Extraction for ML Tuning
//!
//! Implements TunerFeatures, TunerFeaturesBuilder, FeatureExtractor, and RunConfig.

use crate::brick::{BrickCategory, BrickProfiler};
use crate::hardware::HardwareCapability;
use serde::{Deserialize, Serialize};

use super::error::TunerError;
use super::types::{BottleneckClass, KernelType, QuantType};

// ============================================================================
// TunerFeatures
// ============================================================================

/// Feature vector for ML-based kernel tuning.
///
/// All fields normalized to [0, 1] for model input.
/// Total dimension: 42 features.
///
/// # Feature Categories
///
/// - **Static (11)**: Known before execution (model size, batch size, etc.)
/// - **Quant one-hot (8)**: Quantization type encoding
/// - **Kernel one-hot (16)**: Kernel type encoding
/// - **Hardware (5)**: GPU capabilities
/// - **Derived (2)**: Computed features (arithmetic intensity, efficiency)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TunerFeatures {
    // === Static features (11) ===
    /// Model size in billions (log10 normalized)
    pub model_params_b: f32,
    /// Hidden dimension / 16384
    pub hidden_dim_norm: f32,
    /// Number of layers / 128
    pub num_layers_norm: f32,
    /// Number of attention heads / 128
    pub num_heads_norm: f32,
    /// Head dimension / 256
    pub head_dim_norm: f32,
    /// Vocabulary size (log10 normalized)
    pub vocab_size_log: f32,
    /// Batch size M / 64
    pub batch_size_norm: f32,
    /// Sequence length (log2 / 15)
    pub seq_len_log: f32,
    /// CUDA graphs enabled (0 or 1)
    pub cuda_graphs: f32,
    /// Number of KV caches / batch_size (for multi-cache detection)
    pub kv_cache_ratio: f32,
    /// Prefill vs decode (0=decode, 1=prefill)
    pub is_prefill: f32,

    // === Quantization one-hot (8) ===
    pub quant_type_onehot: [f32; 8],

    // === Kernel one-hot (16) ===
    pub kernel_type_onehot: [f32; 16],

    // === Hardware features (5) === [v1.1.0: added L2 cache + zero-copy]
    /// Memory bandwidth / 3000 GB/s
    pub gpu_mem_bw_norm: f32,
    /// Compute TFLOPS / 500
    pub gpu_compute_norm: f32,
    /// SM count / 200
    pub gpu_sm_norm: f32,
    /// L2 cache size / 128 MB (v1.1.0: critical for occupancy)
    pub gpu_l2_cache_norm: f32,
    /// Zero-copy memory path enabled (0 or 1) (v1.1.0: pinned memory)
    pub is_zero_copy: f32,

    // === Derived features (2) ===
    /// Arithmetic intensity (FLOP/byte), normalized
    pub arithmetic_intensity: f32,
    /// Theoretical efficiency (measured / roofline)
    pub theoretical_efficiency: f32,

    // === Labels (for training) ===
    /// Measured throughput (tokens/second) - training label
    #[serde(skip_serializing_if = "Option::is_none")]
    pub measured_tps: Option<f32>,
    /// Best kernel ID - classification label
    #[serde(skip_serializing_if = "Option::is_none")]
    pub best_kernel_id: Option<u8>,
    /// Bottleneck class - classification label
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bottleneck_class: Option<BottleneckClass>,
}

impl Default for TunerFeatures {
    fn default() -> Self {
        Self {
            model_params_b: 0.0,
            hidden_dim_norm: 0.0,
            num_layers_norm: 0.0,
            num_heads_norm: 0.0,
            head_dim_norm: 0.0,
            vocab_size_log: 0.0,
            batch_size_norm: 0.0,
            seq_len_log: 0.0,
            cuda_graphs: 0.0,
            kv_cache_ratio: 1.0,
            is_prefill: 0.0,
            quant_type_onehot: [0.0; 8],
            kernel_type_onehot: [0.0; 16],
            gpu_mem_bw_norm: 0.0,
            gpu_compute_norm: 0.0,
            gpu_sm_norm: 0.0,
            gpu_l2_cache_norm: 0.0,
            is_zero_copy: 0.0,
            arithmetic_intensity: 0.0,
            theoretical_efficiency: 0.0,
            measured_tps: None,
            best_kernel_id: None,
            bottleneck_class: None,
        }
    }
}

impl TunerFeatures {
    /// Total feature dimension (excluding labels)
    /// v1.1.0: 11 static + 8 quant + 16 kernel + 5 hardware + 2 derived = 42
    pub const DIM: usize = 11 + 8 + 16 + 5 + 2; // 42 features (v1.1.0)

    /// Create a new feature builder
    pub fn builder() -> TunerFeaturesBuilder {
        TunerFeaturesBuilder::default()
    }

    /// Convert to flat vector for model input
    pub fn to_vector(&self) -> Vec<f32> {
        let mut v = Vec::with_capacity(Self::DIM);

        // Static features
        v.push(self.model_params_b);
        v.push(self.hidden_dim_norm);
        v.push(self.num_layers_norm);
        v.push(self.num_heads_norm);
        v.push(self.head_dim_norm);
        v.push(self.vocab_size_log);
        v.push(self.batch_size_norm);
        v.push(self.seq_len_log);
        v.push(self.cuda_graphs);
        v.push(self.kv_cache_ratio);
        v.push(self.is_prefill);

        // One-hot encodings
        v.extend_from_slice(&self.quant_type_onehot);
        v.extend_from_slice(&self.kernel_type_onehot);

        // Hardware features (5) [v1.1.0]
        v.push(self.gpu_mem_bw_norm);
        v.push(self.gpu_compute_norm);
        v.push(self.gpu_sm_norm);
        v.push(self.gpu_l2_cache_norm); // v1.1.0
        v.push(self.is_zero_copy); // v1.1.0

        // Derived features
        v.push(self.arithmetic_intensity);
        v.push(self.theoretical_efficiency);

        v
    }

    /// Validate features (F021-F030 falsification criteria)
    pub fn validate(&self) -> Result<(), TunerError> {
        let v = self.to_vector();

        // F021: No NaN features
        if v.iter().any(|x| x.is_nan()) {
            return Err(TunerError::InvalidFeature("NaN value in features".into()));
        }

        // F022: No infinite features
        if v.iter().any(|x| x.is_infinite()) {
            return Err(TunerError::InvalidFeature(
                "Infinite value in features".into(),
            ));
        }

        // F023: All features in [0, 1] (with small tolerance for floating point)
        if v.iter().any(|x| *x < -0.001 || *x > 1.001) {
            return Err(TunerError::InvalidFeature(
                "Feature value outside [0, 1]".into(),
            ));
        }

        // F029: One-hot sums = 1
        let quant_sum: f32 = self.quant_type_onehot.iter().sum();
        if (quant_sum - 1.0).abs() > 0.001 && quant_sum > 0.001 {
            return Err(TunerError::InvalidFeature(
                "Quant one-hot does not sum to 1".into(),
            ));
        }

        let kernel_sum: f32 = self.kernel_type_onehot.iter().sum();
        if (kernel_sum - 1.0).abs() > 0.001 && kernel_sum > 0.001 {
            return Err(TunerError::InvalidFeature(
                "Kernel one-hot does not sum to 1".into(),
            ));
        }

        Ok(())
    }
}

// ============================================================================
// TunerFeaturesBuilder
// ============================================================================

/// Builder for TunerFeatures with automatic normalization.
#[derive(Default)]
pub struct TunerFeaturesBuilder {
    model_params_b: Option<f32>,
    hidden_dim: Option<u32>,
    num_layers: Option<u32>,
    num_heads: Option<u32>,
    head_dim: Option<u32>,
    vocab_size: Option<u32>,
    batch_size: Option<u32>,
    seq_len: Option<u32>,
    cuda_graphs: bool,
    kv_caches: Option<u32>,
    is_prefill: bool,
    quant_type: Option<QuantType>,
    kernel_type: Option<KernelType>,
    gpu_mem_bw_gbs: Option<f32>,
    gpu_compute_tflops: Option<f32>,
    gpu_sm_count: Option<u32>,
    gpu_l2_cache_mb: Option<f32>, // v1.1.0
    is_zero_copy: bool,           // v1.1.0
    measured_tps: Option<f32>,
}

impl TunerFeaturesBuilder {
    /// Set model size in billions of parameters
    pub fn model_params_b(mut self, params: f32) -> Self {
        self.model_params_b = Some(params);
        self
    }

    /// Set hidden dimension
    pub fn hidden_dim(mut self, dim: u32) -> Self {
        self.hidden_dim = Some(dim);
        self
    }

    /// Set number of layers
    pub fn num_layers(mut self, layers: u32) -> Self {
        self.num_layers = Some(layers);
        self
    }

    /// Set number of attention heads
    pub fn num_heads(mut self, heads: u32) -> Self {
        self.num_heads = Some(heads);
        self
    }

    /// Set head dimension
    pub fn head_dim(mut self, dim: u32) -> Self {
        self.head_dim = Some(dim);
        self
    }

    /// Set vocabulary size
    pub fn vocab_size(mut self, size: u32) -> Self {
        self.vocab_size = Some(size);
        self
    }

    /// Set batch size (M)
    pub fn batch_size(mut self, m: u32) -> Self {
        self.batch_size = Some(m);
        self
    }

    /// Set sequence length
    pub fn seq_len(mut self, len: u32) -> Self {
        self.seq_len = Some(len);
        self
    }

    /// Enable CUDA graphs
    pub fn cuda_graphs(mut self, enabled: bool) -> Self {
        self.cuda_graphs = enabled;
        self
    }

    /// Set number of KV caches
    pub fn kv_caches(mut self, count: u32) -> Self {
        self.kv_caches = Some(count);
        self
    }

    /// Set prefill mode
    pub fn is_prefill(mut self, prefill: bool) -> Self {
        self.is_prefill = prefill;
        self
    }

    /// Set quantization type
    pub fn quant_type(mut self, qt: QuantType) -> Self {
        self.quant_type = Some(qt);
        self
    }

    /// Set kernel type
    pub fn kernel_type(mut self, kt: KernelType) -> Self {
        self.kernel_type = Some(kt);
        self
    }

    /// Set GPU memory bandwidth in GB/s
    pub fn gpu_mem_bw_gbs(mut self, bw: f32) -> Self {
        self.gpu_mem_bw_gbs = Some(bw);
        self
    }

    /// Set GPU compute in TFLOPS
    pub fn gpu_compute_tflops(mut self, tflops: f32) -> Self {
        self.gpu_compute_tflops = Some(tflops);
        self
    }

    /// Set GPU SM count
    pub fn gpu_sm_count(mut self, count: u32) -> Self {
        self.gpu_sm_count = Some(count);
        self
    }

    /// Set measured throughput (for training data)
    pub fn measured_tps(mut self, tps: f32) -> Self {
        self.measured_tps = Some(tps);
        self
    }

    /// Set L2 cache size in MB (v1.1.0)
    pub fn gpu_l2_cache_mb(mut self, l2_mb: f32) -> Self {
        self.gpu_l2_cache_mb = Some(l2_mb);
        self
    }

    /// Set zero-copy memory path enabled (v1.1.0)
    pub fn is_zero_copy(mut self, enabled: bool) -> Self {
        self.is_zero_copy = enabled;
        self
    }

    /// Set hardware capability (auto-fills GPU features)
    pub fn hardware(mut self, hw: &HardwareCapability) -> Self {
        if let Some(gpu) = &hw.gpu {
            self.gpu_mem_bw_gbs = Some(gpu.memory_bw_gbps as f32);
            self.gpu_compute_tflops = Some(gpu.peak_tflops_fp32 as f32);
            // SM count not directly available; estimate from compute capability
            self.gpu_sm_count = None;
        }
        self
    }

    /// Build the feature vector with normalization
    pub fn build(self) -> TunerFeatures {
        let batch_size = self.batch_size.unwrap_or(1);
        let kv_caches = self.kv_caches.unwrap_or(batch_size);

        // Create one-hot encodings
        let mut quant_onehot = [0.0f32; 8];
        if let Some(qt) = self.quant_type {
            quant_onehot[qt.to_index()] = 1.0;
        }

        let mut kernel_onehot = [0.0f32; 16];
        if let Some(kt) = self.kernel_type {
            kernel_onehot[kt.to_index()] = 1.0;
        }

        // Calculate derived features
        let hidden_dim = self.hidden_dim.unwrap_or(1536) as f32;
        let batch_size_f = batch_size as f32;
        let quant_bytes = self
            .quant_type
            .map(|q| q.bytes_per_param())
            .unwrap_or(0.5625);

        // Arithmetic intensity for GEMV: 2*N*K FLOPs / (N*K*bytes + K + N) bytes
        // Simplified: ~2 / bytes_per_param for memory-bound inference
        let arithmetic_intensity = (2.0 / quant_bytes).min(10.0) / 10.0;

        // Theoretical efficiency starts at 0 (unknown until measured)
        let theoretical_efficiency = 0.0;

        TunerFeatures {
            // Normalized static features
            model_params_b: self
                .model_params_b
                .map(|p| (p.log10() + 1.0) / 3.0) // log10(0.1)=-1, log10(100)=2 → [0, 1]
                .unwrap_or(0.0)
                .clamp(0.0, 1.0),
            hidden_dim_norm: (hidden_dim / 16384.0).clamp(0.0, 1.0),
            num_layers_norm: (self.num_layers.unwrap_or(28) as f32 / 128.0).clamp(0.0, 1.0),
            num_heads_norm: (self.num_heads.unwrap_or(12) as f32 / 128.0).clamp(0.0, 1.0),
            head_dim_norm: (self.head_dim.unwrap_or(128) as f32 / 256.0).clamp(0.0, 1.0),
            vocab_size_log: self
                .vocab_size
                .map(|v| (v as f32).log10() / 6.0) // log10(1M)=6
                .unwrap_or(0.0)
                .clamp(0.0, 1.0),
            batch_size_norm: (batch_size_f / 64.0).clamp(0.0, 1.0),
            seq_len_log: self
                .seq_len
                .map(|s| (s as f32).log2() / 15.0) // log2(32K)≈15
                .unwrap_or(0.0)
                .clamp(0.0, 1.0),
            cuda_graphs: if self.cuda_graphs { 1.0 } else { 0.0 },
            kv_cache_ratio: (kv_caches as f32 / batch_size_f).clamp(0.0, 1.0),
            is_prefill: if self.is_prefill { 1.0 } else { 0.0 },

            // One-hot encodings
            quant_type_onehot: quant_onehot,
            kernel_type_onehot: kernel_onehot,

            // Hardware features (5) [v1.1.0]
            gpu_mem_bw_norm: (self.gpu_mem_bw_gbs.unwrap_or(1000.0) / 3000.0).clamp(0.0, 1.0),
            gpu_compute_norm: (self.gpu_compute_tflops.unwrap_or(100.0) / 500.0).clamp(0.0, 1.0),
            gpu_sm_norm: (self.gpu_sm_count.unwrap_or(128) as f32 / 200.0).clamp(0.0, 1.0),
            gpu_l2_cache_norm: (self.gpu_l2_cache_mb.unwrap_or(48.0) / 128.0).clamp(0.0, 1.0), // v1.1.0
            is_zero_copy: if self.is_zero_copy { 1.0 } else { 0.0 }, // v1.1.0

            // Derived features
            arithmetic_intensity,
            theoretical_efficiency,

            // Labels
            measured_tps: self.measured_tps,
            best_kernel_id: None,
            bottleneck_class: None,
        }
    }
}

// ============================================================================
// FeatureExtractor
// ============================================================================

/// Extracts features from BrickProfiler and runtime configuration.
#[derive(Debug)]
pub struct FeatureExtractor {
    /// Hardware capability (cached)
    pub(crate) hardware: Option<HardwareCapability>,
}

impl Default for FeatureExtractor {
    fn default() -> Self {
        Self::new()
    }
}

impl FeatureExtractor {
    /// Create a new feature extractor
    pub fn new() -> Self {
        Self { hardware: None }
    }

    /// Create with hardware capability
    pub fn with_hardware(hardware: HardwareCapability) -> Self {
        Self {
            hardware: Some(hardware),
        }
    }

    /// Extract features from profiler and configuration
    pub fn extract(&self, profiler: &BrickProfiler, config: &RunConfig) -> TunerFeatures {
        let mut builder = TunerFeatures::builder()
            .model_params_b(config.model_params_b)
            .hidden_dim(config.hidden_dim)
            .num_layers(config.num_layers)
            .num_heads(config.num_heads)
            .batch_size(config.batch_size)
            .seq_len(config.seq_len)
            .cuda_graphs(config.cuda_graphs)
            .quant_type(config.quant_type)
            .kernel_type(config.kernel_type);

        // Add hardware features if available
        if let Some(hw) = &self.hardware {
            builder = builder.hardware(hw);
        }

        // Add measured throughput if available
        if let Some(tps) = profiler.tokens_per_sec() {
            builder = builder.measured_tps(tps);
        }

        let mut features = builder.build();

        // Update derived features from profiler
        if let Some(efficiency) = self.calculate_efficiency(profiler, config) {
            features.theoretical_efficiency = efficiency;
        }

        // Classify bottleneck from profiler data
        features.bottleneck_class = Some(self.classify_bottleneck(profiler));

        features
    }

    /// Calculate efficiency from profiler data
    pub fn calculate_efficiency(&self, profiler: &BrickProfiler, config: &RunConfig) -> Option<f32> {
        let measured_tps = profiler.tokens_per_sec()?;
        let hw = self.hardware.as_ref()?;
        let gpu = hw.gpu.as_ref()?;

        // Calculate theoretical max based on roofline
        let bytes_per_token = config.model_params_b * 1e9 * config.quant_type.bytes_per_param();
        let theoretical_tps = (gpu.memory_bw_gbps as f32) * 1e9 / bytes_per_token;

        Some((measured_tps / theoretical_tps).clamp(0.0, 1.0))
    }

    /// Classify bottleneck from profiler brick breakdown.
    ///
    /// PAR-200: Uses category_stats() for efficient aggregation.
    pub fn classify_bottleneck(&self, profiler: &BrickProfiler) -> BottleneckClass {
        let cats = profiler.category_stats();
        let total_ns = profiler.total_ns();

        if total_ns == 0 {
            return BottleneckClass::Unknown;
        }

        // Get category percentages
        let attention_pct =
            cats[BrickCategory::Attention as usize].percentage(total_ns) as f32 / 100.0;
        let ffn_pct = cats[BrickCategory::Ffn as usize].percentage(total_ns) as f32 / 100.0;
        let norm_pct = cats[BrickCategory::Norm as usize].percentage(total_ns) as f32 / 100.0;

        // Classify based on dominant component
        if attention_pct > 0.35 {
            BottleneckClass::AttentionBound
        } else if ffn_pct > 0.50 {
            // FFN is memory-bound (large GEMV operations)
            BottleneckClass::MemoryBound
        } else if norm_pct > 0.20 {
            // High norm percentage indicates launch overhead
            BottleneckClass::LaunchBound
        } else {
            BottleneckClass::MemoryBound // Default for inference
        }
    }
}

// ============================================================================
// RunConfig
// ============================================================================

/// Runtime configuration for feature extraction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunConfig {
    pub model_params_b: f32,
    pub hidden_dim: u32,
    pub num_layers: u32,
    pub num_heads: u32,
    pub batch_size: u32,
    pub seq_len: u32,
    pub cuda_graphs: bool,
    pub quant_type: QuantType,
    pub kernel_type: KernelType,
}

impl Default for RunConfig {
    fn default() -> Self {
        Self {
            model_params_b: 1.5,
            hidden_dim: 1536,
            num_layers: 28,
            num_heads: 12,
            batch_size: 1,
            seq_len: 1,
            cuda_graphs: false,
            quant_type: QuantType::Q4K,
            kernel_type: KernelType::VectorizedQ4K,
        }
    }
}

Coverage Report

Created: 2026-01-25 15:05

Line	Count	Source
1		//! Feature Extraction for ML Tuning
2		//!
3		//! Implements TunerFeatures, TunerFeaturesBuilder, FeatureExtractor, and RunConfig.
4
5		use crate::brick::{BrickCategory, BrickProfiler};
6		use crate::hardware::HardwareCapability;
7		use serde::{Deserialize, Serialize};
8
9		use super::error::TunerError;
10		use super::types::{BottleneckClass, KernelType, QuantType};
11
12		// ============================================================================
13		// TunerFeatures
14		// ============================================================================
15
16		/// Feature vector for ML-based kernel tuning.
17		///
18		/// All fields normalized to [0, 1] for model input.
19		/// Total dimension: 42 features.
20		///
21		/// # Feature Categories
22		///
23		/// - Static (11): Known before execution (model size, batch size, etc.)
24		/// - Quant one-hot (8): Quantization type encoding
25		/// - Kernel one-hot (16): Kernel type encoding
26		/// - Hardware (5): GPU capabilities
27		/// - Derived (2): Computed features (arithmetic intensity, efficiency)
28		#[derive(Debug, Clone, Serialize, Deserialize)]
29		pub struct TunerFeatures {
30		// === Static features (11) ===
31		/// Model size in billions (log10 normalized)
32		pub model_params_b: f32,
33		/// Hidden dimension / 16384
34		pub hidden_dim_norm: f32,
35		/// Number of layers / 128
36		pub num_layers_norm: f32,
37		/// Number of attention heads / 128
38		pub num_heads_norm: f32,
39		/// Head dimension / 256
40		pub head_dim_norm: f32,
41		/// Vocabulary size (log10 normalized)
42		pub vocab_size_log: f32,
43		/// Batch size M / 64
44		pub batch_size_norm: f32,
45		/// Sequence length (log2 / 15)
46		pub seq_len_log: f32,
47		/// CUDA graphs enabled (0 or 1)
48		pub cuda_graphs: f32,
49		/// Number of KV caches / batch_size (for multi-cache detection)
50		pub kv_cache_ratio: f32,
51		/// Prefill vs decode (0=decode, 1=prefill)
52		pub is_prefill: f32,
53
54		// === Quantization one-hot (8) ===
55		pub quant_type_onehot: [f32; 8],
56
57		// === Kernel one-hot (16) ===
58		pub kernel_type_onehot: [f32; 16],
59
60		// === Hardware features (5) === [v1.1.0: added L2 cache + zero-copy]
61		/// Memory bandwidth / 3000 GB/s
62		pub gpu_mem_bw_norm: f32,
63		/// Compute TFLOPS / 500
64		pub gpu_compute_norm: f32,
65		/// SM count / 200
66		pub gpu_sm_norm: f32,
67		/// L2 cache size / 128 MB (v1.1.0: critical for occupancy)
68		pub gpu_l2_cache_norm: f32,
69		/// Zero-copy memory path enabled (0 or 1) (v1.1.0: pinned memory)
70		pub is_zero_copy: f32,
71
72		// === Derived features (2) ===
73		/// Arithmetic intensity (FLOP/byte), normalized
74		pub arithmetic_intensity: f32,
75		/// Theoretical efficiency (measured / roofline)
76		pub theoretical_efficiency: f32,
77
78		// === Labels (for training) ===
79		/// Measured throughput (tokens/second) - training label
80		#[serde(skip_serializing_if = "Option::is_none")]
81		pub measured_tps: Option<f32>,
82		/// Best kernel ID - classification label
83		#[serde(skip_serializing_if = "Option::is_none")]
84		pub best_kernel_id: Option<u8>,
85		/// Bottleneck class - classification label
86		#[serde(skip_serializing_if = "Option::is_none")]
87		pub bottleneck_class: Option<BottleneckClass>,
88		}
89
90		impl Default for TunerFeatures {
91	0	fn default() -> Self {
92	0	Self {
93	0	model_params_b: 0.0,
94	0	hidden_dim_norm: 0.0,
95	0	num_layers_norm: 0.0,
96	0	num_heads_norm: 0.0,
97	0	head_dim_norm: 0.0,
98	0	vocab_size_log: 0.0,
99	0	batch_size_norm: 0.0,
100	0	seq_len_log: 0.0,
101	0	cuda_graphs: 0.0,
102	0	kv_cache_ratio: 1.0,
103	0	is_prefill: 0.0,
104	0	quant_type_onehot: [0.0; 8],
105	0	kernel_type_onehot: [0.0; 16],
106	0	gpu_mem_bw_norm: 0.0,
107	0	gpu_compute_norm: 0.0,
108	0	gpu_sm_norm: 0.0,
109	0	gpu_l2_cache_norm: 0.0,
110	0	is_zero_copy: 0.0,
111	0	arithmetic_intensity: 0.0,
112	0	theoretical_efficiency: 0.0,
113	0	measured_tps: None,
114	0	best_kernel_id: None,
115	0	bottleneck_class: None,
116	0	}
117	0	}
118		}
119
120		impl TunerFeatures {
121		/// Total feature dimension (excluding labels)
122		/// v1.1.0: 11 static + 8 quant + 16 kernel + 5 hardware + 2 derived = 42
123		pub const DIM: usize = 11 + 8 + 16 + 5 + 2; // 42 features (v1.1.0)
124
125		/// Create a new feature builder
126	0	pub fn builder() -> TunerFeaturesBuilder {
127	0	TunerFeaturesBuilder::default()
128	0	}
129
130		/// Convert to flat vector for model input
131	0	pub fn to_vector(&self) -> Vec<f32> {
132	0	let mut v = Vec::with_capacity(Self::DIM);
133
134		// Static features
135	0	v.push(self.model_params_b);
136	0	v.push(self.hidden_dim_norm);
137	0	v.push(self.num_layers_norm);
138	0	v.push(self.num_heads_norm);
139	0	v.push(self.head_dim_norm);
140	0	v.push(self.vocab_size_log);
141	0	v.push(self.batch_size_norm);
142	0	v.push(self.seq_len_log);
143	0	v.push(self.cuda_graphs);
144	0	v.push(self.kv_cache_ratio);
145	0	v.push(self.is_prefill);
146
147		// One-hot encodings
148	0	v.extend_from_slice(&self.quant_type_onehot);
149	0	v.extend_from_slice(&self.kernel_type_onehot);
150
151		// Hardware features (5) [v1.1.0]
152	0	v.push(self.gpu_mem_bw_norm);
153	0	v.push(self.gpu_compute_norm);
154	0	v.push(self.gpu_sm_norm);
155	0	v.push(self.gpu_l2_cache_norm); // v1.1.0
156	0	v.push(self.is_zero_copy); // v1.1.0
157
158		// Derived features
159	0	v.push(self.arithmetic_intensity);
160	0	v.push(self.theoretical_efficiency);
161
162	0	v
163	0	}
164
165		/// Validate features (F021-F030 falsification criteria)
166	0	pub fn validate(&self) -> Result<(), TunerError> {
167	0	let v = self.to_vector();
168
169		// F021: No NaN features
170	0	if v.iter().any(\|x\| x.is_nan()) {
171	0	return Err(TunerError::InvalidFeature("NaN value in features".into()));
172	0	}
173
174		// F022: No infinite features
175	0	if v.iter().any(\|x\| x.is_infinite()) {
176	0	return Err(TunerError::InvalidFeature(
177	0	"Infinite value in features".into(),
178	0	));
179	0	}
180
181		// F023: All features in [0, 1] (with small tolerance for floating point)
182	0	if v.iter().any(\|x\| x < -0.001 \|\| x > 1.001) {
183	0	return Err(TunerError::InvalidFeature(
184	0	"Feature value outside [0, 1]".into(),
185	0	));
186	0	}
187
188		// F029: One-hot sums = 1
189	0	let quant_sum: f32 = self.quant_type_onehot.iter().sum();
190	0	if (quant_sum - 1.0).abs() > 0.001 && quant_sum > 0.001 {
191	0	return Err(TunerError::InvalidFeature(
192	0	"Quant one-hot does not sum to 1".into(),
193	0	));
194	0	}
195
196	0	let kernel_sum: f32 = self.kernel_type_onehot.iter().sum();
197	0	if (kernel_sum - 1.0).abs() > 0.001 && kernel_sum > 0.001 {
198	0	return Err(TunerError::InvalidFeature(
199	0	"Kernel one-hot does not sum to 1".into(),
200	0	));
201	0	}
202
203	0	Ok(())
204	0	}
205		}
206
207		// ============================================================================
208		// TunerFeaturesBuilder
209		// ============================================================================
210
211		/// Builder for TunerFeatures with automatic normalization.
212		#[derive(Default)]
213		pub struct TunerFeaturesBuilder {
214		model_params_b: Option<f32>,
215		hidden_dim: Option<u32>,
216		num_layers: Option<u32>,
217		num_heads: Option<u32>,
218		head_dim: Option<u32>,
219		vocab_size: Option<u32>,
220		batch_size: Option<u32>,
221		seq_len: Option<u32>,
222		cuda_graphs: bool,
223		kv_caches: Option<u32>,
224		is_prefill: bool,
225		quant_type: Option<QuantType>,
226		kernel_type: Option<KernelType>,
227		gpu_mem_bw_gbs: Option<f32>,
228		gpu_compute_tflops: Option<f32>,
229		gpu_sm_count: Option<u32>,
230		gpu_l2_cache_mb: Option<f32>, // v1.1.0
231		is_zero_copy: bool, // v1.1.0
232		measured_tps: Option<f32>,
233		}
234
235		impl TunerFeaturesBuilder {
236		/// Set model size in billions of parameters
237	0	pub fn model_params_b(mut self, params: f32) -> Self {
238	0	self.model_params_b = Some(params);
239	0	self
240	0	}
241
242		/// Set hidden dimension
243	0	pub fn hidden_dim(mut self, dim: u32) -> Self {
244	0	self.hidden_dim = Some(dim);
245	0	self
246	0	}
247
248		/// Set number of layers
249	0	pub fn num_layers(mut self, layers: u32) -> Self {
250	0	self.num_layers = Some(layers);
251	0	self
252	0	}
253
254		/// Set number of attention heads
255	0	pub fn num_heads(mut self, heads: u32) -> Self {
256	0	self.num_heads = Some(heads);
257	0	self
258	0	}
259
260		/// Set head dimension
261	0	pub fn head_dim(mut self, dim: u32) -> Self {
262	0	self.head_dim = Some(dim);
263	0	self
264	0	}
265
266		/// Set vocabulary size
267	0	pub fn vocab_size(mut self, size: u32) -> Self {
268	0	self.vocab_size = Some(size);
269	0	self
270	0	}
271
272		/// Set batch size (M)
273	0	pub fn batch_size(mut self, m: u32) -> Self {
274	0	self.batch_size = Some(m);
275	0	self
276	0	}
277
278		/// Set sequence length
279	0	pub fn seq_len(mut self, len: u32) -> Self {
280	0	self.seq_len = Some(len);
281	0	self
282	0	}
283
284		/// Enable CUDA graphs
285	0	pub fn cuda_graphs(mut self, enabled: bool) -> Self {
286	0	self.cuda_graphs = enabled;
287	0	self
288	0	}
289
290		/// Set number of KV caches
291	0	pub fn kv_caches(mut self, count: u32) -> Self {
292	0	self.kv_caches = Some(count);
293	0	self
294	0	}
295
296		/// Set prefill mode
297	0	pub fn is_prefill(mut self, prefill: bool) -> Self {
298	0	self.is_prefill = prefill;
299	0	self
300	0	}
301
302		/// Set quantization type
303	0	pub fn quant_type(mut self, qt: QuantType) -> Self {
304	0	self.quant_type = Some(qt);
305	0	self
306	0	}
307
308		/// Set kernel type
309	0	pub fn kernel_type(mut self, kt: KernelType) -> Self {
310	0	self.kernel_type = Some(kt);
311	0	self
312	0	}
313
314		/// Set GPU memory bandwidth in GB/s
315	0	pub fn gpu_mem_bw_gbs(mut self, bw: f32) -> Self {
316	0	self.gpu_mem_bw_gbs = Some(bw);
317	0	self
318	0	}
319
320		/// Set GPU compute in TFLOPS
321	0	pub fn gpu_compute_tflops(mut self, tflops: f32) -> Self {
322	0	self.gpu_compute_tflops = Some(tflops);
323	0	self
324	0	}
325
326		/// Set GPU SM count
327	0	pub fn gpu_sm_count(mut self, count: u32) -> Self {
328	0	self.gpu_sm_count = Some(count);
329	0	self
330	0	}
331
332		/// Set measured throughput (for training data)
333	0	pub fn measured_tps(mut self, tps: f32) -> Self {
334	0	self.measured_tps = Some(tps);
335	0	self
336	0	}
337
338		/// Set L2 cache size in MB (v1.1.0)
339	0	pub fn gpu_l2_cache_mb(mut self, l2_mb: f32) -> Self {
340	0	self.gpu_l2_cache_mb = Some(l2_mb);
341	0	self
342	0	}
343
344		/// Set zero-copy memory path enabled (v1.1.0)
345	0	pub fn is_zero_copy(mut self, enabled: bool) -> Self {
346	0	self.is_zero_copy = enabled;
347	0	self
348	0	}
349
350		/// Set hardware capability (auto-fills GPU features)
351	0	pub fn hardware(mut self, hw: &HardwareCapability) -> Self {
352	0	if let Some(gpu) = &hw.gpu {
353	0	self.gpu_mem_bw_gbs = Some(gpu.memory_bw_gbps as f32);
354	0	self.gpu_compute_tflops = Some(gpu.peak_tflops_fp32 as f32);
355	0	// SM count not directly available; estimate from compute capability
356	0	self.gpu_sm_count = None;
357	0	}
358	0	self
359	0	}
360
361		/// Build the feature vector with normalization
362	0	pub fn build(self) -> TunerFeatures {
363	0	let batch_size = self.batch_size.unwrap_or(1);
364	0	let kv_caches = self.kv_caches.unwrap_or(batch_size);
365
366		// Create one-hot encodings
367	0	let mut quant_onehot = [0.0f32; 8];
368	0	if let Some(qt) = self.quant_type {
369	0	quant_onehot[qt.to_index()] = 1.0;
370	0	}
371
372	0	let mut kernel_onehot = [0.0f32; 16];
373	0	if let Some(kt) = self.kernel_type {
374	0	kernel_onehot[kt.to_index()] = 1.0;
375	0	}
376
377		// Calculate derived features
378	0	let hidden_dim = self.hidden_dim.unwrap_or(1536) as f32;
379	0	let batch_size_f = batch_size as f32;
380	0	let quant_bytes = self
381	0	.quant_type
382	0	.map(\|q\| q.bytes_per_param())
383	0	.unwrap_or(0.5625);
384
385		// Arithmetic intensity for GEMV: 2NK FLOPs / (NKbytes + K + N) bytes
386		// Simplified: ~2 / bytes_per_param for memory-bound inference
387	0	let arithmetic_intensity = (2.0 / quant_bytes).min(10.0) / 10.0;
388
389		// Theoretical efficiency starts at 0 (unknown until measured)
390	0	let theoretical_efficiency = 0.0;
391
392		TunerFeatures {
393		// Normalized static features
394	0	model_params_b: self
395	0	.model_params_b
396	0	.map(\|p\| (p.log10() + 1.0) / 3.0) // log10(0.1)=-1, log10(100)=2 → [0, 1]
397	0	.unwrap_or(0.0)
398	0	.clamp(0.0, 1.0),
399	0	hidden_dim_norm: (hidden_dim / 16384.0).clamp(0.0, 1.0),
400	0	num_layers_norm: (self.num_layers.unwrap_or(28) as f32 / 128.0).clamp(0.0, 1.0),
401	0	num_heads_norm: (self.num_heads.unwrap_or(12) as f32 / 128.0).clamp(0.0, 1.0),
402	0	head_dim_norm: (self.head_dim.unwrap_or(128) as f32 / 256.0).clamp(0.0, 1.0),
403	0	vocab_size_log: self
404	0	.vocab_size
405	0	.map(\|v\| (v as f32).log10() / 6.0) // log10(1M)=6
406	0	.unwrap_or(0.0)
407	0	.clamp(0.0, 1.0),
408	0	batch_size_norm: (batch_size_f / 64.0).clamp(0.0, 1.0),
409	0	seq_len_log: self
410	0	.seq_len
411	0	.map(\|s\| (s as f32).log2() / 15.0) // log2(32K)≈15
412	0	.unwrap_or(0.0)
413	0	.clamp(0.0, 1.0),
414	0	cuda_graphs: if self.cuda_graphs { 1.0 } else { 0.0 },
415	0	kv_cache_ratio: (kv_caches as f32 / batch_size_f).clamp(0.0, 1.0),
416	0	is_prefill: if self.is_prefill { 1.0 } else { 0.0 },
417
418		// One-hot encodings
419	0	quant_type_onehot: quant_onehot,
420	0	kernel_type_onehot: kernel_onehot,
421
422		// Hardware features (5) [v1.1.0]
423	0	gpu_mem_bw_norm: (self.gpu_mem_bw_gbs.unwrap_or(1000.0) / 3000.0).clamp(0.0, 1.0),
424	0	gpu_compute_norm: (self.gpu_compute_tflops.unwrap_or(100.0) / 500.0).clamp(0.0, 1.0),
425	0	gpu_sm_norm: (self.gpu_sm_count.unwrap_or(128) as f32 / 200.0).clamp(0.0, 1.0),
426	0	gpu_l2_cache_norm: (self.gpu_l2_cache_mb.unwrap_or(48.0) / 128.0).clamp(0.0, 1.0), // v1.1.0
427	0	is_zero_copy: if self.is_zero_copy { 1.0 } else { 0.0 }, // v1.1.0
428
429		// Derived features
430	0	arithmetic_intensity,
431	0	theoretical_efficiency,
432
433		// Labels
434	0	measured_tps: self.measured_tps,
435	0	best_kernel_id: None,
436	0	bottleneck_class: None,
437		}
438	0	}
439		}
440
441		// ============================================================================
442		// FeatureExtractor
443		// ============================================================================
444
445		/// Extracts features from BrickProfiler and runtime configuration.
446		#[derive(Debug)]
447		pub struct FeatureExtractor {
448		/// Hardware capability (cached)
449		pub(crate) hardware: Option<HardwareCapability>,
450		}
451
452		impl Default for FeatureExtractor {
453	0	fn default() -> Self {
454	0	Self::new()
455	0	}
456		}
457
458		impl FeatureExtractor {
459		/// Create a new feature extractor
460	0	pub fn new() -> Self {
461	0	Self { hardware: None }
462	0	}
463
464		/// Create with hardware capability
465	0	pub fn with_hardware(hardware: HardwareCapability) -> Self {
466	0	Self {
467	0	hardware: Some(hardware),
468	0	}
469	0	}
470
471		/// Extract features from profiler and configuration
472	0	pub fn extract(&self, profiler: &BrickProfiler, config: &RunConfig) -> TunerFeatures {
473	0	let mut builder = TunerFeatures::builder()
474	0	.model_params_b(config.model_params_b)
475	0	.hidden_dim(config.hidden_dim)
476	0	.num_layers(config.num_layers)
477	0	.num_heads(config.num_heads)
478	0	.batch_size(config.batch_size)
479	0	.seq_len(config.seq_len)
480	0	.cuda_graphs(config.cuda_graphs)
481	0	.quant_type(config.quant_type)
482	0	.kernel_type(config.kernel_type);
483
484		// Add hardware features if available
485	0	if let Some(hw) = &self.hardware {
486	0	builder = builder.hardware(hw);
487	0	}
488
489		// Add measured throughput if available
490	0	if let Some(tps) = profiler.tokens_per_sec() {
491	0	builder = builder.measured_tps(tps);
492	0	}
493
494	0	let mut features = builder.build();
495
496		// Update derived features from profiler
497	0	if let Some(efficiency) = self.calculate_efficiency(profiler, config) {
498	0	features.theoretical_efficiency = efficiency;
499	0	}
500
501		// Classify bottleneck from profiler data
502	0	features.bottleneck_class = Some(self.classify_bottleneck(profiler));
503
504	0	features
505	0	}
506
507		/// Calculate efficiency from profiler data
508	0	pub fn calculate_efficiency(&self, profiler: &BrickProfiler, config: &RunConfig) -> Option<f32> {
509	0	let measured_tps = profiler.tokens_per_sec()?;
510	0	let hw = self.hardware.as_ref()?;
511	0	let gpu = hw.gpu.as_ref()?;
512
513		// Calculate theoretical max based on roofline
514	0	let bytes_per_token = config.model_params_b * 1e9 * config.quant_type.bytes_per_param();
515	0	let theoretical_tps = (gpu.memory_bw_gbps as f32) * 1e9 / bytes_per_token;
516
517	0	Some((measured_tps / theoretical_tps).clamp(0.0, 1.0))
518	0	}
519
520		/// Classify bottleneck from profiler brick breakdown.
521		///
522		/// PAR-200: Uses category_stats() for efficient aggregation.
523	0	pub fn classify_bottleneck(&self, profiler: &BrickProfiler) -> BottleneckClass {
524	0	let cats = profiler.category_stats();
525	0	let total_ns = profiler.total_ns();
526
527	0	if total_ns == 0 {
528	0	return BottleneckClass::Unknown;
529	0	}
530
531		// Get category percentages
532	0	let attention_pct =
533	0	cats[BrickCategory::Attention as usize].percentage(total_ns) as f32 / 100.0;
534	0	let ffn_pct = cats[BrickCategory::Ffn as usize].percentage(total_ns) as f32 / 100.0;
535	0	let norm_pct = cats[BrickCategory::Norm as usize].percentage(total_ns) as f32 / 100.0;
536
537		// Classify based on dominant component
538	0	if attention_pct > 0.35 {
539	0	BottleneckClass::AttentionBound
540	0	} else if ffn_pct > 0.50 {
541		// FFN is memory-bound (large GEMV operations)
542	0	BottleneckClass::MemoryBound
543	0	} else if norm_pct > 0.20 {
544		// High norm percentage indicates launch overhead
545	0	BottleneckClass::LaunchBound
546		} else {
547	0	BottleneckClass::MemoryBound // Default for inference
548		}
549	0	}
550		}
551
552		// ============================================================================
553		// RunConfig
554		// ============================================================================
555
556		/// Runtime configuration for feature extraction
557		#[derive(Debug, Clone, Serialize, Deserialize)]
558		pub struct RunConfig {
559		pub model_params_b: f32,
560		pub hidden_dim: u32,
561		pub num_layers: u32,
562		pub num_heads: u32,
563		pub batch_size: u32,
564		pub seq_len: u32,
565		pub cuda_graphs: bool,
566		pub quant_type: QuantType,
567		pub kernel_type: KernelType,
568		}
569
570		impl Default for RunConfig {
571	0	fn default() -> Self {
572	0	Self {
573	0	model_params_b: 1.5,
574	0	hidden_dim: 1536,
575	0	num_layers: 28,
576	0	num_heads: 12,
577	0	batch_size: 1,
578	0	seq_len: 1,
579	0	cuda_graphs: false,
580	0	quant_type: QuantType::Q4K,
581	0	kernel_type: KernelType::VectorizedQ4K,
582	0	}
583	0	}
584		}