pub const READING_SPEED_TOK_PER_SEC: f64 = 10.0;
pub const SERIOUS_TOK_PER_SEC: f64 = 10.0;
pub const CRITICAL_TOK_PER_SEC: f64 = 5.0;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ThermalLevel {
Nominal,
Fair,
Serious,
Critical,
}
impl ThermalLevel {
pub fn from_raw(raw: i64) -> Self {
match raw {
1 => Self::Fair,
2 => Self::Serious,
n if n >= 3 => Self::Critical,
_ => Self::Nominal,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
pub enum MemoryPressure {
#[default]
Normal,
Warn,
Critical,
}
impl MemoryPressure {
pub fn from_raw(raw: i64) -> Self {
match raw {
2 => Self::Warn,
4 => Self::Critical,
_ => Self::Normal,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PowerProfile {
#[default]
Performance,
Balanced,
Efficiency,
}
impl PowerProfile {
pub fn parse(text: &str) -> Option<Self> {
match text {
"performance" => Some(Self::Performance),
"balanced" => Some(Self::Balanced),
"efficiency" => Some(Self::Efficiency),
_ => None,
}
}
pub fn as_str(self) -> &'static str {
match self {
Self::Performance => "performance",
Self::Balanced => "balanced",
Self::Efficiency => "efficiency",
}
}
pub fn base_cap(self) -> Option<f64> {
match self {
Self::Efficiency => Some(READING_SPEED_TOK_PER_SEC),
Self::Performance | Self::Balanced => None,
}
}
pub fn thermal_stepping(self) -> bool {
!matches!(self, Self::Performance)
}
}
pub fn thermal_cap(base: Option<f64>, level: ThermalLevel) -> Option<f64> {
let ceiling = match level {
ThermalLevel::Nominal | ThermalLevel::Fair => return base,
ThermalLevel::Serious => SERIOUS_TOK_PER_SEC,
ThermalLevel::Critical => CRITICAL_TOK_PER_SEC,
};
Some(base.map_or(ceiling, |base| base.min(ceiling)))
}
pub fn memory_cap(base: Option<f64>, level: MemoryPressure) -> Option<f64> {
let ceiling = match level {
MemoryPressure::Normal => return base,
MemoryPressure::Warn => SERIOUS_TOK_PER_SEC,
MemoryPressure::Critical => CRITICAL_TOK_PER_SEC,
};
Some(base.map_or(ceiling, |base| base.min(ceiling)))
}
pub fn stepped_cap(
base: Option<f64>,
thermal: ThermalLevel,
memory: MemoryPressure,
) -> Option<f64> {
match (thermal_cap(base, thermal), memory_cap(base, memory)) {
(Some(a), Some(b)) => Some(a.min(b)),
(a, b) => a.or(b),
}
}
#[cfg(target_os = "macos")]
pub fn thermal_level() -> ThermalLevel {
ThermalLevel::from_raw(gpu::thermal_state_raw())
}
#[cfg(target_os = "macos")]
pub fn memory_pressure() -> MemoryPressure {
MemoryPressure::from_raw(gpu::memory_pressure_raw())
}
#[cfg(not(target_os = "macos"))]
pub fn memory_pressure() -> MemoryPressure {
MemoryPressure::Normal
}
#[cfg(not(target_os = "macos"))]
pub fn thermal_level() -> ThermalLevel {
ThermalLevel::Nominal
}
#[cfg(target_os = "macos")]
pub fn low_power_mode_enabled() -> bool {
gpu::low_power_mode_enabled()
}
#[cfg(not(target_os = "macos"))]
pub fn low_power_mode_enabled() -> bool {
false
}
#[cfg(target_os = "macos")]
pub fn physical_memory() -> u64 {
gpu::physical_memory()
}
#[cfg(not(target_os = "macos"))]
pub fn physical_memory() -> u64 {
0
}
#[cfg(target_os = "macos")]
pub fn recommended_max_working_set() -> Option<(u64, String)> {
gpu::recommended_max_working_set()
}
#[cfg(not(target_os = "macos"))]
pub fn recommended_max_working_set() -> Option<(u64, String)> {
None
}
pub fn resolve_profile(explicit: Option<PowerProfile>) -> PowerProfile {
match explicit {
Some(profile) => profile,
None if low_power_mode_enabled() => PowerProfile::Efficiency,
None => PowerProfile::Performance,
}
}
pub fn rate_control_for(profile: PowerProfile, explicit_cap: Option<f64>) -> RateControl {
RateControl {
max_tokens_per_sec: explicit_cap.or_else(|| profile.base_cap()),
thermal_probe: profile
.thermal_stepping()
.then_some(thermal_level as fn() -> ThermalLevel),
memory_probe: profile
.thermal_stepping()
.then_some(memory_pressure as fn() -> MemoryPressure),
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct RateControl {
pub max_tokens_per_sec: Option<f64>,
pub thermal_probe: Option<fn() -> ThermalLevel>,
pub memory_probe: Option<fn() -> MemoryPressure>,
}
impl RateControl {
pub fn is_active(&self) -> bool {
self.max_tokens_per_sec.is_some()
|| self.thermal_probe.is_some()
|| self.memory_probe.is_some()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn raw_thermal_levels_map_and_clamp_toward_the_safe_side() {
assert_eq!(ThermalLevel::from_raw(0), ThermalLevel::Nominal);
assert_eq!(ThermalLevel::from_raw(1), ThermalLevel::Fair);
assert_eq!(ThermalLevel::from_raw(2), ThermalLevel::Serious);
assert_eq!(ThermalLevel::from_raw(3), ThermalLevel::Critical);
assert_eq!(ThermalLevel::from_raw(7), ThermalLevel::Critical);
assert_eq!(ThermalLevel::from_raw(-1), ThermalLevel::Nominal);
}
#[test]
fn profiles_round_trip_through_their_spelling() {
for profile in [
PowerProfile::Performance,
PowerProfile::Balanced,
PowerProfile::Efficiency,
] {
assert_eq!(PowerProfile::parse(profile.as_str()), Some(profile));
}
assert_eq!(PowerProfile::parse("turbo"), None);
assert_eq!(PowerProfile::parse("Performance"), None);
assert_eq!(PowerProfile::default(), PowerProfile::Performance);
}
#[test]
fn only_efficiency_carries_a_base_cap_and_only_performance_skips_stepping() {
assert_eq!(PowerProfile::Performance.base_cap(), None);
assert_eq!(PowerProfile::Balanced.base_cap(), None);
assert_eq!(
PowerProfile::Efficiency.base_cap(),
Some(READING_SPEED_TOK_PER_SEC)
);
assert!(!PowerProfile::Performance.thermal_stepping());
assert!(PowerProfile::Balanced.thermal_stepping());
assert!(PowerProfile::Efficiency.thermal_stepping());
}
#[test]
fn the_thermal_ladder_only_ever_lowers_the_cap() {
assert_eq!(thermal_cap(None, ThermalLevel::Nominal), None);
assert_eq!(thermal_cap(None, ThermalLevel::Fair), None);
assert_eq!(thermal_cap(None, ThermalLevel::Serious), Some(10.0));
assert_eq!(thermal_cap(None, ThermalLevel::Critical), Some(5.0));
assert_eq!(thermal_cap(Some(20.0), ThermalLevel::Fair), Some(20.0));
assert_eq!(thermal_cap(Some(20.0), ThermalLevel::Serious), Some(10.0));
assert_eq!(thermal_cap(Some(20.0), ThermalLevel::Critical), Some(5.0));
assert_eq!(thermal_cap(Some(4.0), ThermalLevel::Serious), Some(4.0));
assert_eq!(thermal_cap(Some(4.0), ThermalLevel::Critical), Some(4.0));
}
#[test]
fn raw_memory_levels_map_the_kernels_bitmask_and_default_to_normal() {
assert_eq!(MemoryPressure::from_raw(1), MemoryPressure::Normal);
assert_eq!(MemoryPressure::from_raw(2), MemoryPressure::Warn);
assert_eq!(MemoryPressure::from_raw(4), MemoryPressure::Critical);
assert_eq!(MemoryPressure::from_raw(0), MemoryPressure::Normal);
assert_eq!(MemoryPressure::from_raw(3), MemoryPressure::Normal);
assert_eq!(MemoryPressure::from_raw(-1), MemoryPressure::Normal);
assert_eq!(MemoryPressure::from_raw(99), MemoryPressure::Normal);
assert_eq!(MemoryPressure::default(), MemoryPressure::Normal);
}
#[test]
fn the_memory_ladder_only_ever_lowers_the_cap() {
assert_eq!(memory_cap(None, MemoryPressure::Normal), None);
assert_eq!(memory_cap(None, MemoryPressure::Warn), Some(10.0));
assert_eq!(memory_cap(None, MemoryPressure::Critical), Some(5.0));
assert_eq!(memory_cap(Some(20.0), MemoryPressure::Warn), Some(10.0));
assert_eq!(memory_cap(Some(4.0), MemoryPressure::Critical), Some(4.0));
}
#[test]
fn the_combined_ladder_takes_whichever_signal_binds_harder() {
assert_eq!(
stepped_cap(None, ThermalLevel::Nominal, MemoryPressure::Normal),
None
);
assert_eq!(
stepped_cap(None, ThermalLevel::Nominal, MemoryPressure::Critical),
Some(5.0)
);
assert_eq!(
stepped_cap(None, ThermalLevel::Critical, MemoryPressure::Normal),
Some(5.0)
);
assert_eq!(
stepped_cap(None, ThermalLevel::Serious, MemoryPressure::Critical),
Some(5.0)
);
assert_eq!(
stepped_cap(None, ThermalLevel::Critical, MemoryPressure::Warn),
Some(5.0)
);
assert_eq!(
stepped_cap(Some(2.0), ThermalLevel::Critical, MemoryPressure::Critical),
Some(2.0)
);
}
#[test]
fn rate_control_pairs_the_cap_with_the_profiles_stepping() {
let performance = rate_control_for(PowerProfile::Performance, None);
assert_eq!(performance.max_tokens_per_sec, None);
assert!(performance.thermal_probe.is_none());
assert!(performance.memory_probe.is_none());
assert!(!performance.is_active());
let balanced = rate_control_for(PowerProfile::Balanced, None);
assert_eq!(balanced.max_tokens_per_sec, None);
assert!(balanced.thermal_probe.is_some());
assert!(balanced.memory_probe.is_some());
assert!(balanced.is_active());
let efficiency = rate_control_for(PowerProfile::Efficiency, None);
assert_eq!(
efficiency.max_tokens_per_sec,
Some(READING_SPEED_TOK_PER_SEC)
);
assert!(efficiency.thermal_probe.is_some());
assert!(efficiency.memory_probe.is_some());
let overridden = rate_control_for(PowerProfile::Efficiency, Some(3.0));
assert_eq!(overridden.max_tokens_per_sec, Some(3.0));
let capped_performance = rate_control_for(PowerProfile::Performance, Some(8.0));
assert_eq!(capped_performance.max_tokens_per_sec, Some(8.0));
assert!(capped_performance.thermal_probe.is_none());
assert!(capped_performance.memory_probe.is_none());
assert!(capped_performance.is_active());
}
#[test]
fn an_explicit_profile_wins_over_the_low_power_mode_default() {
assert_eq!(
resolve_profile(Some(PowerProfile::Performance)),
PowerProfile::Performance
);
assert_eq!(
resolve_profile(Some(PowerProfile::Efficiency)),
PowerProfile::Efficiency
);
let expected = if low_power_mode_enabled() {
PowerProfile::Efficiency
} else {
PowerProfile::Performance
};
assert_eq!(resolve_profile(None), expected);
}
}