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

Source
//! Tuner Type Definitions
//!
//! Core enums for quantization, kernel selection, and bottleneck classification.

use crate::brick::BrickBottleneck;
use serde::{Deserialize, Serialize};

// ============================================================================
// QuantType
// ============================================================================

/// Quantization type for feature encoding.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
pub enum QuantType {
    Q4_0,
    Q4_1,
    #[default]
    Q4K,
    Q5K,
    Q6K,
    Q8_0,
    F16,
    F32,
}

impl QuantType {
    /// One-hot encoding index (0-7)
    pub fn to_index(self) -> usize {
        match self {
            QuantType::Q4_0 => 0,
            QuantType::Q4_1 => 1,
            QuantType::Q4K => 2,
            QuantType::Q5K => 3,
            QuantType::Q6K => 4,
            QuantType::Q8_0 => 5,
            QuantType::F16 => 6,
            QuantType::F32 => 7,
        }
    }

    /// Bytes per parameter (approximate)
    pub fn bytes_per_param(self) -> f32 {
        match self {
            QuantType::Q4_0 | QuantType::Q4_1 | QuantType::Q4K => 0.5625, // 4.5 bits
            QuantType::Q5K => 0.6875,                                     // 5.5 bits
            QuantType::Q6K => 0.8125,                                     // 6.5 bits
            QuantType::Q8_0 => 1.0,
            QuantType::F16 => 2.0,
            QuantType::F32 => 4.0,
        }
    }
}

// ============================================================================
// KernelType
// ============================================================================

/// Kernel type for feature encoding.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
pub enum KernelType {
    // Q4K variants
    #[default]
    TiledQ4K,
    CoalescedQ4K,
    VectorizedQ4K,
    BatchedQ4K,
    Dp4aQ4K,
    FusedRmsNormQ4K,
    // Q6K variants
    CoalescedQ6K,
    // Attention variants
    IncrementalAttention,
    MultiWarpAttention,
    BatchedAttention,
    // Normalization
    RmsNorm,
    VectorizedRmsNorm,
    BatchedRmsNorm,
    // Other
    Generic,
    Unknown,
}

impl KernelType {
    /// One-hot encoding index (0-15)
    pub fn to_index(self) -> usize {
        match self {
            KernelType::TiledQ4K => 0,
            KernelType::CoalescedQ4K => 1,
            KernelType::VectorizedQ4K => 2,
            KernelType::BatchedQ4K => 3,
            KernelType::Dp4aQ4K => 4,
            KernelType::FusedRmsNormQ4K => 5,
            KernelType::CoalescedQ6K => 6,
            KernelType::IncrementalAttention => 7,
            KernelType::MultiWarpAttention => 8,
            KernelType::BatchedAttention => 9,
            KernelType::RmsNorm => 10,
            KernelType::VectorizedRmsNorm => 11,
            KernelType::BatchedRmsNorm => 12,
            KernelType::Generic => 13,
            KernelType::Unknown => 14,
        }
    }

    /// Convert kernel index to type (inverse of to_index())
    pub fn from_index(idx: usize) -> Self {
        match idx {
            0 => KernelType::TiledQ4K,
            1 => KernelType::CoalescedQ4K,
            2 => KernelType::VectorizedQ4K,
            3 => KernelType::BatchedQ4K,
            4 => KernelType::Dp4aQ4K,
            5 => KernelType::FusedRmsNormQ4K,
            6 => KernelType::CoalescedQ6K,
            7 => KernelType::IncrementalAttention,
            8 => KernelType::MultiWarpAttention,
            9 => KernelType::BatchedAttention,
            10 => KernelType::RmsNorm,
            11 => KernelType::VectorizedRmsNorm,
            12 => KernelType::BatchedRmsNorm,
            13 => KernelType::Generic,
            _ => KernelType::Unknown,
        }
    }

    /// Number of kernel types
    pub const COUNT: usize = 16;
}

// ============================================================================
// BottleneckClass
// ============================================================================

/// Bottleneck classification for ML model.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
pub enum BottleneckClass {
    #[default]
    Unknown,
    MemoryBound,
    ComputeBound,
    LaunchBound,
    AttentionBound,
}

impl BottleneckClass {
    /// Convert from BrickBottleneck
    pub fn from_brick_bottleneck(b: BrickBottleneck) -> Self {
        match b {
            BrickBottleneck::Memory => BottleneckClass::MemoryBound,
            BrickBottleneck::Compute => BottleneckClass::ComputeBound,
            BrickBottleneck::Unknown => BottleneckClass::Unknown,
        }
    }

