1use std::collections::HashMap;
17
18#[derive(Debug, Clone, Copy, PartialEq)]
22pub enum CalibMethod {
23 MinMax,
25 Percentile(f32), Aciq,
29 Mse,
31}
32
33#[derive(Debug, Clone)]
37pub struct LayerCalibStats {
38 pub layer_name: String,
39 pub num_samples: usize,
40 pub running_min: f32,
41 pub running_max: f32,
42 pub running_mean: f32,
43 pub running_var: f32, histogram: Vec<u64>,
46 histogram_max: f32,
47}
48
49impl LayerCalibStats {
50 pub fn new(layer_name: impl Into<String>) -> Self {
52 Self {
53 layer_name: layer_name.into(),
54 num_samples: 0,
55 running_min: f32::INFINITY,
56 running_max: f32::NEG_INFINITY,
57 running_mean: 0.0,
58 running_var: 0.0,
59 histogram: vec![0u64; 256],
60 histogram_max: 0.0,
61 }
62 }
63
64 pub fn update(&mut self, activations: &[f32]) {
68 for &x in activations {
69 if !x.is_finite() {
70 continue;
71 }
72
73 if x < self.running_min {
75 self.running_min = x;
76 }
77 if x > self.running_max {
78 self.running_max = x;
79 }
80
81 self.num_samples += 1;
83 let n = self.num_samples as f32;
84 let delta = x - self.running_mean;
85 self.running_mean += delta / n;
86 let delta2 = x - self.running_mean;
87 self.running_var += delta * delta2;
88
89 let abs_x = x.abs();
91 if abs_x > self.histogram_max {
92 self.rebuild_histogram_max(abs_x);
94 }
95 }
96
97 for &x in activations {
99 if !x.is_finite() {
100 continue;
101 }
102 let abs_x = x.abs();
103 self.insert_histogram(abs_x);
104 }
105 }
106
107 fn rebuild_histogram_max(&mut self, new_max: f32) {
109 if self.histogram_max > 0.0 && new_max > self.histogram_max {
111 let scale = self.histogram_max / new_max;
112 let mut new_hist = vec![0u64; 256];
113 for (old_bin, &count) in self.histogram.iter().enumerate() {
114 if count == 0 {
115 continue;
116 }
117 let old_frac = (old_bin as f32 + 0.5) / 256.0;
119 let new_frac = old_frac * scale;
120 let new_bin = (new_frac * 256.0) as usize;
121 let new_bin = new_bin.min(255);
122 new_hist[new_bin] += count;
123 }
124 self.histogram = new_hist;
125 }
126 self.histogram_max = new_max;
127 }
128
129 fn insert_histogram(&mut self, abs_x: f32) {
131 if self.histogram_max <= 0.0 {
132 return;
133 }
134 let frac = abs_x / self.histogram_max;
135 let bin = (frac * 256.0) as usize;
136 let bin = bin.min(255);
137 self.histogram[bin] += 1;
138 }
139
140 pub fn std_dev(&self) -> f32 {
142 if self.num_samples < 2 {
143 return 0.0;
144 }
145 let variance = self.running_var / self.num_samples as f32;
146 variance.max(0.0).sqrt()
147 }
148
149 pub fn percentile_abs(&self, p: f32) -> f32 {
153 if self.num_samples == 0 || self.histogram_max <= 0.0 {
154 return 0.0;
155 }
156
157 let p_clamped = p.clamp(0.0, 1.0);
158 let target_count = (p_clamped * self.num_samples as f32).ceil() as u64;
159 let target_count = target_count.max(1);
160
161 let mut cumulative = 0u64;
162 for (bin_idx, &count) in self.histogram.iter().enumerate() {
163 cumulative += count;
164 if cumulative >= target_count {
165 let upper = (bin_idx as f32 + 1.0) / 256.0 * self.histogram_max;
167 return upper.min(self.histogram_max);
168 }
169 }
170
171 self.histogram_max
172 }
173
174 pub fn aciq_clip(&self) -> f32 {
178 2.83 * self.std_dev()
179 }
180
181 pub fn compute_scale(&self, method: CalibMethod) -> f32 {
183 const CLIP_VAL: f32 = 127.0;
184 self.compute_scale_with_clip(method, CLIP_VAL)
185 }
186
187 pub fn compute_scale_int4(&self, method: CalibMethod) -> f32 {
189 const CLIP_VAL: f32 = 7.0;
190 self.compute_scale_with_clip(method, CLIP_VAL)
191 }
192
193 fn compute_scale_with_clip(&self, method: CalibMethod, clip_val: f32) -> f32 {
195 if self.num_samples == 0 {
196 return 0.0;
197 }
198
199 let abs_max = self.running_min.abs().max(self.running_max.abs());
200
201 let clipping_value = match method {
202 CalibMethod::MinMax => abs_max,
203 CalibMethod::Percentile(p) => {
204 let pv = self.percentile_abs(p);
205 if pv <= 0.0 {
206 abs_max
207 } else {
208 pv
209 }
210 }
211 CalibMethod::Aciq => {
212 let aciq = self.aciq_clip();
213 if aciq <= 0.0 {
214 abs_max
215 } else {
216 aciq.min(abs_max)
217 }
218 }
219 CalibMethod::Mse => {
220 self.mse_optimal_clip(abs_max, clip_val)
223 }
224 };
225
226 if clipping_value <= 0.0 {
227 return 0.0;
228 }
229
230 clipping_value / clip_val
231 }
232
233 fn mse_optimal_clip(&self, abs_max: f32, clip_val: f32) -> f32 {
238 if abs_max <= 0.0 || self.num_samples == 0 {
239 return abs_max;
240 }
241
242 let n_steps = 100usize;
243 let mut best_clip = abs_max;
244 let mut best_mse = f32::INFINITY;
245
246 for step in 1..=n_steps {
247 let frac = step as f32 / n_steps as f32;
248 let candidate_clip = abs_max * frac;
249 if candidate_clip <= 0.0 {
250 continue;
251 }
252
253 let scale = candidate_clip / clip_val;
255 let rounding_var = scale * scale / 3.0;
256
257 let clip_frac = candidate_clip / self.histogram_max;
260 let clip_bin = (clip_frac * 256.0) as usize;
261
262 let mut clip_error_sq = 0.0f32;
263 for (bin_idx, &count) in self.histogram.iter().enumerate().skip(clip_bin.min(255)) {
264 if count == 0 {
265 continue;
266 }
267 let bin_center = (bin_idx as f32 + 0.5) / 256.0 * self.histogram_max;
269 let excess = (bin_center - candidate_clip).max(0.0);
270 clip_error_sq += count as f32 * excess * excess;
271 }
272
273 let n = self.num_samples as f32;
274 let clip_mse = if n > 0.0 { clip_error_sq / n } else { 0.0 };
275 let total_mse = rounding_var + clip_mse;
276
277 if total_mse < best_mse {
278 best_mse = total_mse;
279 best_clip = candidate_clip;
280 }
281 }
282
283 best_clip
284 }
285
286 pub fn summary(&self) -> CalibSummary {
288 let p99 = self.percentile_abs(0.99);
289 let p9999 = self.percentile_abs(0.9999);
290 let int8_scale = self.compute_scale(CalibMethod::Percentile(0.9999));
291
292 let (min, max) = if self.num_samples == 0 {
293 (0.0, 0.0)
294 } else {
295 (self.running_min, self.running_max)
296 };
297
298 CalibSummary {
299 layer_name: self.layer_name.clone(),
300 num_samples: self.num_samples,
301 min,
302 max,
303 mean: self.running_mean,
304 std_dev: self.std_dev(),
305 p99,
306 p9999,
307 suggested_int8_scale: int8_scale,
308 }
309 }
310}
311
312#[derive(Debug, Clone)]
316pub struct CalibSummary {
317 pub layer_name: String,
318 pub num_samples: usize,
319 pub min: f32,
320 pub max: f32,
321 pub mean: f32,
322 pub std_dev: f32,
323 pub p99: f32,
324 pub p9999: f32,
325 pub suggested_int8_scale: f32,
326}
327
328impl CalibSummary {
329 pub fn summary_line(&self) -> String {
331 format!(
332 "[{}] n={} min={:.4} max={:.4} mean={:.4} std={:.4} p99={:.4} p9999={:.4} scale_int8={:.6}",
333 self.layer_name,
334 self.num_samples,
335 self.min,
336 self.max,
337 self.mean,
338 self.std_dev,
339 self.p99,
340 self.p9999,
341 self.suggested_int8_scale,
342 )
343 }
344}
345
346pub struct CalibrationDb {
350 layers: HashMap<String, LayerCalibStats>,
351 method: CalibMethod,
352}
353
354impl CalibrationDb {
355 pub fn new(method: CalibMethod) -> Self {
357 Self {
358 layers: HashMap::new(),
359 method,
360 }
361 }
362
363 pub fn new_minmax() -> Self {
365 Self::new(CalibMethod::MinMax)
366 }
367
368 pub fn new_percentile(p: f32) -> Self {
370 Self::new(CalibMethod::Percentile(p.clamp(0.0, 1.0)))
371 }
372
373 pub fn record(&mut self, layer_name: &str, activations: &[f32]) {
375 let stats = self
376 .layers
377 .entry(layer_name.to_owned())
378 .or_insert_with(|| LayerCalibStats::new(layer_name));
379 stats.update(activations);
380 }
381
382 pub fn get_stats(&self, layer_name: &str) -> Option<&LayerCalibStats> {
384 self.layers.get(layer_name)
385 }
386
387 pub fn scale_for_layer(&self, layer_name: &str) -> Option<f32> {
389 self.layers
390 .get(layer_name)
391 .map(|s| s.compute_scale(self.method))
392 }
393
394 pub fn export_scales(&self) -> HashMap<String, f32> {
396 self.layers
397 .iter()
398 .map(|(name, stats)| (name.clone(), stats.compute_scale(self.method)))
399 .collect()
400 }
401
402 pub fn summaries(&self) -> Vec<CalibSummary> {
404 let mut summaries: Vec<CalibSummary> = self.layers.values().map(|s| s.summary()).collect();
405 summaries.sort_by(|a, b| a.layer_name.cmp(&b.layer_name));
406 summaries
407 }
408
409 pub fn num_layers(&self) -> usize {
411 self.layers.len()
412 }
413
414 pub fn report(&self) -> String {
416 let method_str = match self.method {
417 CalibMethod::MinMax => "MinMax".to_owned(),
418 CalibMethod::Percentile(p) => format!("Percentile({:.4})", p),
419 CalibMethod::Aciq => "ACIQ".to_owned(),
420 CalibMethod::Mse => "MSE".to_owned(),
421 };
422
423 let mut lines = Vec::new();
424 lines.push("=== PTQ Calibration Report ===".to_string());
425 lines.push(format!("Method: {}", method_str));
426 lines.push(format!("Layers: {}", self.layers.len()));
427 lines.push(String::new());
428
429 let summaries = self.summaries();
430 for s in &summaries {
431 lines.push(s.summary_line());
432 }
433
434 lines.push(String::new());
435 lines.push("=== Scales ===".to_string());
436 let scales = self.export_scales();
437 let mut scale_entries: Vec<_> = scales.iter().collect();
438 scale_entries.sort_by_key(|(k, _)| k.as_str());
439 for (name, scale) in scale_entries {
440 lines.push(format!(" {}: {:.8}", name, scale));
441 }
442
443 lines.join("\n")
444 }
445}
446
447fn lcg_f32(state: &mut u64) -> f32 {
451 *state = state
452 .wrapping_mul(6_364_136_223_846_793_005)
453 .wrapping_add(1_442_695_040_888_963_407);
454 ((*state >> 32) as f32) / (u32::MAX as f32 + 1.0) * 2.0 - 1.0
455}
456
457pub fn simulate_calibration(
463 db: &mut CalibrationDb,
464 layer_names: &[&str],
465 samples_per_layer: usize,
466 seed: u64,
467) {
468 let mut state = seed;
469
470 for (layer_idx, &layer_name) in layer_names.iter().enumerate() {
471 let amplitude = 1.0 + layer_idx as f32 * 0.5;
473 let mut activations = Vec::with_capacity(samples_per_layer);
474
475 for _ in 0..samples_per_layer {
476 let v = lcg_f32(&mut state) * amplitude;
477 activations.push(v);
478 }
479
480 db.record(layer_name, &activations);
481 }
482}
483
484#[derive(Debug, Clone)]
488pub struct CalibValidation {
489 pub layer_name: String,
490 pub scale: f32,
491 pub issues: Vec<String>,
492 pub is_valid: bool,
493}
494
495impl CalibValidation {
496 pub fn validate(layer_name: &str, stats: &LayerCalibStats, scale: f32) -> Self {
498 let mut issues = Vec::new();
499
500 if !scale.is_finite() {
502 issues.push(format!("scale is not finite: {}", scale));
503 } else if scale <= 0.0 {
504 issues.push(format!("scale is non-positive: {}", scale));
505 }
506
507 if stats.num_samples == 0 {
509 issues.push("no calibration samples recorded".to_owned());
510 }
511
512 if scale.is_finite() && scale > 0.0 && stats.num_samples > 0 {
515 let abs_max = stats.running_min.abs().max(stats.running_max.abs());
516 if abs_max > 0.0 {
517 let minmax_scale = abs_max / 127.0;
519 if scale > minmax_scale * 10.0 {
520 issues.push(format!(
521 "scale {:.6} is >10x the MinMax scale {:.6} (possible overflow)",
522 scale, minmax_scale
523 ));
524 }
525 if scale < minmax_scale / 1000.0 {
527 issues.push(format!(
528 "scale {:.8} is <1/1000 of MinMax scale {:.6} (possible underflow)",
529 scale, minmax_scale
530 ));
531 }
532 } else {
533 issues.push("all activations are zero".to_owned());
535 }
536 }
537
538 if stats.num_samples > 0 && stats.running_min > stats.running_max {
540 issues.push(format!(
541 "running_min ({}) > running_max ({})",
542 stats.running_min, stats.running_max
543 ));
544 }
545
546 let is_valid = issues.is_empty();
547
548 Self {
549 layer_name: layer_name.to_owned(),
550 scale,
551 issues,
552 is_valid,
553 }
554 }
555}
556
557pub fn validate_calibration(db: &CalibrationDb) -> Vec<CalibValidation> {
559 let mut results: Vec<CalibValidation> = db
560 .layers
561 .iter()
562 .map(|(name, stats)| {
563 let scale = stats.compute_scale(db.method);
564 CalibValidation::validate(name, stats, scale)
565 })
566 .collect();
567 results.sort_by(|a, b| a.layer_name.cmp(&b.layer_name));
568 results
569}
570
571#[cfg(test)]
574mod tests {
575 use super::*;
576
577 #[test]
578 fn layer_calib_stats_new() {
579 let stats = LayerCalibStats::new("layer0");
580 assert_eq!(stats.layer_name, "layer0");
581 assert_eq!(stats.num_samples, 0);
582 assert_eq!(stats.running_var, 0.0);
583 assert_eq!(stats.running_mean, 0.0);
584 }
585
586 #[test]
587 fn layer_calib_stats_update_single() {
588 let mut stats = LayerCalibStats::new("layer0");
589 let data: Vec<f32> = (0..100).map(|i| i as f32 * 0.01 - 0.5).collect();
590 stats.update(&data);
591 assert_eq!(stats.num_samples, data.len());
592 }
593
594 #[test]
595 fn layer_calib_stats_running_min_max() {
596 let mut stats = LayerCalibStats::new("layer0");
597 stats.update(&[-3.0, 1.0, 2.5]);
598 stats.update(&[0.0, 5.0, -1.0]);
599 assert!((stats.running_min - (-3.0)).abs() < 1e-6);
600 assert!((stats.running_max - 5.0).abs() < 1e-6);
601 }
602
603 #[test]
604 fn layer_calib_stats_std_dev() {
605 let mut stats = LayerCalibStats::new("layer0");
606 let data: Vec<f32> = (0..1000).map(|i| (i as f32 - 500.0) * 0.01).collect();
608 stats.update(&data);
609 let sd = stats.std_dev();
610 assert!(
611 sd > 0.0,
612 "std_dev should be positive for non-constant data, got {sd}"
613 );
614 }
615
616 #[test]
617 fn layer_calib_stats_percentile_abs() {
618 let mut stats = LayerCalibStats::new("layer0");
619 let data: Vec<f32> = (1..=100).map(|i| i as f32).collect();
620 stats.update(&data);
621 let p100 = stats.percentile_abs(1.0);
623 assert!(
624 p100 >= 99.0,
625 "p=1.0 percentile should be near 100, got {p100}"
626 );
627 }
628
629 #[test]
630 fn layer_calib_stats_aciq_clip() {
631 let mut stats = LayerCalibStats::new("layer0");
632 let data: Vec<f32> = (0..500).map(|i| (i as f32 - 250.0) * 0.1).collect();
634 stats.update(&data);
635 let sd = stats.std_dev();
636 let aciq = stats.aciq_clip();
637 let expected = 2.83 * sd;
638 assert!(
639 (aciq - expected).abs() < 1e-4,
640 "aciq_clip={aciq}, expected 2.83*std_dev={expected}"
641 );
642 }
643
644 #[test]
645 fn layer_calib_stats_compute_scale_minmax() {
646 let mut stats = LayerCalibStats::new("layer0");
647 let data: Vec<f32> = vec![-2.54, 1.0, 0.5, -0.3, std::f32::consts::PI];
648 stats.update(&data);
649 let scale = stats.compute_scale(CalibMethod::MinMax);
650 assert!(scale > 0.0, "MinMax scale must be positive, got {scale}");
651 }
652
653 #[test]
654 fn layer_calib_stats_compute_scale_percentile() {
655 let mut stats = LayerCalibStats::new("layer0");
656 let mut data: Vec<f32> = (0..999).map(|i| (i as f32) * 0.001).collect();
658 data.push(100.0); stats.update(&data);
660
661 let scale_minmax = stats.compute_scale(CalibMethod::MinMax);
662 let scale_p99 = stats.compute_scale(CalibMethod::Percentile(0.99));
663 assert!(
665 scale_p99 <= scale_minmax,
666 "Percentile scale {scale_p99} should be <= MinMax scale {scale_minmax}"
667 );
668 }
669
670 #[test]
671 fn calib_summary_summary_line_nonempty() {
672 let mut stats = LayerCalibStats::new("fc1");
673 stats.update(&[0.1, -0.2, 0.3]);
674 let summary = stats.summary();
675 let line = summary.summary_line();
676 assert!(!line.is_empty(), "summary_line should not be empty");
677 assert!(
678 line.contains("fc1"),
679 "summary_line should contain layer name"
680 );
681 }
682
683 #[test]
684 fn calib_db_new_minmax() {
685 let db = CalibrationDb::new_minmax();
686 assert_eq!(db.num_layers(), 0);
689 }
690
691 #[test]
692 fn calib_db_record_creates_layer() {
693 let mut db = CalibrationDb::new_minmax();
694 assert_eq!(db.num_layers(), 0);
695 db.record("attn.q", &[0.1, -0.2, 0.5]);
696 assert_eq!(db.num_layers(), 1);
697 db.record("attn.k", &[0.3, 0.4]);
698 assert_eq!(db.num_layers(), 2);
699 db.record("attn.q", &[0.9, -0.1]);
701 assert_eq!(db.num_layers(), 2);
702 }
703
704 #[test]
705 fn calib_db_scale_for_unknown_layer() {
706 let db = CalibrationDb::new_minmax();
707 let result = db.scale_for_layer("nonexistent_layer");
708 assert!(result.is_none(), "Should return None for unknown layer");
709 }
710
711 #[test]
712 fn calib_db_export_scales_all_layers() {
713 let mut db = CalibrationDb::new_minmax();
714 db.record("layer_a", &[1.0, 2.0, 3.0]);
715 db.record("layer_b", &[-1.0, 0.5]);
716 db.record("layer_c", &[0.1, 0.2]);
717 let scales = db.export_scales();
718 assert_eq!(
719 scales.len(),
720 db.num_layers(),
721 "export_scales count should match num_layers"
722 );
723 assert!(scales.contains_key("layer_a"));
724 assert!(scales.contains_key("layer_b"));
725 assert!(scales.contains_key("layer_c"));
726 }
727
728 #[test]
729 fn calib_db_report_nonempty() {
730 let mut db = CalibrationDb::new_percentile(0.999);
731 db.record("block0.ffn", &[0.1, -0.5, 0.3, 0.9]);
732 let report = db.report();
733 assert!(
734 !report.is_empty(),
735 "report() should return a non-empty string"
736 );
737 assert!(
738 report.contains("block0.ffn"),
739 "report should contain layer name"
740 );
741 }
742
743 #[test]
744 fn simulate_calibration_fills_db() {
745 let mut db = CalibrationDb::new_minmax();
746 let layer_names = ["layer0", "layer1", "layer2", "layer3"];
747 simulate_calibration(&mut db, &layer_names, 256, 42);
748 assert_eq!(db.num_layers(), layer_names.len());
749 for &name in &layer_names {
750 let stats = db.get_stats(name).expect("layer should exist");
751 assert_eq!(stats.num_samples, 256);
752 }
753 }
754
755 #[test]
756 fn simulate_calibration_deterministic() {
757 let layer_names = ["attn.q", "attn.k", "attn.v", "ffn.up"];
758 let seed = 12345u64;
759
760 let mut db1 = CalibrationDb::new_minmax();
761 simulate_calibration(&mut db1, &layer_names, 128, seed);
762 let scales1 = db1.export_scales();
763
764 let mut db2 = CalibrationDb::new_minmax();
765 simulate_calibration(&mut db2, &layer_names, 128, seed);
766 let scales2 = db2.export_scales();
767
768 for name in &layer_names {
769 let s1 = scales1[*name];
770 let s2 = scales2[*name];
771 assert!(
772 (s1 - s2).abs() < 1e-8,
773 "scales should be identical for same seed: {name}: {s1} vs {s2}"
774 );
775 }
776 }
777
778 #[test]
779 fn validate_calibration_valid_layer() {
780 let mut stats = LayerCalibStats::new("block0.attn");
781 let data: Vec<f32> = (0..200).map(|i| (i as f32 - 100.0) * 0.05).collect();
783 stats.update(&data);
784 let scale = stats.compute_scale(CalibMethod::MinMax);
785 let val = CalibValidation::validate("block0.attn", &stats, scale);
786 assert!(
787 val.is_valid,
788 "should be valid for reasonable data: issues={:?}",
789 val.issues
790 );
791 assert!(val.issues.is_empty());
792 }
793
794 #[test]
795 fn validate_calibration_all_valid() {
796 let mut db = CalibrationDb::new_minmax();
797 let layer_names = [
798 "embed",
799 "block0.attn.q",
800 "block0.attn.k",
801 "block0.ffn.up",
802 "block0.ffn.down",
803 "lm_head",
804 ];
805 simulate_calibration(&mut db, &layer_names, 512, 999);
806 let validations = validate_calibration(&db);
807 assert_eq!(validations.len(), layer_names.len());
808 for v in &validations {
809 assert!(
810 v.is_valid,
811 "layer '{}' should be valid after simulate_calibration: issues={:?}",
812 v.layer_name, v.issues
813 );
814 }
815 }
816}