cherenkov/
options.rs

1use anyhow::{Result, ensure};
2use clap::Args;
3use serde::{Deserialize, Deserializer, Serialize, Serializer, de};
4use std::{fmt, str::FromStr};
5
6#[derive(Debug, Clone, Copy, Default, PartialEq)]
7pub enum PoolBudget {
8    #[default]
9    Adaptive,
10    Max,
11    Gb(f64),
12}
13
14impl FromStr for PoolBudget {
15    type Err = String;
16
17    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
18        match s {
19            "adaptive" => return Ok(Self::Adaptive),
20            "max" => return Ok(Self::Max),
21            _ => {}
22        }
23
24        match s.parse::<f64>() {
25            Ok(n) if n.is_finite() && n > 0.0 => Ok(Self::Gb(n)),
26            _ => Err("pool must be a positive number of decimal GB, adaptive, or max".into()),
27        }
28    }
29}
30
31// TOML keeps its compact number-or-mode syntax; resolved settings use one enum.
32impl Serialize for PoolBudget {
33    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
34        match self {
35            Self::Adaptive => serializer.serialize_str("adaptive"),
36            Self::Max => serializer.serialize_str("max"),
37            Self::Gb(n) => serializer.serialize_f64(*n),
38        }
39    }
40}
41
42impl<'de> Deserialize<'de> for PoolBudget {
43    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
44        struct PoolVisitor;
45
46        impl<'de> de::Visitor<'de> for PoolVisitor {
47            type Value = PoolBudget;
48
49            fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50                f.write_str("a positive number of decimal GB, 'adaptive', or 'max'")
51            }
52
53            fn visit_str<E: de::Error>(self, value: &str) -> Result<PoolBudget, E> {
54                value.parse().map_err(E::custom)
55            }
56
57            fn visit_f64<E: de::Error>(self, value: f64) -> Result<PoolBudget, E> {
58                if value.is_finite() && value > 0.0 {
59                    Ok(PoolBudget::Gb(value))
60                } else {
61                    Err(E::invalid_value(de::Unexpected::Float(value), &self))
62                }
63            }
64
65            fn visit_i64<E: de::Error>(self, value: i64) -> Result<PoolBudget, E> {
66                self.visit_f64(value as f64)
67            }
68
69            fn visit_u64<E: de::Error>(self, value: u64) -> Result<PoolBudget, E> {
70                self.visit_f64(value as f64)
71            }
72        }
73
74        deserializer.deserialize_any(PoolVisitor)
75    }
76}
77
78pub(crate) fn positive(s: &str) -> std::result::Result<usize, String> {
79    match s.parse::<usize>() {
80        Ok(n) if n > 0 => Ok(n),
81        _ => Err("must be a positive integer".into()),
82    }
83}
84
85pub(crate) fn cut(s: &str) -> std::result::Result<f32, String> {
86    match s.parse::<f32>() {
87        Ok(n) if n.is_finite() && (0.0..=1.0).contains(&n) => Ok(n),
88        _ => Err("weight must be finite and between 0 and 1".into()),
89    }
90}
91
92/// Measured user choices; developer diagnostics remain environment-only.
93#[derive(Args, Debug, Clone)]
94pub struct Options {
95    #[command(flatten)]
96    pub sampling: crate::sampling::Sampling,
97    /// Expert precision: 4, 3 or 2 bits
98    #[arg(long, default_value_t = 4, value_parser = clap::value_parser!(u32).range(2..=4))]
99    pub experts: u32,
100    /// Mid-step fetch precision; follows --experts (mixed mode needs resident 4-bit)
101    #[arg(long, value_parser = clap::value_parser!(u32).range(2..=4))]
102    pub miss_experts: Option<u32>,
103    /// Skip late weak experts; nonzero makes output non-reproducible
104    #[arg(long, default_value_t = 0.0, hide_default_value = true, value_parser = cut)]
105    pub cut_weak: f32,
106    /// Adaptive MTP drafts, 0..3; 0 unloads the draft head
107    #[arg(long, default_value_t = 2, value_parser = clap::value_parser!(u8).range(0..=3))]
108    pub drafts: u8,
109    /// Wired expert pool in decimal GB, or max (default: adaptive)
110    #[arg(
111        long,
112        default_value = "adaptive",
113        hide_default_value = true,
114        value_name = "N|max|adaptive"
115    )]
116    pub pool_gb: PoolBudget,
117    /// Rebuild the selected low-bit store
118    #[arg(long)]
119    pub repack: bool,
120    /// Server policy; CLI runs retain automatic store construction.
121    #[arg(skip = true)]
122    pub build_missing_store: bool,
123    /// Output token limit
124    #[arg(long, default_value_t = 64, value_parser = positive)]
125    pub max_tokens: usize,
126    /// Context capacity; ~22.5 KB/token comes out of the expert pool
127    #[arg(long, default_value_t = 2048, value_parser = positive)]
128    pub max_ctx: usize,
129    /// Skip the chat template
130    #[arg(long)]
131    pub raw: bool,
132    /// Compare GPU rows with the exact CPU reference (slow)
133    #[arg(long)]
134    pub check: bool,
135    /// Repeat using loaded weights, resetting sequence state
136    #[arg(long, default_value_t = 1, value_parser = positive)]
137    pub repeat: usize,
138    /// Ignore EOS for benchmarking
139    #[arg(long)]
140    pub no_eos: bool,
141}
142
143impl Default for Options {
144    fn default() -> Self {
145        Self {
146            sampling: crate::sampling::Sampling::default(),
147            experts: 4,
148            miss_experts: None,
149            cut_weak: 0.0,
150            drafts: 2,
151            pool_gb: PoolBudget::Adaptive,
152            repack: false,
153            build_missing_store: true,
154            max_tokens: 64,
155            max_ctx: 2048,
156            raw: false,
157            check: false,
158            repeat: 1,
159            no_eos: false,
160        }
161    }
162}
163
164impl Options {
165    pub fn effective_drafts(&self) -> usize {
166        if self.sampling.greedy() {
167            self.drafts as usize
168        } else {
169            0
170        }
171    }
172
173    pub fn miss_bits(&self) -> u32 {
174        self.miss_experts.unwrap_or(self.experts)
175    }
176
177    pub fn validate(&self) -> Result<()> {
178        self.sampling.validate()?;
179        ensure!(
180            (2..=4).contains(&self.experts) && (2..=4).contains(&self.miss_bits()),
181            "expert precision must be 4, 3 or 2"
182        );
183        ensure!(
184            self.experts == 4 || self.miss_bits() == self.experts,
185            "mixed precision requires --experts 4; only one low-bit store is attached at a time"
186        );
187        ensure!(
188            !self.repack || self.miss_bits() < 4,
189            "--repack requires a 2- or 3-bit expert precision"
190        );
191        ensure!(self.drafts <= 3, "--drafts must be 0..3");
192        ensure!(
193            self.cut_weak.is_finite() && (0.0..=1.0).contains(&self.cut_weak),
194            "invalid cut weight"
195        );
196        ensure!(
197            self.max_tokens > 0 && self.max_ctx > 0 && self.repeat > 0,
198            "token limits and repeat must be positive"
199        );
200
201        if let PoolBudget::Gb(n) = self.pool_gb {
202            ensure!(n.is_finite() && n > 0.0, "invalid pool size");
203        }
204
205        Ok(())
206    }
207}
208
209#[cfg(test)]
210#[path = "../tests/unit/options.rs"]
211mod tests;
cherenkov/options.rs

cherenkov/
options.rs