    /// Recommended action for this bottleneck
    pub fn recommended_action(self) -> &'static str {
        match self {
            BottleneckClass::MemoryBound => {
                "Increase batch size (M) to amortize weight reads across sequences"
            }
            BottleneckClass::ComputeBound => {
                "Rare for inference; check for redundant computation or use tensor cores"
            }
            BottleneckClass::LaunchBound => {
                "Enable CUDA graphs or fuse kernels to reduce launch overhead"
            }
            BottleneckClass::AttentionBound => {
                "Use Flash Decoding, reduce sequence length, or use batched attention"
            }
            BottleneckClass::Unknown => "Run profiling to identify bottleneck",
        }
    }

    /// One-hot encoding index (0-4)
    pub fn to_index(self) -> usize {
        match self {
            BottleneckClass::Unknown => 0,
            BottleneckClass::MemoryBound => 1,
            BottleneckClass::ComputeBound => 2,
            BottleneckClass::LaunchBound => 3,
            BottleneckClass::AttentionBound => 4,
        }
    }
}

impl std::fmt::Display for BottleneckClass {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BottleneckClass::Unknown => write!(f, "Unknown"),
            BottleneckClass::MemoryBound => write!(f, "MemoryBound"),
            BottleneckClass::ComputeBound => write!(f, "ComputeBound"),
            BottleneckClass::LaunchBound => write!(f, "LaunchBound"),
            BottleneckClass::AttentionBound => write!(f, "AttentionBound"),
        }
    }
}

