1use crate::{
2 bench::append,
3 capture, metrics,
4 report::{median, text},
5 util,
6};
7use anyhow::{Context, Result, ensure};
8use clap::Args;
9use serde_json::{Value, json};
10use std::{
11 collections::HashSet,
12 fmt::Write as _,
13 fs,
14 path::{Path, PathBuf},
15 time::Duration,
16};
17
18#[derive(Args)]
19pub struct Options {
20 #[arg(long)]
21 pub before: PathBuf,
22 #[arg(long)]
23 pub after: PathBuf,
24 #[arg(long)]
25 pub model: PathBuf,
26 #[arg(long)]
27 pub out: PathBuf,
28 #[arg(long, default_value_t = 2)]
29 pub rounds: usize,
30 #[arg(long,num_args=1..,default_values=["4","3","2"],value_parser=["4","3","2","4/3","4/2"])]
31 pub configs: Vec<String>,
32}
33
34fn precision(row: &Value) -> (u64, u64) {
35 let bits = row["bits"].as_u64().unwrap_or(4);
36
37 (bits, row["miss_bits"].as_u64().unwrap_or(bits))
38}
39
40pub fn write_summary(out: &Path, report: &Value) -> Result<()> {
41 let runs: Vec<_> = report["runs"]
42 .as_array()
43 .context("runs")?
44 .iter()
45 .filter(|r| r["valid"] == true)
46 .collect();
47 let mut configurations = Vec::new();
48
49 for r in &runs {
50 if !configurations.contains(&precision(r)) {
51 configurations.push(precision(r));
52 }
53 }
54
55 let mut s = format!(
56 "# Low-bit prefill comparison\n\nThe report contains {} valid samples. The medians exclude loading and conversion.\n\n| Prompt tokens | Resident / miss bits | Before pp/s | After pp/s | Change | Pairs |\n| ---: | ---: | ---: | ---: | ---: | ---: |\n",
57 runs.len()
58 );
59 let mut prompts: Vec<_> = report["prompts"]
60 .as_object()
61 .context("prompts")?
62 .keys()
63 .collect();
64
65 prompts.sort_by_key(|p| p.parse::<usize>().unwrap_or(usize::MAX));
66
67 for prompt in prompts {
68 for &(bits, miss) in &configurations {
69 let selected: Vec<_> = runs
70 .iter()
71 .filter(|r| r["prompt"] == *prompt && precision(r) == (bits, miss))
72 .copied()
73 .collect();
74
75 append_median(&mut s, &selected, bits, miss)?;
76 }
77 }
78
79 let controls: Vec<_> = runs
80 .iter()
81 .filter(|r| precision(r) == (4, 4) && r["binary"] == "before")
82 .collect();
83
84 if !controls.is_empty() {
85 s.push_str("\n## Output checks\n\n");
86 }
87
88 for before in controls {
89 let after = runs.iter().find(|r| {
90 r["round"] == before["round"]
91 && r["prompt"] == before["prompt"]
92 && precision(r) == (4, 4)
93 && r["binary"] == "after"
94 });
95 let Some(after) = after else {
96 continue;
97 };
98 let same = fs::read(out.join(text(&before["output"])))?
99 == fs::read(out.join(text(&after["output"])))?;
100
101 writeln!(
102 s,
103 "- Round {}, {} tokens: the Q4 outputs {}.",
104 before["round"],
105 before["metrics"]["prompt_tokens"],
106 if same { "are identical" } else { "differ" }
107 )?;
108 }
109
110 s.push_str(
111 "\nBinary hashes, timings, source and power readings are in [report.json](report.json).\n",
112 );
113
114 if out.join("README.md").is_file() {
115 s.push_str("\n[Run observations](README.md).\n");
116 }
117
118 fs::write(out.join("summary.md"), s)?;
119
120 Ok(())
121}
122
123fn append_median(s: &mut String, runs: &[&Value], bits: u64, miss: u64) -> Result<()> {
124 let rates = |label: &str| {
125 runs.iter()
126 .filter(|r| r["binary"] == label)
127 .filter_map(|r| r["metrics"]["pp_s"].as_f64())
128 .collect::<Vec<_>>()
129 };
130 let (a, b) = (rates("before"), rates("after"));
131
132 if a.is_empty() || b.is_empty() {
133 return Ok(());
134 }
135
136 let pairs = a.len().min(b.len());
137 let (a, b) = (median(a), median(b));
138 let label = if bits == miss {
139 bits.to_string()
140 } else {
141 format!("{bits} / {miss}")
142 };
143
144 writeln!(
145 s,
146 "| {} | {label} | {a:.1} | {b:.1} | {:+.1}% | {pairs} |",
147 runs[0]["metrics"]["prompt_tokens"],
148 100.0 * (b / a - 1.0)
149 )?;
150
151 Ok(())
152}
153
154struct Sample<'a> {
155 options: &'a Options,
156 binary: &'a Path,
157 label: &'a str,
158 round: usize,
159 length: usize,
160 config: &'a str,
161 prompt: &'a str,
162}
163
164impl Sample<'_> {
165 fn run(&self) -> Result<Value> {
166 let (bits, miss) = self
167 .config
168 .split_once('/')
169 .unwrap_or((self.config, self.config));
170 let name = format!(
171 "r{}-repeat{}-q{}-{}",
172 self.round + 1,
173 self.length,
174 self.config.replace('/', "-miss"),
175 self.label
176 );
177 let args = vec![
178 self.binary.display().to_string(),
179 self.options.model.canonicalize()?.display().to_string(),
180 self.prompt.into(),
181 "--experts".into(),
182 bits.into(),
183 "--miss-experts".into(),
184 miss.into(),
185 "--max-ctx".into(),
186 "8192".into(),
187 "--max-tokens".into(),
188 "8".into(),
189 "--no-eos".into(),
190 ];
191 let c = capture::run(
192 &args,
193 &self.options.out.join(format!("{name}.txt")),
194 false,
195 Some(Duration::from_secs(300)),
196 )?;
197 let mut row = json!({
198 "id": name,
199 "round": self.round + 1,
200 "binary": self.label,
201 "bits": bits.parse::<u8>()?,
202 "miss_bits": miss.parse::<u8>()?,
203 "prompt": self.length.to_string(),
204 "args": args,
205 "wall_seconds": c.wall_seconds,
206 "power_before": c.power_before,
207 "power_after": c.power_after,
208 "power_samples": c.power_samples,
209 "exit_code": c.code,
210 "stderr": c.stderr,
211 "output": format!("{name}.txt"),
212 "valid": false,
213 });
214
215 if let Err(error) = validate(&c, &mut row) {
216 row["error"] = json!(error.to_string());
217 }
218
219 eprintln!(
220 "{name}: {} pp/s, valid={}",
221 row["metrics"]["pp_s"], row["valid"]
222 );
223
224 Ok(row)
225 }
226}
227
228fn validate(c: &capture::Captured, row: &mut Value) -> Result<()> {
229 ensure!(
230 c.code == 0 && !c.timed_out && c.cycle.is_none(),
231 "engine failed, timed out, or cycled"
232 );
233
234 let m = metrics::parse(&c.stderr)?;
235 row["metrics"] = m.clone();
236
237 ensure!(m["store_build_seconds"].is_null(), "sample built a store");
238 ensure!(
239 c.stable_power() && c.power_after["source"] == "ac",
240 "power changed"
241 );
242 ensure!(
243 m["metal_gb"].as_f64().unwrap() <= 25.0,
244 "reported Metal allocations exceeded 25 GB"
245 );
246 ensure!(
247 m["output_tokens"] == 8,
248 "decode handoff did not emit eight tokens"
249 );
250
251 row["valid"] = json!(true);
252
253 Ok(())
254}
255
256pub fn run(options: Options) -> Result<()> {
257 ensure!(options.rounds > 0, "rounds must be positive");
258 ensure!(
259 options.configs.iter().collect::<HashSet<_>>().len() == options.configs.len(),
260 "duplicate configurations"
261 );
262 fs::create_dir_all(&options.out)?;
263 ensure!(
264 !options.out.join("report.json").exists(),
265 "output already contains a report"
266 );
267
268 let before = options.before.canonicalize()?;
269 let after = options.after.canonicalize()?;
270 let mut prompts = json!({});
271 let sentence = "The quick brown fox jumps over the lazy dog. Each sentence is part of a repeated document used to measure prompt processing.\n";
272
273 for n in [4, 20, 128] {
274 prompts[n.to_string()] =
275 json!(sentence.repeat(n) + "\nSummarize this document in three sentences.");
276 }
277
278 let mut report = json!({
279 "binaries": {
280 "before": {"path": before, "sha256": util::digest(&before)?},
281 "after": {"path": after, "sha256": util::digest(&after)?},
282 },
283 "source_commit": util::output(&["git", "rev-parse", "HEAD"])?,
284 "source_diff": util::output(&["git", "diff", "HEAD"])?,
285 "prompts": prompts,
286 "configurations": options.configs,
287 "model": options.model.canonicalize()?,
288 "runs": [],
289 });
290
291 for round in 0..options.rounds {
292 for length in [4, 20, 128] {
293 for index in 0..options.configs.len() {
294 let config = &options.configs[(index + round) % options.configs.len()];
295 let order = if round % 2 == 0 {
296 [("before", &before), ("after", &after)]
297 } else {
298 [("after", &after), ("before", &before)]
299 };
300
301 for (label, binary) in order {
302 let sample = Sample {
303 options: &options,
304 binary,
305 label,
306 round,
307 length,
308 config,
309 prompt: prompts[&length.to_string()].as_str().unwrap(),
310 };
311 let row = sample.run()?;
312 let valid = row["valid"] == true;
313 let error = row["error"].clone();
314
315 append(&mut report, "runs", row);
316 util::write_json(&options.out.join("report.json"), &report)?;
317 write_summary(&options.out, &report)?;
318 ensure!(valid, "{error}; details saved in {}", options.out.display());
319 }
320 }
321 }
322 }
323
324 Ok(())
325}