Line	Count	Source
1		//! Tuner Type Definitions
2		//!
3		//! Core enums for quantization, kernel selection, and bottleneck classification.
4
5		use crate::brick::BrickBottleneck;
6		use serde::{Deserialize, Serialize};
7
8		// ============================================================================
9		// QuantType
10		// ============================================================================
11
12		/// Quantization type for feature encoding.
13		#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
14		pub enum QuantType {
15		Q4_0,
16		Q4_1,
17		#[default]
18		Q4K,
19		Q5K,
20		Q6K,
21		Q8_0,
22		F16,
23		F32,
24		}
25
26		impl QuantType {
27		/// One-hot encoding index (0-7)
28	0	pub fn to_index(self) -> usize {
29	0	match self {
30	0	QuantType::Q4_0 => 0,
31	0	QuantType::Q4_1 => 1,
32	0	QuantType::Q4K => 2,
33	0	QuantType::Q5K => 3,
34	0	QuantType::Q6K => 4,
35	0	QuantType::Q8_0 => 5,
36	0	QuantType::F16 => 6,
37	0	QuantType::F32 => 7,
38		}
39	0	}
40
41		/// Bytes per parameter (approximate)
42	0	pub fn bytes_per_param(self) -> f32 {
43	0	match self {
44	0	QuantType::Q4_0 \| QuantType::Q4_1 \| QuantType::Q4K => 0.5625, // 4.5 bits
45	0	QuantType::Q5K => 0.6875, // 5.5 bits
46	0	QuantType::Q6K => 0.8125, // 6.5 bits
47	0	QuantType::Q8_0 => 1.0,
48	0	QuantType::F16 => 2.0,
49	0	QuantType::F32 => 4.0,
50		}
51	0	}
52		}
53
54		// ============================================================================
55		// KernelType
56		// ============================================================================
57
58		/// Kernel type for feature encoding.
59		#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
60		pub enum KernelType {
61		// Q4K variants
62		#[default]
63		TiledQ4K,
64		CoalescedQ4K,
65		VectorizedQ4K,
66		BatchedQ4K,
67		Dp4aQ4K,
68		FusedRmsNormQ4K,
69		// Q6K variants
70		CoalescedQ6K,
71		// Attention variants
72		IncrementalAttention,
73		MultiWarpAttention,
74		BatchedAttention,
75		// Normalization
76		RmsNorm,
77		VectorizedRmsNorm,
78		BatchedRmsNorm,
79		// Other
80		Generic,
81		Unknown,
82		}
83
84		impl KernelType {
85		/// One-hot encoding index (0-15)
86	0	pub fn to_index(self) -> usize {
87	0	match self {
88	0	KernelType::TiledQ4K => 0,
89	0	KernelType::CoalescedQ4K => 1,
90	0	KernelType::VectorizedQ4K => 2,
91	0	KernelType::BatchedQ4K => 3,
92	0	KernelType::Dp4aQ4K => 4,
93	0	KernelType::FusedRmsNormQ4K => 5,
94	0	KernelType::CoalescedQ6K => 6,
95	0	KernelType::IncrementalAttention => 7,
96	0	KernelType::MultiWarpAttention => 8,
97	0	KernelType::BatchedAttention => 9,
98	0	KernelType::RmsNorm => 10,
99	0	KernelType::VectorizedRmsNorm => 11,
100	0	KernelType::BatchedRmsNorm => 12,
101	0	KernelType::Generic => 13,
102	0	KernelType::Unknown => 14,
103		}
104	0	}
105
106		/// Convert kernel index to type (inverse of to_index())
107	0	pub fn from_index(idx: usize) -> Self {
108	0	match idx {
109	0	0 => KernelType::TiledQ4K,
110	0	1 => KernelType::CoalescedQ4K,
111	0	2 => KernelType::VectorizedQ4K,
112	0	3 => KernelType::BatchedQ4K,
113	0	4 => KernelType::Dp4aQ4K,
114	0	5 => KernelType::FusedRmsNormQ4K,
115	0	6 => KernelType::CoalescedQ6K,
116	0	7 => KernelType::IncrementalAttention,
117	0	8 => KernelType::MultiWarpAttention,
118	0	9 => KernelType::BatchedAttention,
119	0	10 => KernelType::RmsNorm,
120	0	11 => KernelType::VectorizedRmsNorm,
121	0	12 => KernelType::BatchedRmsNorm,
122	0	13 => KernelType::Generic,
123	0	_ => KernelType::Unknown,
124		}
125	0	}
126
127		/// Number of kernel types
128		pub const COUNT: usize = 16;
129		}
130
131		// ============================================================================
132		// BottleneckClass
133		// ============================================================================
134
135		/// Bottleneck classification for ML model.
136		#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
137		pub enum BottleneckClass {
138		#[default]
139		Unknown,
140		MemoryBound,
141		ComputeBound,
142		LaunchBound,
143		AttentionBound,
144		}
145
146		impl BottleneckClass {
147		/// Convert from BrickBottleneck
148	0	pub fn from_brick_bottleneck(b: BrickBottleneck) -> Self {
149	0	match b {
150	0	BrickBottleneck::Memory => BottleneckClass::MemoryBound,
151	0	BrickBottleneck::Compute => BottleneckClass::ComputeBound,
152	0	BrickBottleneck::Unknown => BottleneckClass::Unknown,
153		}
154	0	}
155
156		/// Recommended action for this bottleneck
157	0	pub fn recommended_action(self) -> &'static str {
158	0	match self {
159		BottleneckClass::MemoryBound => {
160	0	"Increase batch size (M) to amortize weight reads across sequences"
161		}
162		BottleneckClass::ComputeBound => {
163	0	"Rare for inference; check for redundant computation or use tensor cores"
164		}
165		BottleneckClass::LaunchBound => {
166	0	"Enable CUDA graphs or fuse kernels to reduce launch overhead"
167		}
168		BottleneckClass::AttentionBound => {
169	0	"Use Flash Decoding, reduce sequence length, or use batched attention"
170		}
171	0	BottleneckClass::Unknown => "Run profiling to identify bottleneck",
172		}
173	0	}
174
175		/// One-hot encoding index (0-4)
176	0	pub fn to_index(self) -> usize {
177	0	match self {
178	0	BottleneckClass::Unknown => 0,
179	0	BottleneckClass::MemoryBound => 1,
180	0	BottleneckClass::ComputeBound => 2,
181	0	BottleneckClass::LaunchBound => 3,
182	0	BottleneckClass::AttentionBound => 4,
183		}
184	0	}
185		}
186
187		impl std::fmt::Display for BottleneckClass {
188	0	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189	0	match self {
190	0	BottleneckClass::Unknown => write!(f, "Unknown"),
191	0	BottleneckClass::MemoryBound => write!(f, "MemoryBound"),
192	0	BottleneckClass::ComputeBound => write!(f, "ComputeBound"),
193	0	BottleneckClass::LaunchBound => write!(f, "LaunchBound"),
194	0	BottleneckClass::AttentionBound => write!(f, "AttentionBound"),
195		}
196	0	}
197		}

Coverage Report

Created: 2026-01-25 15:05