diff --git a/crates/core/src/perfmodel/fpm/estimator.rs b/crates/core/src/perfmodel/fpm/estimator.rs index ed0482056a..655b6c6357 100644 --- a/crates/core/src/perfmodel/fpm/estimator.rs +++ b/crates/core/src/perfmodel/fpm/estimator.rs @@ -133,6 +133,10 @@ pub enum RegressionFitKind { pub struct RegressionFitConfig { pub kind: RegressionFitKind, pub singular_ridge_scale: f64, + /// Refresh centered statistics after this many accepted insertions and + /// actual evictions per workload store. The default `None` disables periodic + /// refreshes; numerical recovery and batch-fit fallbacks remain enabled. + pub rebuild_interval: Option, } impl Default for RegressionFitConfig { @@ -140,6 +144,7 @@ impl Default for RegressionFitConfig { Self { kind: RegressionFitKind::StandardizedNnls, singular_ridge_scale: 1e-9, + rebuild_interval: None, } } } @@ -244,6 +249,7 @@ impl EstimatorConfig { } pub(crate) fn validate(&self) -> Result<(), AicError> { + validate_rebuild_interval(self.fpm_regression.fit.rebuild_interval)?; crate::config::validate_fpm_parquet_path( self.fpm_interpolation.fpm_parquet_path.as_deref(), true, @@ -301,3 +307,56 @@ impl EstimatorConfig { }) } } + +pub(super) fn validate_rebuild_interval(interval: Option) -> Result<(), AicError> { + if interval == Some(0) { + return Err(AicError::InvalidEngineConfig( + "estimator_config.fpm_regression.fit.rebuild_interval must be positive or null".into(), + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rebuild_interval_defaults_and_explicit_null_survive_serde() { + for (json, expected) in [ + ("{}", None), + (r#"{"rebuild_interval":7}"#, Some(7)), + (r#"{"rebuild_interval":4096}"#, Some(4096)), + (r#"{"rebuild_interval":null}"#, None), + ] { + let fit: RegressionFitConfig = serde_json::from_str(json).unwrap(); + assert_eq!(fit.rebuild_interval, expected); + let encoded = serde_json::to_value(&fit).unwrap(); + assert_eq!(encoded["rebuild_interval"], serde_json::json!(expected)); + let decoded: RegressionFitConfig = serde_json::from_value(encoded).unwrap(); + assert_eq!(decoded, fit); + } + } + + #[test] + fn rebuild_interval_is_validated_before_estimator_selection() { + let mut config = EstimatorConfig::default(); + for interval in [Some(1), Some(4096), Some(usize::MAX), None] { + config.fpm_regression.fit.rebuild_interval = interval; + config.validate().unwrap(); + } + config.fpm_regression.fit.rebuild_interval = Some(0); + assert!(config.validate().unwrap_err().to_string().contains( + "estimator_config.fpm_regression.fit.rebuild_interval must be positive or null" + )); + } + + #[test] + fn legacy_conversion_uses_the_canonical_rebuild_default() { + let config = EstimatorConfig::from_legacy(ForwardPassPerfOptions::default()).unwrap(); + assert_eq!( + config.fpm_regression.fit.rebuild_interval, + RegressionFitConfig::default().rebuild_interval + ); + } +} diff --git a/crates/core/src/perfmodel/fpm/model.rs b/crates/core/src/perfmodel/fpm/model.rs index 4e524df1f3..52b52d591e 100644 --- a/crates/core/src/perfmodel/fpm/model.rs +++ b/crates/core/src/perfmodel/fpm/model.rs @@ -128,7 +128,11 @@ struct RegressionStores { } impl RegressionStores { - fn new(worker_type: ForwardPassWorkerType, options: &ForwardPassPerfOptions) -> Self { + fn new( + worker_type: ForwardPassWorkerType, + options: &ForwardPassPerfOptions, + rebuild_interval: Option, + ) -> Self { use ForwardPassRegressionWorkloadKind::*; let kinds: &[ForwardPassRegressionWorkloadKind] = match worker_type { ForwardPassWorkerType::Prefill => &[PurePrefill], @@ -143,7 +147,7 @@ impl RegressionStores { Self { stores: kinds .iter() - .map(|kind| (*kind, BucketedRegression::new(options))) + .map(|kind| (*kind, BucketedRegression::new(options, rebuild_interval))) .collect(), } } @@ -281,10 +285,7 @@ impl ForwardPassPerfModel { } } - /// API: - /// `ForwardPassPerfModel::from_regression(worker_type, options) -> Result` - /// - /// Description: create a regression-only forward-pass model. + /// Create a regression-only forward-pass model after resolving configuration. /// /// This mode is for native-AIC-unsupported models. It returns `None` from /// `estimate_forward_pass_time_ms` for non-empty iterations until the @@ -294,12 +295,14 @@ impl ForwardPassPerfModel { pub(crate) fn from_regression( worker_type: ForwardPassWorkerType, options: ForwardPassPerfOptions, + rebuild_interval: Option, ) -> Result { validate_regression_options(&options)?; + super::estimator::validate_rebuild_interval(rebuild_interval)?; Ok(Self { mode: ForwardPassPerfMode::Regression { worker_type, - regression: RegressionStores::new(worker_type, &options), + regression: RegressionStores::new(worker_type, &options, rebuild_interval), }, options, last_warning: None, @@ -348,7 +351,11 @@ impl ForwardPassPerfModel { } if mode == EstimationMode::FpmRegression { let options = config.estimator_config.regression_options(); - let mut model = Self::from_regression(config.worker_type, options)?; + let mut model = Self::from_regression( + config.worker_type, + options, + config.estimator_config.fpm_regression.fit.rebuild_interval, + )?; let mut resolved = config.clone(); resolved.estimation_mode = mode; resolved.fallback_policy = ForwardPassFallbackPolicy::Deny; @@ -1165,6 +1172,72 @@ fn can_fallback_to_regression(err: &AicError) -> bool { ) } +#[cfg(test)] +mod rebuild_tests { + use super::*; + use ForwardPassRegressionWorkloadKind::*; + + #[test] + fn workload_stores_and_clones_have_independent_rebuild_clocks() { + let options = ForwardPassPerfOptions::default(); + let mut stores = + RegressionStores::new(ForwardPassWorkerType::Aggregated, &options, Some(3)); + assert!( + stores + .store_mut(PurePrefill) + .add_observation([1.0, 2.0], 3.0) + ); + for i in 1..=2 { + assert!( + stores + .store_mut(PureDecode) + .add_observation([i as f64, 1.0], 4.0) + ); + } + let mut cloned = stores.clone(); + assert!( + stores + .store_mut(PureDecode) + .add_observation([3.0, 4.0], 5.0) + ); + assert_eq!(stores.store(PureDecode).mutations_since_rebuild(), 0); + assert_eq!(stores.store(PurePrefill).mutations_since_rebuild(), 1); + assert_eq!( + stores.store(ContainsLocallyMixed).mutations_since_rebuild(), + 0 + ); + assert_eq!( + stores.store(CrossRankAggregated).mutations_since_rebuild(), + 0 + ); + assert_eq!(cloned.store(PureDecode).mutations_since_rebuild(), 2); + assert!( + cloned + .store_mut(PurePrefill) + .add_observation([2.0, 3.0], 4.0) + ); + assert_eq!(cloned.store(PurePrefill).mutations_since_rebuild(), 2); + assert_eq!(stores.store(PurePrefill).mutations_since_rebuild(), 1); + } + + #[test] + fn default_rebuild_interval_disables_periodic_refresh_for_every_workload_store() { + let mut stores = RegressionStores::new( + ForwardPassWorkerType::Aggregated, + &ForwardPassPerfOptions::default(), + super::super::estimator::RegressionFitConfig::default().rebuild_interval, + ); + // Constant features avoid fit work without numerical damage. A change + // back to a 4096-operation default would reset each clock at 4096. + for (_, store) in &mut stores.stores { + for _ in 0..2081 { + assert!(store.add_observation([1.0, 2.0], 3.0)); + } + assert_eq!(store.mutations_since_rebuild(), 4098); + } + } +} + #[cfg(test)] mod fallback_errors { use super::*; diff --git a/crates/core/src/perfmodel/fpm/regression.rs b/crates/core/src/perfmodel/fpm/regression.rs index 751418903f..39afb640fb 100644 --- a/crates/core/src/perfmodel/fpm/regression.rs +++ b/crates/core/src/perfmodel/fpm/regression.rs @@ -8,7 +8,10 @@ //! coordinates, while fitting and prediction use standardized raw features. use super::options::ForwardPassPerfOptions; -use super::samples::{BucketedSamples, StoreStats}; +use super::samples::{BucketedSamples, SampleInsertion, StoreStats}; + +mod recursive; +use recursive::RecursiveFit; const FEATURE_DIMENSION: usize = 2; const INACTIVE_SCALE_RELATIVE_TOLERANCE: f64 = 1e-12; @@ -25,7 +28,7 @@ pub(crate) struct BucketedRegression { samples: BucketedSamples, min_observations: usize, fit: Option, - ridge_scale: f64, + recursive: RecursiveFit, } impl StoreStats for BucketedRegression { @@ -39,16 +42,16 @@ impl StoreStats for BucketedRegression { } impl BucketedRegression { - pub(crate) fn new(options: &ForwardPassPerfOptions) -> Self { + pub(crate) fn new(options: &ForwardPassPerfOptions, rebuild_interval: Option) -> Self { Self { samples: BucketedSamples::new_dynamic(options, FEATURE_DIMENSION), min_observations: options.min_observations, fit: None, - ridge_scale: options.regression_ridge_scale, + recursive: RecursiveFit::new(options.regression_ridge_scale, rebuild_interval), } } - /// Retain an observation and refit from the currently retained raw data. + /// Retain an observation and update the fit for the retained raw data. /// /// Returns `false` without changing the model when a feature is negative /// or non-finite, or when the target is not finite and strictly positive. @@ -63,17 +66,28 @@ impl BucketedRegression { let bucket_x = raw_x.map(f64::ln_1p); let observation = RegressionObservation { raw_x, observed_ms }; - if !self.samples.add(bucket_x.to_vec(), observation) { + let SampleInsertion::Accepted { evicted } = self + .samples + .add_with_eviction(bucket_x.to_vec(), observation) + else { return false; - } + }; - let retained = self - .samples - .observations() - .into_iter() - .map(|(_, observation)| observation) - .collect::>(); - self.fit = fit_regression_with_ridge(&retained, self.min_observations, self.ridge_scale); + self.recursive.add(observation); + if let Some(evicted) = evicted { + self.recursive.remove(evicted); + } + debug_assert_eq!( + self.recursive.observation_count(), + self.samples.total_observations + ); + self.fit = self.recursive.fit_lazy(self.min_observations, || { + self.samples + .observations() + .into_iter() + .map(|(_, observation)| observation) + .collect() + }); true } @@ -84,6 +98,11 @@ impl BucketedRegression { let prediction = self.fit.as_ref()?.predict(raw_x)?; Some(prediction.max(MIN_POSITIVE_PREDICTION_MS)) } + + #[cfg(test)] + pub(crate) fn mutations_since_rebuild(&self) -> usize { + self.recursive.mutations_since_rebuild() + } } fn valid_features(x: &[f64; FEATURE_DIMENSION]) -> bool { @@ -188,14 +207,6 @@ struct StandardizedObservation { observed_ms: f64, } -#[cfg(test)] -fn fit_regression( - observations: &[RegressionObservation], - min_observations: usize, -) -> Option { - fit_regression_with_ridge(observations, min_observations, 1e-9) -} - fn fit_regression_with_ridge( observations: &[RegressionObservation], min_observations: usize, @@ -379,6 +390,7 @@ mod tests { RegressionObservation, Standardization, StandardizedObservation, StoreStats, fit_linear_active_set, }; + use crate::RegressionFitConfig; use crate::fpm::options::ForwardPassPerfOptions; fn regression_options() -> ForwardPassPerfOptions { @@ -398,7 +410,10 @@ mod tests { #[test] fn buckets_on_log_coordinates_but_retains_raw_observation() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); let raw_x = [99.0, 9_999.0]; assert!(regression.add_observation(raw_x, 12.0)); @@ -415,7 +430,8 @@ mod tests { bucket_count: 4, ..regression_options() }; - let mut regression = BucketedRegression::new(&options); + let mut regression = + BucketedRegression::new(&options, RegressionFitConfig::default().rebuild_interval); assert!(regression.add_observation([0.0, 0.0], 1.0)); assert!(regression.add_observation([1.0, 1.0], 2.0)); @@ -444,7 +460,10 @@ mod tests { #[test] fn rejects_invalid_raw_features_and_targets_without_mutation() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); assert!(!regression.add_observation([-1.0, 1.0], 1.0)); assert!(!regression.add_observation([f64::NAN, 1.0], 1.0)); assert!(!regression.add_observation([1.0, f64::INFINITY], 1.0)); @@ -457,7 +476,10 @@ mod tests { #[test] fn fits_standardized_raw_features_instead_of_log_features() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); for raw_x in [ [1.0, 2.0], [2.0, 7.0], @@ -476,7 +498,10 @@ mod tests { #[test] fn standardization_handles_large_feature_scale_disparity() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); for raw_x in [ [1.0e12, 1.0], [2.0e12, 8.0], @@ -499,7 +524,10 @@ mod tests { #[test] fn constant_axis_is_inactive_and_predicts_with_fit_snapshot() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); for attention in 1..=6 { let raw_x = [attention as f64, 7.0]; assert!(regression.add_observation(raw_x, 2.0 + 4.0 * raw_x[0])); @@ -521,7 +549,8 @@ mod tests { bucket_count: 1, ..ForwardPassPerfOptions::default() }; - let mut regression = BucketedRegression::new(&options); + let mut regression = + BucketedRegression::new(&options, RegressionFitConfig::default().rebuild_interval); for value in 0..=5 { let raw_x = [value as f64, (value * value) as f64]; assert!(regression.add_observation(raw_x, 1.0 + raw_x[0] + raw_x[1])); @@ -554,7 +583,10 @@ mod tests { #[test] fn all_constant_features_never_make_regression_ready() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); for observed_ms in 1..=6 { assert!(regression.add_observation([4.0, 9.0], observed_ms as f64)); } @@ -565,7 +597,10 @@ mod tests { #[test] fn identifiable_intercept_only_boundary_is_not_ready() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); for attention in 1..=6 { let value = attention as f64; assert!(regression.add_observation([value, 0.0], 100.0 - value)); @@ -573,9 +608,35 @@ mod tests { assert!(!regression.is_ready()); } + #[test] + fn negative_underlying_slope_selects_the_nonnegative_boundary() { + let options = ForwardPassPerfOptions { + regression_ridge_scale: 0.25, + ..regression_options() + }; + let mut regression = BucketedRegression::new(&options, None); + // A complete Cartesian grid has independent feature columns. For + // y=20-2*x0+3*x1, constraining the first slope to zero replaces -2*x0 + // by its mean -2. The constrained optimum is therefore 18+3*x1. + for x0 in 0..3 { + for x1 in 0..3 { + let raw_x = [x0 as f64, x1 as f64]; + assert!(regression.add_observation(raw_x, 20.0 - 2.0 * raw_x[0] + 3.0 * raw_x[1])); + } + } + let fit = regression.fit.as_ref().unwrap(); + assert_eq!(fit.coefficients[0], 0.0); + assert!(fit.coefficients[1] > 0.0); + assert_close(regression.predict(&[100.0, 4.0]).unwrap(), 30.0, 1e-10); + assert_close(regression.predict(&[0.0, 0.0]).unwrap(), 18.0, 1e-10); + } + #[test] fn positive_slope_extrapolation_is_clamped_to_prediction_floor() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); for attention in 6..=10 { let attention = attention as f64; assert!(regression.add_observation([attention, 0.0], attention - 5.0)); @@ -636,8 +697,9 @@ mod tests { 1e-12, ); - let mut regression = BucketedRegression::new(&options); - assert_eq!(regression.ridge_scale, 0.25); + let mut regression = + BucketedRegression::new(&options, RegressionFitConfig::default().rebuild_interval); + assert_eq!(options.regression_ridge_scale, 0.25); for observation in observations { assert!(regression.add_observation(observation.raw_x, observation.observed_ms)); } @@ -655,7 +717,8 @@ mod tests { bucket_count: 4, ..ForwardPassPerfOptions::default() }; - let mut regression = BucketedRegression::new(&options); + let mut regression = + BucketedRegression::new(&options, RegressionFitConfig::default().rebuild_interval); // The global oldest sample lives in the lower cell. The extreme is // instead the oldest sample in the upper cell, which is made uniquely @@ -703,7 +766,8 @@ mod tests { ..regression_options() }; - let mut two_active_axes = BucketedRegression::new(&options); + let mut two_active_axes = + BucketedRegression::new(&options, RegressionFitConfig::default().rebuild_interval); assert!(two_active_axes.add_observation([0.0, 0.0], 2.0)); assert!(two_active_axes.add_observation([1.0, 1.0], 1.0)); assert_eq!( @@ -712,7 +776,8 @@ mod tests { ); assert!(two_active_axes.is_ready()); - let mut one_active_axis = BucketedRegression::new(&options); + let mut one_active_axis = + BucketedRegression::new(&options, RegressionFitConfig::default().rebuild_interval); assert!(one_active_axis.add_observation([0.0, 0.0], 2.0)); assert!(one_active_axis.add_observation([1.0, 0.0], 1.0)); assert!(!one_active_axis.is_ready()); @@ -720,7 +785,10 @@ mod tests { #[test] fn prediction_rejects_invalid_features() { - let mut regression = BucketedRegression::new(®ression_options()); + let mut regression = BucketedRegression::new( + ®ression_options(), + RegressionFitConfig::default().rebuild_interval, + ); for raw_x in [[1.0, 2.0], [2.0, 7.0], [4.0, 3.0], [8.0, 11.0], [16.0, 5.0]] { assert!(regression.add_observation(raw_x, 5.0 + raw_x[0] + raw_x[1])); } diff --git a/crates/core/src/perfmodel/fpm/regression/recursive.rs b/crates/core/src/perfmodel/fpm/regression/recursive.rs new file mode 100644 index 0000000000..6a7df2814c --- /dev/null +++ b/crates/core/src/perfmodel/fpm/regression/recursive.rs @@ -0,0 +1,735 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Centered sufficient statistics for the existing constrained regression fit. +//! +//! Insertions and evictions update the mean and scatter of `(x0, x1, y)`. +//! Standardization and the small nonnegative slope search are then recomputed +//! from those statistics. Numerically ambiguous cases use the original batch +//! fitter on the exact retained observations. + +use super::{ + FEATURE_DIMENSION, INACTIVE_SCALE_RELATIVE_TOLERANCE, LinearFit, RegressionObservation, + Standardization, fit_regression_with_ridge, solve_linear_system, + solve_regularized_linear_system, +}; +const SCORE_RELATIVE_TOLERANCE: f64 = 1e-10; +const MIN_COVARIANCE_DETERMINANT_RATIO: f64 = 1e-8; + +#[derive(Clone, Debug)] +pub(super) struct RecursiveFit { + count: usize, + means: [f64; 3], + scatter: [[f64; 3]; 3], + mutations_since_rebuild: usize, + damaged_downdate: bool, + ridge_scale: f64, + rebuild_interval: Option, +} + +struct Candidate { + fit: LinearFit, + squared_error: f64, + error_tolerance: f64, +} + +impl RecursiveFit { + pub(super) fn new(ridge_scale: f64, rebuild_interval: Option) -> Self { + Self { + count: 0, + means: [0.0; 3], + scatter: [[0.0; 3]; 3], + mutations_since_rebuild: 0, + damaged_downdate: false, + ridge_scale, + rebuild_interval, + } + } + + pub(super) fn observation_count(&self) -> usize { + self.count + } + + #[cfg(test)] + pub(super) fn mutations_since_rebuild(&self) -> usize { + self.mutations_since_rebuild + } + + // The caller supplies validated samples and accounts for each eviction. + pub(super) fn add(&mut self, observation: RegressionObservation) { + self.insert_statistics(observation); + self.mutations_since_rebuild = self.mutations_since_rebuild.saturating_add(1); + } + + fn insert_statistics(&mut self, observation: RegressionObservation) { + let z = [ + observation.raw_x[0], + observation.raw_x[1], + observation.observed_ms, + ]; + if self.count == 0 { + self.count = 1; + self.means = z; + self.scatter = [[0.0; 3]; 3]; + return; + } + let delta = std::array::from_fn::<_, 3, _>(|i| z[i] - self.means[i]); + self.count += 1; + for (mean, difference) in self.means.iter_mut().zip(delta) { + *mean += difference / self.count as f64; + } + let remaining = std::array::from_fn::<_, 3, _>(|i| z[i] - self.means[i]); + for i in 0..3 { + for j in i..3 { + // The diagonal matches production's Welford operation. The + // symmetric cross term avoids privileging either coordinate. + let increment = if i == j { + delta[i] * remaining[i] + } else { + 0.5 * delta[i] * remaining[j] + 0.5 * delta[j] * remaining[i] + }; + self.scatter[i][j] += increment; + self.scatter[j][i] = self.scatter[i][j]; + } + } + } + + pub(super) fn remove(&mut self, observation: RegressionObservation) { + self.mutations_since_rebuild = self.mutations_since_rebuild.saturating_add(1); + if self.count == 0 { + self.damaged_downdate = true; + return; + } + if self.count == 1 { + self.count = 0; + self.means = [0.0; 3]; + self.scatter = [[0.0; 3]; 3]; + return; + } + let z = [ + observation.raw_x[0], + observation.raw_x[1], + observation.observed_ms, + ]; + let delta = std::array::from_fn::<_, 3, _>(|i| z[i] - self.means[i]); + self.count -= 1; + for (mean, difference) in self.means.iter_mut().zip(delta) { + *mean -= difference / self.count as f64; + } + let remaining = std::array::from_fn::<_, 3, _>(|i| z[i] - self.means[i]); + for i in 0..3 { + for j in i..3 { + let previous = self.scatter[i][j]; + let decrement = if i == j { + delta[i] * remaining[i] + } else { + 0.5 * delta[i] * remaining[j] + 0.5 * delta[j] * remaining[i] + }; + self.scatter[i][j] -= decrement; + self.scatter[j][i] = self.scatter[i][j]; + // Losing almost all variance can expose previously hidden + // cancellation, including a constant axis becoming active. + if i == j + && (self.scatter[i][i] < 0.0 + || (previous > 0.0 && self.scatter[i][i] < previous * 1e-8)) + { + self.damaged_downdate = true; + } + } + } + } + + fn rebuild(&mut self, retained: &[RegressionObservation]) { + self.count = 0; + self.means = [0.0; 3]; + self.scatter = [[0.0; 3]; 3]; + for &observation in retained { + self.insert_statistics(observation); + } + self.mutations_since_rebuild = 0; + self.damaged_downdate = false; + } + + fn statistics_are_finite(&self) -> bool { + self.means.iter().all(|value| value.is_finite()) + && self.scatter.iter().flatten().all(|value| value.is_finite()) + && (0..3).all(|i| self.scatter[i][i] >= 0.0) + } + + fn fallback( + &self, + retained: &mut impl FnMut() -> Vec, + min_observations: usize, + ) -> Option { + // A batch fallback does not reset the statistics or rebuild clock. + fit_regression_with_ridge(&retained(), min_observations, self.ridge_scale) + } + + #[cfg(test)] + pub(super) fn fit( + &mut self, + retained: &[RegressionObservation], + min_observations: usize, + ) -> Option { + if self.count != retained.len() { + self.rebuild(retained); + } + self.fit_lazy(min_observations, || retained.to_vec()) + } + + /// Fit after the complete insertion/eviction transaction. One insertion + /// and one eviction count as two mutations; rebucketing counts as none. + /// The interval is checked once here, so an odd threshold can be exceeded + /// by one mutation before a rebuild resets the clock. Disabling periodic + /// rebuilds leaves numerical recovery and batch fallbacks enabled. + /// + /// The healthy path does not materialize or inspect retained samples. + /// The caller must account for every insertion and actual eviction. + pub(super) fn fit_lazy( + &mut self, + min_observations: usize, + mut retained: impl FnMut() -> Vec, + ) -> Option { + let periodic = self + .rebuild_interval + .is_some_and(|interval| self.mutations_since_rebuild >= interval); + if self.damaged_downdate || periodic || !self.statistics_are_finite() { + self.rebuild(&retained()); + } + if self.count < min_observations || self.count == 0 { + return None; + } + if !self.statistics_are_finite() { + return self.fallback(&mut retained, min_observations); + } + + let mut standardization = Standardization { + means: [self.means[0], self.means[1]], + scales: [0.0; FEATURE_DIMENSION], + active: [false; FEATURE_DIMENSION], + }; + let mut varying_axes = Vec::with_capacity(FEATURE_DIMENSION); + for axis in 0..FEATURE_DIMENSION { + let scale = (self.scatter[axis][axis] / self.count as f64) + .max(0.0) + .sqrt(); + let threshold = INACTIVE_SCALE_RELATIVE_TOLERANCE * self.means[axis].abs().max(1.0); + standardization.scales[axis] = scale; + standardization.active[axis] = scale.is_finite() && scale > threshold; + // Very small nonzero spread is sensitive to mean rounding and + // updates can disagree with batch fitting about active axes. + if scale > 0.0 && scale <= threshold * 1e4 { + return self.fallback(&mut retained, min_observations); + } + if standardization.active[axis] { + varying_axes.push(axis); + } + } + if varying_axes.is_empty() { + return None; + } + + if varying_axes.len() == 2 { + let correlation = (self.scatter[0][1] / standardization.scales[0]) + / standardization.scales[1] + / self.count as f64; + if !correlation.is_finite() + || 1.0 - correlation * correlation <= MIN_COVARIANCE_DETERMINANT_RATIO + { + // Production can pick a different face under near-perfect + // collinearity; matching only in-sample predictions is unsafe. + return self.fallback(&mut retained, min_observations); + } + } + + let target_scale = (self.scatter[2][2] / self.count as f64).sqrt(); + let slope_tolerance = target_scale * 1e-10 + self.means[2].abs() * f64::EPSILON * 32.0; + let mut candidates = Vec::with_capacity(1 << varying_axes.len()); + for mask in 0..(1usize << varying_axes.len()) { + let fitted_axes = varying_axes + .iter() + .enumerate() + .filter_map(|(i, &axis)| ((mask & (1 << i)) != 0).then_some(axis)) + .collect::>(); + let size = fitted_axes.len() + 1; + let mut lhs = vec![vec![0.0; size]; size]; + let mut rhs = vec![0.0; size]; + lhs[0][0] = self.count as f64; + rhs[0] = self.count as f64 * self.means[2]; + for (i, &axis_i) in fitted_axes.iter().enumerate() { + rhs[i + 1] = self.scatter[axis_i][2] / standardization.scales[axis_i]; + for (j, &axis_j) in fitted_axes.iter().enumerate() { + lhs[i + 1][j + 1] = (self.scatter[axis_i][axis_j] + / standardization.scales[axis_i]) + / standardization.scales[axis_j]; + } + } + let solution = solve_linear_system(lhs.clone(), rhs.clone()).or_else(|| { + solve_regularized_linear_system(lhs.clone(), rhs.clone(), self.ridge_scale) + }); + let Some(solution) = solution else { + return self.fallback(&mut retained, min_observations); + }; + if !solution.iter().all(|value| value.is_finite()) { + return self.fallback(&mut retained, min_observations); + } + if solution[1..] + .iter() + .any(|value| value.abs() <= slope_tolerance) + { + return self.fallback(&mut retained, min_observations); + } + if solution[1..].iter().any(|value| *value < 0.0) { + continue; + } + + let mut coefficients = [0.0; FEATURE_DIMENSION]; + let mut linear_term = 0.0; + let mut quadratic_term = 0.0; + let mut magnitude = self.scatter[2][2].abs(); + for (i, &axis) in fitted_axes.iter().enumerate() { + coefficients[axis] = solution[i + 1]; + let term = 2.0 * solution[i + 1] * rhs[i + 1]; + linear_term += term; + magnitude += term.abs(); + for j in 0..fitted_axes.len() { + let term = solution[i + 1] * lhs[i + 1][j + 1] * solution[j + 1]; + quadratic_term += term; + magnitude += term.abs(); + } + } + let intercept_error = solution[0] - self.means[2]; + // Score every face by unpenalized SSE, including ridge retries: + // Cyy - 2 b' D^-1 Cxy + b' D^-1 Cxx D^-1 b + n(a - mean_y)^2. + // Near ties require batch residuals because these terms can cancel. + let squared_error = self.scatter[2][2] - linear_term + + quadratic_term + + self.count as f64 * intercept_error * intercept_error; + let error_tolerance = SCORE_RELATIVE_TOLERANCE * magnitude.max(1e-24); + if !squared_error.is_finite() || squared_error < -error_tolerance { + return self.fallback(&mut retained, min_observations); + } + candidates.push(Candidate { + fit: LinearFit { + intercept: solution[0], + coefficients, + standardization, + }, + squared_error, + error_tolerance, + }); + } + + let mut best_index = 0; + for i in 1..candidates.len() { + if candidates[i].squared_error < candidates[best_index].squared_error { + best_index = i; + } + } + let best = candidates.get(best_index)?; + if candidates.iter().enumerate().any(|(index, other)| { + index != best_index + && (other.squared_error - best.squared_error).abs() + <= other.error_tolerance + best.error_tolerance + }) { + return self.fallback(&mut retained, min_observations); + } + let best = candidates.swap_remove(best_index).fit; + let is_underdetermined = self.count <= varying_axes.len(); + let has_load_signal = best.coefficients.iter().any(|value| *value > 0.0); + (is_underdetermined || has_load_signal).then_some(best) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::RegressionFitConfig; + use crate::fpm::options::ForwardPassPerfOptions; + use crate::fpm::regression::BucketedRegression; + + fn observation(i: usize) -> RegressionObservation { + let x = [(i % 43) as f64 + 1.0, ((i * 29) % 67) as f64 + 1.0]; + RegressionObservation { + raw_x: x, + observed_ms: 4.0 + 1.7 * x[0] + 0.4 * x[1] + (i % 7) as f64 * 0.01, + } + } + + fn default_recursive() -> RecursiveFit { + let config = RegressionFitConfig::default(); + RecursiveFit::new(config.singular_ridge_scale, config.rebuild_interval) + } + + fn close(actual: f64, expected: f64) { + assert!( + (actual - expected).abs() <= 1e-8 + 1e-8 * expected.abs(), + "expected {expected:.17e}, got {actual:.17e}" + ); + } + + fn compare_fits( + incremental: Option<&LinearFit>, + batch: Option<&LinearFit>, + retained: &[RegressionObservation], + ) { + assert_eq!(incremental.is_some(), batch.is_some()); + if let (Some(incremental), Some(batch)) = (incremental, batch) { + assert_eq!( + incremental.standardization.active, + batch.standardization.active + ); + close(incremental.intercept, batch.intercept); + for axis in 0..FEATURE_DIMENSION { + assert_eq!( + incremental.coefficients[axis] > 0.0, + batch.coefficients[axis] > 0.0 + ); + close(incremental.coefficients[axis], batch.coefficients[axis]); + close( + incremental.standardization.means[axis], + batch.standardization.means[axis], + ); + close( + incremental.standardization.scales[axis], + batch.standardization.scales[axis], + ); + } + for x in retained.iter().map(|sample| sample.raw_x).chain([ + [0.0, 0.0], + [17.0, 91.0], + [137.0, 4.0], + [1000.0, 1000.0], + ]) { + close(incremental.predict(&x).unwrap(), batch.predict(&x).unwrap()); + } + } + } + + fn compare(recursive: &mut RecursiveFit, retained: &[RegressionObservation], min: usize) { + let incremental = recursive.fit(retained, min); + let batch = fit_regression_with_ridge(retained, min, recursive.ridge_scale); + compare_fits(incremental.as_ref(), batch.as_ref(), retained); + } + + #[test] + fn centered_updates_track_batch_across_evictions_and_periodic_rebuilds() { + for interval in [RegressionFitConfig::default().rebuild_interval, Some(4096)] { + let mut recursive = RecursiveFit::new(1e-9, interval); + let mut retained = Vec::new(); + for i in 0..4300 { + let incoming = observation(i); + retained.push(incoming); + recursive.add(incoming); + if retained.len() > 64 { + recursive.remove(retained.remove(0)); + } + compare(&mut recursive, &retained, 5); + } + } + } + + #[test] + fn sliding_window_tracks_scale_changes_and_axes_becoming_inactive() { + for ridge_scale in [0.0, 0.25] { + let mut recursive = RecursiveFit::new(ridge_scale, None); + let mut retained = Vec::new(); + for i in 0..128 { + let raw_x = match i / 32 { + 0 => [((i % 11) + 1) as f64 * 1e12, 7.0], + 1 => [ + ((i % 11) + 1) as f64 * 1e12, + (((i * 7) % 19) + 1) as f64 * 0.01, + ], + 2 => [3e12, (((i * 7) % 19) + 1) as f64 * 0.01], + _ => [((i % 11) + 1) as f64 * 1000.0, 7.0], + }; + let sample = RegressionObservation { + raw_x, + observed_ms: 5.0 + 2.5e-9 * raw_x[0] + 4.0 * raw_x[1], + }; + retained.push(sample); + recursive.add(sample); + if retained.len() > 16 { + recursive.remove(retained.remove(0)); + } + compare(&mut recursive, &retained, 5); + if i == 31 || i == 95 || i == 127 { + let fit = recursive.fit(&retained, 5).unwrap(); + assert_eq!( + fit.standardization.active, + if i == 95 { + [false, true] + } else { + [true, false] + } + ); + } + } + } + } + + #[test] + fn periodic_schedule_counts_mutations_and_checks_once_per_transaction() { + let cases = [ + (Some(1), 4, 12, (1..=12).collect::>()), + // Four warmup insertions, then two mutations per replacement: + // threshold five is first crossed at six and reset to zero. + (Some(5), 4, 12, vec![5, 8, 11]), + (Some(4096), 64, 4300, vec![2080, 4128]), + ( + RegressionFitConfig::default().rebuild_interval, + 64, + 4300, + vec![], + ), + ]; + for (interval, capacity, updates, expected) in cases { + let mut recursive = RecursiveFit::new(0.25, interval); + let mut retained = Vec::new(); + let mut rebuilt_at = Vec::new(); + for input in 1..=updates { + let sample = observation(input); + retained.push(sample); + recursive.add(sample); + if retained.len() > capacity { + recursive.remove(retained.remove(0)); + } + // An unreachable readiness threshold isolates rebuild work + // from numerical candidate fallbacks without test counters. + let mut reads = 0; + assert!( + recursive + .fit_lazy(usize::MAX, || { + reads += 1; + retained.clone() + }) + .is_none() + ); + if reads > 0 { + assert_eq!(reads, 1); + assert_eq!(recursive.mutations_since_rebuild, 0); + rebuilt_at.push(input); + } + } + assert_eq!(rebuilt_at, expected, "interval {interval:?}"); + assert_eq!(recursive.rebuild_interval, interval); + assert_eq!(recursive.ridge_scale, 0.25); + } + } + + #[test] + fn healthy_fit_does_not_materialize_retained_samples() { + let mut recursive = default_recursive(); + for sample in (0..64).map(observation) { + recursive.add(sample); + } + assert!( + recursive + .fit_lazy(5, || panic!("healthy fit must be incremental")) + .is_some() + ); + } + + #[test] + fn collinear_fallback_preserves_configured_ridge_and_rebuild_clock() { + let mut recursive = RecursiveFit::new(0.25, None); + let retained = (0..16) + .map(|i| RegressionObservation { + raw_x: [i as f64, 2.0 * i as f64], + observed_ms: 3.0 + 7.0 * i as f64, + }) + .collect::>(); + for &sample in &retained { + recursive.add(sample); + } + let mut reads = 0; + let incremental = recursive.fit_lazy(5, || { + reads += 1; + retained.clone() + }); + let batch = fit_regression_with_ridge(&retained, 5, 0.25); + compare_fits(incremental.as_ref(), batch.as_ref(), &retained); + assert_eq!(reads, 1); + assert_eq!(recursive.mutations_since_rebuild, retained.len()); + } + + #[test] + fn zero_slope_boundary_uses_batch_readiness() { + let mut recursive = default_recursive(); + let retained = (0..16) + .map(|i| RegressionObservation { + raw_x: [i as f64, 0.0], + observed_ms: 3.0, + }) + .collect::>(); + for &sample in &retained { + recursive.add(sample); + } + let mut reads = 0; + assert!( + recursive + .fit_lazy(5, || { + reads += 1; + retained.clone() + }) + .is_none() + ); + assert_eq!(reads, 1); + assert!(fit_regression_with_ridge(&retained, 5, recursive.ridge_scale).is_none()); + } + + #[test] + fn disabled_periodic_rebuilds_still_recover_after_outlier_removal() { + let mut recursive = RecursiveFit::new(0.25, None); + let retained = (0..16) + .map(|i| RegressionObservation { + raw_x: [10.0, i as f64], + observed_ms: 3.0 + 7.0 * i as f64, + }) + .collect::>(); + let outlier = RegressionObservation { + raw_x: [1e12, 1e6], + observed_ms: 1e12, + }; + for &sample in &retained { + recursive.add(sample); + } + recursive.add(outlier); + recursive.remove(outlier); + assert!(recursive.damaged_downdate); + compare(&mut recursive, &retained, 5); + assert!(!recursive.damaged_downdate); + assert_eq!(recursive.mutations_since_rebuild, 0); + assert_eq!(recursive.scatter[0][0], 0.0); + assert_eq!(recursive.rebuild_interval, None); + assert_eq!(recursive.ridge_scale, 0.25); + } + + #[test] + fn disabled_periodic_rebuilds_still_recover_nonfinite_statistics() { + let mut recursive = RecursiveFit::new(0.25, None); + let retained = (0..64).map(observation).collect::>(); + for &sample in &retained { + recursive.add(sample); + } + recursive.scatter[0][1] = f64::INFINITY; + compare(&mut recursive, &retained, 5); + assert!(recursive.statistics_are_finite()); + assert_eq!(recursive.mutations_since_rebuild, 0); + } + + #[test] + fn mutation_clock_saturates_when_periodic_rebuilds_are_disabled() { + let mut recursive = RecursiveFit::new(0.25, None); + for sample in (0..64).map(observation) { + recursive.add(sample); + } + recursive.mutations_since_rebuild = usize::MAX - 1; + let sample = observation(100); + recursive.add(sample); + recursive.remove(sample); + assert_eq!(recursive.mutations_since_rebuild, usize::MAX); + assert!( + recursive + .fit_lazy(usize::MAX, || panic!("disabled periodic rebuild")) + .is_none() + ); + // A maximum finite interval must still trigger at the saturated value. + recursive.rebuild_interval = Some(usize::MAX); + let retained = (0..64).map(observation).collect::>(); + recursive.fit_lazy(usize::MAX, || retained.clone()); + assert_eq!(recursive.mutations_since_rebuild, 0); + } + + #[test] + fn clone_preserves_configuration_and_has_independent_statistics() { + let mut original = RecursiveFit::new(0.25, Some(37)); + for sample in (0..20).map(observation) { + original.add(sample); + } + let means = original.means; + let scatter = original.scatter; + let mut cloned = original.clone(); + assert_eq!(cloned.rebuild_interval, Some(37)); + assert_eq!(cloned.ridge_scale, 0.25); + assert_eq!(cloned.mutations_since_rebuild, 20); + cloned.add(observation(100)); + assert_eq!(original.count, 20); + assert_eq!(original.mutations_since_rebuild, 20); + assert_eq!(original.means, means); + assert_eq!(original.scatter, scatter); + assert_eq!(cloned.count, 21); + assert_ne!(cloned.means, means); + } + + #[test] + fn removing_last_observation_preserves_configuration_and_allows_reuse() { + let mut recursive = RecursiveFit::new(0.25, None); + let first = observation(0); + recursive.add(first); + recursive.remove(first); + assert!(recursive.fit(&[], 1).is_none()); + assert_eq!(recursive.ridge_scale, 0.25); + assert_eq!(recursive.rebuild_interval, None); + let retained = (1..20).map(observation).collect::>(); + for &sample in &retained { + recursive.add(sample); + } + compare(&mut recursive, &retained, 5); + } + + #[test] + fn bucketed_updates_match_exact_retained_batch_with_rectangular_grid() { + let options = ForwardPassPerfOptions { + min_observations: 5, + max_observations: 32, + bucket_shape: Some([2, 3]), + regression_ridge_scale: 0.25, + ..ForwardPassPerfOptions::default() + }; + for interval in [None, Some(17)] { + let mut regression = BucketedRegression::new(&options, interval); + for i in 0..300 { + let sample = observation(i); + assert!(regression.add_observation(sample.raw_x, sample.observed_ms)); + let retained = regression + .samples + .observations() + .into_iter() + .map(|(_, sample)| sample) + .collect::>(); + assert_eq!(regression.recursive.count, retained.len()); + let batch = fit_regression_with_ridge(&retained, 5, 0.25); + compare_fits(regression.fit.as_ref(), batch.as_ref(), &retained); + } + } + } + + #[test] + fn rejected_inputs_and_rebucketing_do_not_add_mutations() { + let options = ForwardPassPerfOptions { + min_observations: 5, + max_observations: 4, + bucket_shape: Some([2, 3]), + ..ForwardPassPerfOptions::default() + }; + let mut regression = BucketedRegression::new(&options, None); + for (index, x) in [[0.0, 0.0], [1.0, 1.0], [15.0, 15.0]] + .into_iter() + .enumerate() + { + assert!(regression.add_observation(x, 1.0 + index as f64)); + assert_eq!(regression.recursive.mutations_since_rebuild, index + 1); + assert_eq!(regression.recursive.count, index + 1); + } + assert!(!regression.add_observation([f64::NAN, 1.0], 2.0)); + assert!(!regression.add_observation([1.0, 1.0], -1.0)); + assert_eq!(regression.recursive.mutations_since_rebuild, 3); + assert_eq!(regression.recursive.count, 3); + assert!(regression.add_observation([3.0, 0.0], 4.0)); + assert!(regression.add_observation([5.0, 5.0], 5.0)); + assert_eq!(regression.recursive.count, 4); + assert_eq!(regression.recursive.mutations_since_rebuild, 6); + } +} diff --git a/crates/core/src/perfmodel/fpm/samples.rs b/crates/core/src/perfmodel/fpm/samples.rs index 8a6665db20..4e11b58a70 100644 --- a/crates/core/src/perfmodel/fpm/samples.rs +++ b/crates/core/src/perfmodel/fpm/samples.rs @@ -24,6 +24,15 @@ pub(crate) trait StoreStats { fn is_ready(&self) -> bool; } +/// The accepted sample may immediately evict an observation from a full store. +/// Returning the actual value keeps incremental consumers in sync with the +/// sampler's retention policy, including ties between equally full buckets. +#[derive(Debug, PartialEq)] +pub(crate) enum SampleInsertion { + Accepted { evicted: Option }, + Rejected, +} + #[derive(Clone, Debug)] pub(crate) struct BucketedSamples { pub(crate) buckets: HashMap, Vec<(Vec, T)>>, @@ -83,13 +92,20 @@ impl BucketedSamples { } pub(crate) fn add(&mut self, x: Vec, y: T) -> bool { + matches!( + self.add_with_eviction(x, y), + SampleInsertion::Accepted { .. } + ) + } + + pub(crate) fn add_with_eviction(&mut self, x: Vec, y: T) -> SampleInsertion { if x.len() != self.axis_min.len() || !x.iter().all(|value| value.is_finite()) { - return false; + return SampleInsertion::Rejected; } if self.fixed_bounds { if !self.is_in_bounds(&x) { - return false; + return SampleInsertion::Rejected; } } else { let bounds_changed = self.update_axis_bounds(&x); @@ -102,10 +118,12 @@ impl BucketedSamples { self.buckets.entry(key).or_default().push((x, y)); self.total_observations += 1; - if self.total_observations > self.max_observations { - self.retire_from_fattest_bucket(); - } - true + let evicted = if self.total_observations > self.max_observations { + self.retire_from_fattest_bucket() + } else { + None + }; + SampleInsertion::Accepted { evicted } } pub(crate) fn observations(&self) -> Vec<(Vec, T)> { @@ -200,28 +218,27 @@ impl BucketedSamples { } } - fn retire_from_fattest_bucket(&mut self) { + fn retire_from_fattest_bucket(&mut self) -> Option { // Eviction removes samples only. Dynamic regression bounds remain // monotonic; fixed correction-grid workload ranges are configured at // model creation. - let Some(key) = self + let key = self .buckets .iter() .max_by_key(|(_, bucket)| bucket.len()) - .map(|(key, _)| key.clone()) - else { - return; - }; + .map(|(key, _)| key.clone())?; + let mut evicted = None; if let Some(bucket) = self.buckets.get_mut(&key) { if !bucket.is_empty() { - bucket.remove(0); + evicted = Some(bucket.remove(0).1); self.total_observations -= 1; } if bucket.is_empty() { self.buckets.remove(&key); } } + evicted } } @@ -244,3 +261,118 @@ pub(crate) fn median_ratio(values: impl Iterator) -> Option { pub(crate) fn integer_sqrt(value: usize) -> usize { (value as f64).sqrt() as usize } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn insertion_returns_the_actual_fattest_bucket_eviction() { + let options = ForwardPassPerfOptions { + max_observations: 4, + bucket_shape: Some([2, 3]), + ..ForwardPassPerfOptions::default() + }; + let ranges = [AxisRange::from_zero_to(10); 2]; + let mut samples = BucketedSamples::new_fixed(&options, &ranges); + for (index, x) in [[0.0, 0.0], [10.0, 10.0], [7.0, 7.0], [8.0, 8.0]] + .into_iter() + .enumerate() + { + assert_eq!( + samples.add_with_eviction(x.to_vec(), index), + SampleInsertion::Accepted { evicted: None } + ); + } + // The upper bucket is uniquely fattest (three rows). Its oldest row + // is [10,10], whereas the global oldest row [0,0] must be retained. + assert_eq!( + samples.add_with_eviction(vec![2.0, 0.0], 4), + SampleInsertion::Accepted { evicted: Some(1) } + ); + let mut retained = samples + .observations() + .into_iter() + .map(|(_, value)| value) + .collect::>(); + retained.sort_unstable(); + assert_eq!(retained, [0, 2, 3, 4]); + assert_eq!(samples.total_observations, 4); + assert_eq!(samples.axis_max, [10.0, 10.0]); + } + + #[test] + fn rejected_insertions_do_not_change_samples_or_bounds() { + let options = ForwardPassPerfOptions { + max_observations: 2, + ..ForwardPassPerfOptions::default() + }; + let mut samples = BucketedSamples::new_fixed(&options, &[AxisRange::from_zero_to(10); 2]); + assert!(samples.add(vec![2.0, 3.0], 1)); + let before = samples.clone(); + for x in [ + vec![4.0], + vec![f64::NAN, 3.0], + vec![2.0, f64::INFINITY], + vec![11.0, 3.0], + ] { + assert_eq!(samples.add_with_eviction(x, 2), SampleInsertion::Rejected); + assert_eq!(samples.buckets, before.buckets); + assert_eq!(samples.axis_min, before.axis_min); + assert_eq!(samples.axis_max, before.axis_max); + assert_eq!(samples.total_observations, before.total_observations); + } + assert_eq!( + samples.add_with_eviction(vec![3.0, 4.0], 2), + SampleInsertion::Accepted { evicted: None } + ); + } + + #[test] + fn dynamic_rebucketing_preserves_rows_and_rectangular_geometry() { + let options = ForwardPassPerfOptions { + max_observations: 8, + bucket_shape: Some([2, 3]), + ..ForwardPassPerfOptions::default() + }; + let mut samples = BucketedSamples::new_dynamic(&options, 2); + for (index, x) in [[0.0, 0.0], [1.0, 1.0], [15.0, 15.0]] + .into_iter() + .enumerate() + { + assert_eq!( + samples.add_with_eviction(x.to_vec(), index), + SampleInsertion::Accepted { evicted: None } + ); + } + assert_eq!(samples.total_observations, 3); + assert_eq!(samples.buckets[&vec![0, 0]].len(), 2); + assert_eq!(samples.buckets[&vec![1, 2]], vec![(vec![15.0, 15.0], 2)]); + assert!(samples.buckets[&vec![0, 0]].contains(&(vec![1.0, 1.0], 1))); + } + + #[test] + fn boolean_add_preserves_retention_including_bucket_ties() { + let options = ForwardPassPerfOptions { + max_observations: 16, + bucket_shape: Some([2, 3]), + ..ForwardPassPerfOptions::default() + }; + let mut with_events = BucketedSamples::new_dynamic(&options, 2); + // Clone the initial hasher so tied bucket choices are comparable. + let mut boolean_only = with_events.clone(); + for i in 0..100 { + let x = vec![(i % 17) as f64, ((i * 7) % 23) as f64]; + assert!(boolean_only.add(x.clone(), i)); + let SampleInsertion::Accepted { evicted } = with_events.add_with_eviction(x, i) else { + panic!("valid sample rejected"); + }; + assert_eq!(evicted.is_some(), i >= 16); + assert_eq!(with_events.buckets, boolean_only.buckets); + assert_eq!( + with_events.total_observations, + boolean_only.total_observations + ); + } + } +} diff --git a/crates/core/src/perfmodel/fpm/tests.rs b/crates/core/src/perfmodel/fpm/tests.rs index e03562ec3f..ea807e34b1 100644 --- a/crates/core/src/perfmodel/fpm/tests.rs +++ b/crates/core/src/perfmodel/fpm/tests.rs @@ -170,7 +170,11 @@ fn regression_model( worker_type: ForwardPassWorkerType, options: ForwardPassPerfOptions, ) -> Result { - ForwardPassPerfModel::from_regression(worker_type, options) + ForwardPassPerfModel::from_regression( + worker_type, + options, + super::estimator::RegressionFitConfig::default().rebuild_interval, + ) } fn fixture_engine() -> Arc { diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index 5157021eea..b9895f18f2 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -246,6 +246,17 @@ mod tests { #[test] fn regression_constructor_is_environment_independent() { let model = regression_model().expect("construct regression model"); + assert_eq!( + model + .provenance() + .unwrap() + .config + .estimator_config + .fpm_regression + .fit + .rebuild_interval, + None + ); let stores = regression_stores(&model); assert_eq!(stores.len(), 4); assert_eq!( @@ -270,6 +281,74 @@ mod tests { assert_eq!(options.features.ffn_token_weight, 4.0); } + #[test] + fn regression_rebuild_interval_is_public_and_preserved_by_canonical_reload() { + use aisimulate_core::RegressionFitConfig; + + assert_eq!(RegressionFitConfig::default().rebuild_interval, None); + for interval in [Some(1), Some(17), Some(4096), None] { + let mut config = ForwardPassPerfModelConfig::new( + "test/model", + "test-system", + BackendKind::Vllm, + ForwardPassWorkerType::Decode, + ); + config.estimation_mode = EstimationMode::FpmRegression; + config.estimator_config.fpm_regression.fit = RegressionFitConfig { + rebuild_interval: interval, + ..RegressionFitConfig::default() + }; + let model = ForwardPassPerfModel::best_available(config).unwrap(); + let resolved = &model.provenance().unwrap().config; + assert_eq!( + resolved + .estimator_config + .fpm_regression + .fit + .rebuild_interval, + interval + ); + let reloaded = ForwardPassPerfModel::best_available(resolved.clone()).unwrap(); + assert_eq!(&reloaded.provenance().unwrap().config, resolved); + assert_eq!(reloaded.regression_store_diagnostics().len(), 1); + assert!(!reloaded.regression_store_diagnostics()[0].ready); + } + } + + #[test] + fn zero_regression_rebuild_interval_is_invalid_before_fallback() { + use aisimulate_core::ForwardPassFallbackPolicy; + + for mode in [EstimationMode::Auto, EstimationMode::FpmRegression] { + let mut config = ForwardPassPerfModelConfig::new( + "test/model", + "test-system", + BackendKind::Vllm, + ForwardPassWorkerType::Decode, + ); + config.estimation_mode = mode; + config.fallback_policy = ForwardPassFallbackPolicy::Allow; + config.estimator_config.fpm_regression.fit.rebuild_interval = Some(0); + let error = ForwardPassPerfModel::best_available(config) + .err() + .expect("zero interval must fail"); + assert!(matches!(error, AicError::InvalidEngineConfig(_))); + assert!(error + .to_string() + .contains("estimator_config.fpm_regression.fit.rebuild_interval")); + } + } + + #[test] + fn legacy_regression_options_keep_rust_owned_fit_defaults() { + use aisimulate_core::ForwardPassPerfOptions; + + let migrated = EstimatorConfig::from_legacy(ForwardPassPerfOptions::default()).unwrap(); + assert_eq!(migrated.fpm_regression.fit.rebuild_interval, None); + assert_eq!(migrated.fpm_regression.sampling.max_observations, 64); + assert_eq!(migrated.fpm_regression.sampling.bins_per_axis, [4, 4]); + } + #[test] fn agentic_snapshot_preparation_and_execution_are_public() { use aisimulate_core::replay::loadgen::{ diff --git a/docs/core-api.md b/docs/core-api.md index 05fce8a944..08243ea6fa 100644 --- a/docs/core-api.md +++ b/docs/core-api.md @@ -175,6 +175,8 @@ config = ForwardPassPerfModelConfig( "fpm_regression": { "sampling": {"bins_per_axis": [4, 16], "max_observations": 128}, "min_observations": 5, + # Opt into periodic rebuilding; omission leaves it disabled. + "fit": {"rebuild_interval": 4096}, }, "correction": {"enabled": True}, }, @@ -313,7 +315,9 @@ nested paths. The supported namespaces are: - `fpm_regression`: independent `sampling`, `min_observations` (5), and `fit`. The fit kind is `standardized_nnls`, with a free intercept and nonnegative slopes. `singular_ridge_scale` defaults to `1e-9` and applies only when retrying - a singular equation. + a singular equation. `rebuild_interval` defaults to JSON `null` / Python + `None`, disabling periodic rebuilding. A positive integer opts into + rebuilding after that many retained-sample mutations. - `correction`: `enabled` (true), independent `sampling`, `min_observations` (5), `factor_bounds` (min 0.5, max 2.0), and the existing `max_num_tokens` (8192), `max_batch_size` (512), and `max_kv_tokens` (2000000) ranges. @@ -347,6 +351,76 @@ return `None`. `tune_with_fpms()` preserves the established FPM observation contract. Native construction still uses Python model compilation; estimator selection, regression, correction, and latency computation are owned by Rust. +### Recursive regression and statistics rebuilding + +Regression maintains centered sufficient statistics for the retained samples +and applies the existing standardized nonnegative least-squares fit. The +objective and readiness rules stay the same. The retention grid still controls +which samples are kept; it does not create separate fitted planes within a +workload store. + +See the [recursive regression walkthrough](fpm-recursive-regression.md) for +the update equations, numerical guards, and measured fitting costs. + +Set `estimator_config.fpm_regression.fit.rebuild_interval` through the canonical +constructor. Rust owns its default and validation. Omission means `None` +(`null` in JSON), so periodic rebuilding is disabled by default. A positive +integer, such as 4096, opts into a periodic interval. Zero, negative values, +booleans, floating-point values, strings, arrays and objects are rejected with +the nested field path. The explicit setting survives normalization, provenance, +saved configuration and reload. No flat legacy option is added. + +When enabled, the interval counts **one insertion and one eviction as separate +mutations**. A rebuild runs after the complete retained-sample update transaction +and resets the mutation counter to zero. With an explicit interval of 4096, a +capacity of 64, an initially empty store, and no earlier recovery rebuild, the +first rebuild occurs after 2,080 accepted observations: 64 initial insertions, +then 2,016 insert/evict pairs. Further rebuilds occur every 2,048 accepted +observations while the store stays full. This counts accepted observations per +store, not prediction queries or wall-clock time. + +There is at most one periodic rebuild after an update transaction. A full-store +insert/evict pair can cross an odd interval by one mutation; it still produces +one rebuild and a reset to zero. Rejected observations and spatial rebucketing +do not advance the counter. A batch-fit fallback alone does not reset it; +rebuilding the statistics does. The setting is fixed for each store when the +model is constructed. + +The default and explicit Python `None` both disable only the periodic schedule. +Numerical recovery rebuilds and conservative batch fallbacks remain enabled: + +```python +config = ForwardPassPerfModelConfig( + model="Qwen/Qwen3-32B", + system="h200_sxm", + backend="vllm", + worker_type="decode", + estimation_mode="fpm_regression", + estimator_config={"fpm_regression": {"fit": {"rebuild_interval": None}}}, +) +model = RustForwardPassPerfModel.best_available(config) +saved = model.diagnostics()["provenance"]["config"] +restored = RustForwardPassPerfModel.best_available(saved) +``` + +Raw JSON uses `"fit": {"rebuild_interval": null}` at the same nested path. +In Rust the control is the public `RegressionFitConfig::rebuild_interval` field: + +```rust +use aisimulate_core::{ + BackendKind, EstimationMode, ForwardPassPerfModel, ForwardPassPerfModelConfig, + ForwardPassWorkerType, +}; + +let mut config = ForwardPassPerfModelConfig::new( + "Qwen/Qwen3-32B", "h200_sxm", BackendKind::Vllm, ForwardPassWorkerType::Decode, +); +config.estimation_mode = EstimationMode::FpmRegression; +// Opt into periodic rebuilding; the default None retains numerical recovery. +config.estimator_config.fpm_regression.fit.rebuild_interval = Some(4096); +let model = ForwardPassPerfModel::best_available(config)?; +``` + ### DCP self-benchmark profiles `dcp` is an optional recorded decode-context-parallel dimension within TP. diff --git a/docs/fpm-recursive-regression.md b/docs/fpm-recursive-regression.md new file mode 100644 index 0000000000..dbfa9b3a27 --- /dev/null +++ b/docs/fpm-recursive-regression.md @@ -0,0 +1,509 @@ + + +# Recursive regression in the forward-pass performance model + +AISimulate updates its regression fit when a retained observation is inserted +or evicted. It maintains small, centered sufficient statistics and solves the +existing standardized, nonnegative linear regression from those statistics. +The usual fitting path therefore does not scan the retained observations. +Numerically ambiguous cases still use a fresh batch fit on the exact retained +data. + +The model has two features and one fitted plane per workload store. Recursive +updates make maintaining that plane cheaper; sample selection and the plane's +ability to describe a workload remain separate concerns. Periodic statistics +rebuilding is disabled by default (`None` in Rust/Python, `null` in JSON). +Numerical recovery and conservative batch fallbacks remain enabled. + +## What one store learns + +An observation contains two finite, nonnegative raw features and a positive, +finite observed latency in milliseconds: + +$$ +x = [x_0, x_1], \qquad y = \text{observed latency in ms}. +$$ + +`RegressionIterationFeatures` computes $x_0$ as critical attention work: a +maximum over attention-DP ranks after composing each rank's attention score. +It computes $x_1$ as global FFN/MoE work: a weighted sum of fresh prefill tokens +and/or decode requests across ranks, according to the worker's role. In +particular, the aggregated role uses the sum of prefill tokens **plus decode +requests**. Feature weights are fixed at construction. The +[feature reference](../python/aisimulate/docs/fpm/aic-fpm-regression-design.md#reference-symbols-and-features) +gives the full formulas. + +Tuning obtains $y$ from the maximum finite, positive rank wall time, converted +from seconds to milliseconds. Prediction uses scheduled workload fields and +does not consume the observed wall time. + +Each logical store owns its retained samples, statistics, fit, and rebuild +clock: + +| Worker role | Independent regression stores | +|---|---| +| Prefill | `pure_prefill` | +| Decode | `pure_decode` | +| Aggregated | `pure_decode`, `contains_locally_mixed`, `cross_rank_aggregated`, `pure_prefill` | + +A query selects its workload store. Training another store does not change +its answer. The default capacity of 64 applies to **each store**; the default +four-by-four retention grid inside a store does not create 16 fitted planes. + +## The fitted plane and standardization + +For the current $n$ retained observations, the population mean and scale of +feature $j$ are + +$$ +\mu_j = \frac{1}{n}\sum_i x_{ij}, \qquad +s_j = \sqrt{\frac{1}{n}\sum_i (x_{ij}-\mu_j)^2}. +$$ + +The denominator is $n$, not $n-1$. A feature is active only when its scale is +finite and + +$$ +s_j > 10^{-12}\max(1, |\mu_j|). +$$ + +An active feature is transformed to $u_j=(x_j-\mu_j)/s_j$. An inactive feature +has transformed value zero and fitted coefficient zero. The target stays in +milliseconds; it is not divided by a target standard deviation. + +The fitted prediction is + +$$ +\widehat y = a + b_0 u_0 + b_1 u_1, \qquad b_j \ge 0. +$$ + +The intercept $a$ is free, including negative values. The slopes are +constrained to be nonnegative. For active axes, the corresponding raw plane +has coefficients + +$$ +\theta_j = b_j/s_j, \qquad +\theta_{\mathrm{intercept}} = a - \sum_{j\text{ active}} \theta_j\mu_j. +$$ + +Thus nonnegative standardized slopes also mean nonnegative raw slopes. +The stored intercept $a$ is the value at the retained feature means, not +necessarily the raw plane's value at the origin. Each fit keeps the means, +scales, and active-axis flags that produced its coefficients, so prediction +uses a consistent snapshot. + +## Statistics that support insertion and eviction + +Combine the two raw features and target into a three-dimensional vector: + +$$ +z_i = [x_{i0}, x_{i1}, y_i]^\mathsf T. +$$ + +`RecursiveFit` maintains a count $n$, mean vector $m$, and symmetric centered +scatter matrix $C$: + +$$ +m = \frac{1}{n}\sum_i z_i, \qquad +C = \sum_i (z_i-m)(z_i-m)^\mathsf T + = \begin{bmatrix} C_{xx} & C_{xy} \\ C_{xy}^\mathsf T & C_{yy}\end{bmatrix}. +$$ + +Here $C_{xx}$ is two-by-two, $C_{xy}$ has two entries, and $C_{yy}$ is a scalar. +They contain everything needed to standardize the features, build the normal +equations, and score a candidate's residual error. Keeping centered scatter +avoids routinely subtracting two large raw quantities such as +$\sum x_i^2-n\mu^2$ to recover a small variance. Floating-point roundoff is +still possible, especially during removal. + +### Adding one observation + +Let $z$ be the incoming observation and $\delta=z-m$ before insertion. Then + +$$ +\begin{aligned} +n' &= n+1,\\ +m' &= m + \frac{\delta}{n+1},\\ +C' &= C + \frac{n}{n+1}\delta\delta^\mathsf T. +\end{aligned} +$$ + +To see the scatter update, recenter the old points at $m'$. Their centered +deviations sum to zero, leaving $C+n(m-m')(m-m')^\mathsf T$. Adding the new +point's deviation $(z-m')(z-m')^\mathsf T$ gives the factor $n/(n+1)$ above. +The first insertion initializes $m=z$ and $C=0$. + +The code uses the equivalent Welford expression +$\delta(z-m')^\mathsf T$. For off-diagonal entries it averages the two +coordinate orders and mirrors the result to retain symmetry. + +### Removing one retained observation + +For $n>1$, let $z$ be the actual evicted observation and $\delta=z-m$ before +removal. Inverting the insertion identity gives + +$$ +\begin{aligned} +n' &= n-1,\\ +m' &= m - \frac{\delta}{n-1},\\ +C' &= C - \frac{n}{n-1}\delta\delta^\mathsf T. +\end{aligned} +$$ + +Removing the last observation clears the count, mean, and scatter. A full +store performs insertion followed by removal; the removal formula uses the +mean **after insertion**. The sampler returns the actual evicted value, so +the update does not assume that the globally oldest observation was removed. + +For example, start with $(x_0,x_1,y)=(0,0,1)$ and $(2,0,5)$: + +$$ +n=2,\quad m=[1,0,3]^\mathsf T,\quad +C=\begin{bmatrix}2&0&4\\0&0&0\\4&0&8\end{bmatrix}. +$$ + +Inserting $(4,0,9)$ gives $m'=[2,0,5]^\mathsf T$ and +$C'=\left[\begin{smallmatrix}8&0&16\\0&0&0\\16&0&32\end{smallmatrix}\right]$. +Removing $(0,0,1)$ then gives mean $[3,0,7]^\mathsf T$ and the original +scatter matrix. The remaining points have feature scale $s_0=1$, so their +line is $7+2(x_0-3)=1+2x_0$. This illustrates the statistics; with only two +retained points, the default minimum of five still prevents a ready fit. + +## Solving the constrained fit from the statistics + +For active axes, let $D=\operatorname{diag}(s_j)$ and define + +$$ +G=D^{-1}C_{xx}D^{-1}, \qquad h=D^{-1}C_{xy}. +$$ + +These are the standardized feature Gram matrix and feature-target cross +products. Centered features have zero sum, so the intercept separates from +the slopes. For a selected subset of fitted axes $S$, the normal equations are + +$$ +\begin{bmatrix}n&0\\0&G_{SS}\end{bmatrix} +\begin{bmatrix}a\\b_S\end{bmatrix} += +\begin{bmatrix}n m_y\\h_S\end{bmatrix}. +$$ + +In exact arithmetic, the free intercept is $a=m_y$. Axes outside $S$ have +coefficient zero. With two active features the implementation considers four +faces of the nonnegative constraint: no slopes, only the first slope, only +the second slope, and both slopes. With one active feature there are two. +It solves each candidate, discards candidates with negative fitted slopes, +and chooses the smallest residual sum of squares (SSE). It does not simply +clip a negative unconstrained coefficient to zero: the remaining coefficients +must be refitted on that face. + +For a hand-derived example, take all nine combinations $x_0,x_1\in\{0,1,2\}$ +and labels $y=20-2x_0+3x_1$. The independent feature columns make the +nonnegative optimum $\widehat y=18+3x_1$: the forbidden negative term is +replaced by its mean $-2$. This plane predicts 30 at $(100,4)$, regardless of +the first coordinate. A production test anchors this behavior independently +of the batch comparator. + +### Ridge is a singular-solve fallback + +Each candidate first attempts an unregularized solve. Only a failed solve +retries with a penalty on its fitted slopes: + +$$ +\left(H+\lambda\operatorname{diag}(0,1,\ldots,1)\right)c=r, +\qquad +\lambda=\texttt{singular\_ridge\_scale}\, +\max\!\left(1,\sum_k |H_{kk}|\right). +$$ + +$H$ is that candidate's normal-equation matrix, including the intercept +entry $n$. `singular_ridge_scale` defaults to $10^{-9}$. The intercept remains +unpenalized, and the penalty applies in standardized feature coordinates. +The configured scale also reaches every batch fallback. + +This preserves the existing algorithm: some candidates can use ridge and +others can use ordinary least squares, and all are ranked by **unpenalized +SSE**. It is not a single always-regularized ridge objective. The small linear +solver treats a pivot below $10^{-12}$ in absolute magnitude as singular. + +### Scoring without another pass over the observations + +For a candidate $a,b$, with omitted slopes filled with zeros, expansion of +$\sum_i(y_i-a-b^\mathsf T u_i)^2$ yields + +$$ +\operatorname{SSE} +=C_{yy}-2b^\mathsf T h+b^\mathsf T G b+n(a-m_y)^2. +$$ + +The centered cross terms disappear because $\sum_i u_i=0$ and +$\sum_i(y_i-m_y)=0$. This avoids a residual scan for each face. Its terms can +nearly cancel for a good fit, so the implementation checks negative/nonfinite +scores and close candidate scores before trusting the selected face. + +## Numerical recovery, readiness, and prediction + +The recurrence and fresh batch fitting are equivalent in exact arithmetic; +floating-point accumulation orders differ. The implementation retains the +original batch fitter to handle decisions that are sensitive to that +difference, including: + +- feature spread close to the active-axis threshold; +- nearly collinear active features, including $1-\rho^2\le10^{-8}$; +- fitted slopes close to the zero constraint boundary; +- failed or nonfinite solves, invalid SSE, and numerically tied candidate SSE. + +The batch fallback rebuilds standardized observations from the retained raw +rows and scores direct residuals, using the same ridge configuration. It +returns a fit for the current update without replacing the incremental +statistics or resetting their mutation clock. + +A **statistics rebuild** is different: it reaccumulates the means and scatter +from the retained raw rows and resets the clock. Recovery rebuilds run when +the statistics become nonfinite, a scatter diagonal becomes negative, or a +downdate removes almost all previously represented variance. These guards +remain active when periodic rebuilding is disabled. They reduce numerical +risk but do not constitute a bound on roundoff for every possible stream. + +A store is ready only if it has enough retained observations (five by default) +and at least one varying feature. After selecting a fit, it also requires a +positive slope, except for the existing low-observation case $n\le d$, where +$d$ is the number of varying axes. This count is not the rank of the feature +matrix. With the default minimum, an intercept-only solution is unready. A +store can lose readiness after an eviction. + +For a valid nonempty query, a ready store transforms features using its fit +snapshot. If transformation and evaluation are finite, it returns the prediction +floored at $10^{-6}$ ms; otherwise it returns no prediction. The floor +is applied after fitting; training SSE uses raw predictions. A valid iteration +with no scheduled work returns zero through the outer model and adds no +observation. A cold workload store returns no prediction even if another store +is ready. + +## Configuring periodic rebuilding + +The canonical field is +`estimator_config.fpm_regression.fit.rebuild_interval`. Rust owns its default +and validation: + +| Setting | Behavior | +|---|---| +| Omitted, Rust/Python `None`, or JSON `null` | No periodic statistics rebuild | +| Positive integer $I$ | Rebuild after at least $I$ retained-sample mutations | +| Zero, negative, noninteger, or boolean | Field-specific configuration error | + +An accepted insertion counts as one mutation and an actual eviction counts +as another. The check runs once after the complete transaction. A full-store +replacement therefore contributes two mutations; it can cross an odd interval +by one before the clock resets to zero. Interval 1 produces one rebuild per +accepted transaction, including when that transaction inserts and evicts. +Rejected input and spatial rebucketing do not advance the clock. Each actual +statistics rebuild resets it; an ordinary batch fallback does not. The clock +saturates instead of overflowing when left running indefinitely. + +For an explicitly configured interval of 4096 and capacity 64, with no earlier +recovery rebuild, the first periodic rebuild is after 2080 accepted inputs: +64 initial insertions plus 2016 insert/evict pairs. Later rebuilds occur every +2048 accepted inputs while the store stays full. This schedule is per store, +not per prediction query or elapsed time. + +Construct regression through the same public API as other estimators: + +```python +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + +config = ForwardPassPerfModelConfig( + model="Qwen/Qwen3-32B", + system="h200_sxm", + backend="vllm", + worker_type="aggregated", + estimation_mode="fpm_regression", + estimator_config={ + "fpm_regression": { + "sampling": {"bins_per_axis": [4, 4], "max_observations": 64}, + "fit": {"rebuild_interval": None}, + }, + }, +) +model = RustForwardPassPerfModel.best_available(config) +``` + +The explicit `None` documents the choice; omitting the field has the same +default behavior. Set it to `4096`, for example, to enable periodic rebuilding +on a newly constructed model. JSON uses `null` rather than Python's `None`. +In Rust, set the same typed field before construction: + +```rust +use aisimulate_core::{ + BackendKind, EstimationMode, ForwardPassPerfModel, ForwardPassPerfModelConfig, + ForwardPassWorkerType, +}; + +let mut config = ForwardPassPerfModelConfig::new( + "Qwen/Qwen3-32B", "h200_sxm", BackendKind::Vllm, ForwardPassWorkerType::Decode, +); +config.estimation_mode = EstimationMode::FpmRegression; +config.estimator_config.fpm_regression.fit.rebuild_interval = Some(4096); +let model = ForwardPassPerfModel::best_available(config)?; +``` + +The setting is fixed at construction and preserved in resolved configuration +and provenance. Reloading a saved explicit interval keeps that interval; +changing the default does not overwrite it. There is no separate flat legacy +option for this control. See the [canonical API](core-api.md#choosing-a-forward-pass-api) +for selection, diagnostics, migration, and saved-configuration behavior. + +## Sample retention and total tuning cost + +The sampler uses $\log(1+x_j)$ coordinates to choose retention cells. Fitting +uses the raw $x_j$. Grid bounds expand with incoming observations and do not +shrink on eviction. If the store exceeds its capacity, it evicts the oldest +observation in a most-populated cell. Hash-map ordering can affect tied +choices. Each retained observation has weight one; an evicted observation +has weight zero. This is spatially balanced retention, not a FIFO window or +exponential forgetting. + +Let $M$ be retained capacity, $B$ the bucket-map storage scanned to select an +eviction, and $k$ the population of the selected cell. With two fixed features: + +| Work in a tuning update | Cost | +|---|---| +| Feature validation and reduction | Linear in the supplied rank count | +| Centered insertion/removal and small constrained fit | Constant in $M$ on the healthy path | +| Finding a fattest retention cell | Scan of the bucket map, $O(B)$ | +| Removing the oldest row from that cell's vector | $O(k)$, up to $O(M)$ | +| Rebucketing after dynamic bounds expand | $O(M)$ | +| Statistics rebuild or batch-fit fallback | $O(M)$ | +| Stored observations | $O(M)$, plus constant-sized fitting state | + +The healthy fitter requests retained rows lazily, so it avoids both a scan +and a temporary observation copy. Periodic rebuilding, when enabled, adds +an amortized term proportional to $M/I$ for interval $I$; a full store uses +approximately two mutations per accepted input. Recovery and batch-fallback +frequency depend on the data. + +Consequently, increasing capacity can make fitting cheaper than repeated +batch refits while still increasing memory, eviction, rebucketing, and +recovery work. It also changes the training sample distribution. One plane +per workload store can average over a curved or changing latency surface; +retaining more of the landscape does not turn it into a local tangent model +or a nonlinear model. Capacity remains an accuracy and retention choice as +well as a resource choice. + +## Relationship to classic inverse-matrix RLS + +For fixed features $\phi=[1,x_0,x_1]^\mathsf T$, ordinary RLS often maintains +the inverse normal matrix $P=(\sum_i\phi_i\phi_i^\mathsf T)^{-1}$. For one +insertion, its rank-one update is + +$$ +P'=P-\frac{P\phi\phi^\mathsf T P}{1+\phi^\mathsf T P\phi}. +$$ + +That is useful when the feature coordinates and fitting objective stay fixed. +Our implementation uses recursive **sufficient statistics** instead. Means +and scales change with the retained observations, active axes can change, +the nonnegative slope face can change, and ridge is conditional on a failed +candidate solve. Recomputing the small constrained solve from updated scatter +preserves those behaviors directly, without maintaining and transforming an +inverse for each possible constraint face. It still performs recursive +least-squares estimation; it does not use a forgetting factor or the classic +inverse-matrix coefficient update. + +## Speed measurements + +The following CPU measurements time one **retained-sample update**: +`BucketedRegression::add_observation` from already computed raw features. +This includes validation, log-grid coordinates, retention/eviction, recursive +statistics, and fitting, including any recovery or batch fallback the data +triggers. It excludes rank-level feature extraction, Python/JSON overhead, +prediction queries, and GPU execution. It is not an end-to-end serving speedup. + +The baseline uses the original full-bucket update from `8e763d4` and the same +sampler and batch-fitting arithmetic. The recursive path uses the production +fitter with `rebuild_interval=None`, including all numerical guards. Both +paths use a four-by-four retention grid, minimum five observations, and +`singular_ridge_scale=1e-9`. + +Times are medians of 21 paired trials; brackets show the minimum–maximum +trial range. Speedup is the batch median divided by the recursive median. + +| Input data | Capacity | Batch update (µs) | Recursive update (µs) | Speedup | +|---|---:|---:|---:|---:| +| Toy | 64 | 7.078 [7.006–7.136] | 0.812 [0.788–0.823] | 8.7× | +| H200 / MiniMax-M2.7 / prefill | 64 | 7.093 [6.995–7.299] | 0.818 [0.781–0.840] | 8.7× | +| B200 / MiniMax-M3 NVFP4 / decode | 64 | 6.909 [6.815–7.027] | 0.807 [0.780–0.831] | 8.6× | +| Toy | 8,192 | 736.844 [717.487–750.102] | 1.120 [1.034–1.257] | 657.9× | +| H200 / MiniMax-M2.7 / prefill | 8,192 | 738.917 [710.320–774.915] | 1.770 [1.706–1.889] | 417.6× | +| B200 / MiniMax-M3 NVFP4 / decode | 8,192 | 737.600 [720.991–754.895] | 2.288 [2.195–2.434] | 322.3× | + +Outside the timed region, **6,150 fits and 50,835,900 raw/floored +prediction comparisons passed**, with identical retained rows and no readiness, +active-feature, or slope-constraint failures. The audit covers the warm state +and every update in the 1,024-row tail for each case. It checks retained points +and five boundary/extrapolation queries, including 1.5 times the feature maxima. +Maximum absolute prediction difference was **2.724e-10 ms**, +within `1e-8 + 1e-8 * abs(batch prediction)` ms. Maximum scaled raw-coefficient +and training-SSE differences were 1.065e-12 and +1.549e-14, respectively, using `abs(batch-recursive) / max(1, abs(batch))`. +This tests recovery of the batch fit on these rows, not an unbounded-stream +roundoff guarantee. + +Measured on 2026-09-18 on an Apple M5 Pro, macOS 26.6.2, arm64, +18 logical CPUs, with Rust 1.97.1. The isolated benchmark compiles the production +regression and sampling modules with `opt-level=3`, `target-cpu=native`, and +`codegen-units=1`; timing is single-threaded. The recursive arithmetic is +unchanged from `14ac37e4`. +These are fresh measurements of the production path with the disabled setting, +not the earlier prototype timings. + +Each case warms the model with 8,192 observations, then measures a distinct +1,024-observation tail. Each method starts from a clone of the same retained +state, preserving identical sampling decisions. There are 21 paired trials, +alternating which method runs first, after one discarded trial per method. +Each trial repeats that tail from the same warm state 16 times at capacity 64 +or four times at capacity 8,192. +Cloning, setup, and destruction are outside the timer; every updated fit is +passed to `black_box`. This measures repeated warm windows, not continuous +long-stream behavior. Reported ranges describe timing variation for the fixed +retained state; new processes can resolve sampling ties differently. + +Inputs are shuffled with seed `0xA15120260917` for this CPU comparison: + +- **Toy:** independent log-uniform features and + $y=\max(0.01,3+3\times10^{-5}x_0+0.025x_1+\epsilon)$ ms, with + $\epsilon\sim N(0,0.03^2)$ and generator seed `0xA15170`. +- **H200 prefill:** MiniMax-M2.7, vLLM 0.25.1, TP4/DP1; 9,505 measured + offline shape-grid rows. Feature coordinates use the balanced-request proxy. + Source `fpm_forward_perf.parquet` at AISimulate revision + `b790939feb2fc681c5f1c3b55c23e4388a56dd6e`, collector cell + `fpm-574efc8a51db0a0b`. +- **B200 decode:** MiniMax-M3 NVFP4, vLLM 0.28.0, TP4/DP1; the 310,492 + pure-decode rows from recorded serving iterations in + `nvidia/aisimulate-fpm-dataset` revision + `5487a4599a7fbc012c07bcd3699754bdf4a8bef7`. + +Each case uses 9,216 unique source row IDs for warmup and its measured tail. +The shuffle and repeated trials are a timing protocol; they do not measure +chronological adaptation or held-out prediction accuracy. + +## Source map and validation anchors + +| Source | Responsibility | +|---|---| +| [model.rs](../crates/core/src/perfmodel/fpm/model.rs) | `RegressionIterationFeatures`, workload routing, independent `RegressionStores`, observation targets | +| [samples.rs](../crates/core/src/perfmodel/fpm/samples.rs) | Retention grid and `SampleInsertion` carrying the actual eviction | +| [regression.rs](../crates/core/src/perfmodel/fpm/regression.rs) | `BucketedRegression`, original batch fit, small linear solvers, prediction floor | +| [regression/recursive.rs](../crates/core/src/perfmodel/fpm/regression/recursive.rs) | Centered updates, constrained fit from statistics, guards, rebuild clock | +| [estimator.rs](../crates/core/src/perfmodel/fpm/estimator.rs) | Rust-owned `RegressionFitConfig` defaults and validation | + +Tests compare recursive and batch fits on identical retained rows, including +evictions, changing scales, inactive axes, and nondefault ridge strengths. +Independent plane and nonnegative-boundary examples check expected numerical +answers. Scheduling tests cover custom intervals, disabled periodic rebuilding, +recovery, clone independence, and counter saturation. These checks complement +measured workload validation; they do not establish a universal bound on +long-run floating-point drift. diff --git a/python/aisimulate/docs/fpm/aic-fpm-regression-design.md b/python/aisimulate/docs/fpm/aic-fpm-regression-design.md index 0e5cacdb0b..7ee0542e91 100644 --- a/python/aisimulate/docs/fpm/aic-fpm-regression-design.md +++ b/python/aisimulate/docs/fpm/aic-fpm-regression-design.md @@ -72,7 +72,7 @@ formulas and explain how each rank contributes. |---|---| | Gym caller | Infer each worker's offline role, create one predictor per `worker_id`, and preserve replay order | | AISim predictor | Validate the iteration, extract features, select the workload bucket, predict, and tune | -| Workload bucket | Retain its own observations, rebuild its fit, and track its readiness | +| Workload bucket | Retain its own observations, update its fit recursively, and track its readiness | For example, the caller could own the following predictors: @@ -170,7 +170,8 @@ iteration: an error. 4. Record the outcome and score an available estimate against observed latency. 5. Tune the same worker's predictor with the observation. AISim uses the same - classifier, retains the sample in the selected bucket, and refits that bucket. + classifier, retains the sample in the selected bucket, and updates that + bucket's statistics and fit. The current iteration's latency is introduced at the tuning step. Prediction does not consume it. An unavailable prediction can still be followed by a @@ -556,9 +557,12 @@ cap. With the default `sampling.bins_per_axis=[4, 4]`, the grid is $4\times4$. B choose which observations survive; they are not local predictors and are not queried during estimation. -After insertion and eviction, the fit is rebuilt from the retained **raw** -features. For each axis, population mean and standard deviation are computed -with stable Welford accumulation: +Each accepted insertion adds the **raw** features and observed latency to +centered sufficient statistics; an actual eviction removes that observation's +contribution. After both operations, the fitter derives population means and +standard deviations from the updated statistics and solves the small +nonnegative regression. Healthy fitting updates do not scan the retained +observations. For each feature axis, the statistics represent: $$ \mu_j=\frac1n\sum_i x_{ij},\qquad @@ -566,6 +570,15 @@ $$ z_{ij}=\frac{x_{ij}-\mu_j}{\sigma_j}. $$ +The fitter reads retained observations only for a statistics rebuild or a +batch-fit fallback. Periodic rebuilds are disabled by default +(`estimator_config.fpm_regression.fit.rebuild_interval=None` in Python, +`null` in JSON); a positive integer enables them. Numerical recovery rebuilds +and batch fallbacks remain enabled regardless of that setting. Retention-grid +rebucketing is separate from fitting and may still scan stored observations. +See the [recursive regression walkthrough](../../../../docs/fpm-recursive-regression.md) +for the update equations, rebuild-counter semantics, and numerical guards. + An axis is inactive when $$ @@ -587,5 +600,5 @@ fit remains unready. The existing special handling for an explicitly configured lower minimum is preserved when there are too few observations to identify the active slopes. -Each workload bucket runs this pipeline independently. Refitting one bucket +Each workload bucket runs this pipeline independently. Updating one bucket leaves the other buckets and other predictor instances unchanged. diff --git a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py index 612be02bdc..148e3edffd 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py +++ b/python/aisimulate/src/aisimulate_core/sdk/rust_engine_step.py @@ -101,7 +101,15 @@ class RustEngineUnsupportedError(RuntimeError): @dataclass(frozen=True) class ForwardPassPerfModelConfig: - """Canonical immutable identity and selection policy for the estimator.""" + """Canonical immutable identity and selection policy for the estimator. + + Estimator controls pass through to Rust unchanged. For regression, + ``estimator_config["fpm_regression"]["fit"]["rebuild_interval"]`` accepts + a positive integer mutation count or ``None`` to disable periodic + statistics rebuilding. Omitting it uses the Rust default of ``None``; + numerical recovery and batch fallbacks remain enabled. An explicit + positive interval, such as 4096, opts into periodic rebuilding. + """ model: str system: str @@ -222,6 +230,14 @@ class RustForwardPassPerfModel: own fit and retention state. ``max_observations`` (default ``64``) and ``min_observations`` (default ``5``) apply independently to each store. + Regression updates centered sufficient statistics as retained samples are + inserted or evicted. The canonical ``fpm_regression.fit.rebuild_interval`` + control counts one mutation per insertion and one per eviction. A scheduled + rebuild occurs after the complete update transaction, then resets its + counter to zero. The Rust default is ``None``, which disables only scheduled + rebuilds, preserving numerical recovery and batch fallbacks. Set a positive + interval, such as 4096 mutations, to enable scheduled rebuilds. + Queued request fields are accepted for schema compatibility but ignored by this AIC forward-pass model. ``estimate_forward_pass_time_ms()`` treats FPM as a workload descriptor: scheduled request fields are used, while @@ -265,6 +281,8 @@ def best_available(cls, config: ForwardPassPerfModelConfig | Mapping[str, Any]) Auto searches op_level, fpm_interpolation, then fpm_regression even with fallback_policy=deny. Explicit modes default to strict selection. Native correction retains its existing workload feature space. + Nested estimator controls, including an explicit ``None`` rebuild + interval, are preserved in resolved configuration and provenance. """ import aisimulate_core diff --git a/python/aisimulate/tests/cross_package/test_core_public_api.py b/python/aisimulate/tests/cross_package/test_core_public_api.py index b7b0c8f66a..e7ee928096 100644 --- a/python/aisimulate/tests/cross_package/test_core_public_api.py +++ b/python/aisimulate/tests/cross_package/test_core_public_api.py @@ -236,6 +236,69 @@ def test_raw_fpm_binding_regression_round_trip(worker_type: str) -> None: assert prediction is not None and prediction > 0.0 +@pytest.mark.parametrize( + ("fit", "expected_interval"), + [ + ({}, None), + ({"rebuild_interval": 31}, 31), + ({"rebuild_interval": 4096}, 4096), + ({"rebuild_interval": None}, None), + ], +) +def test_raw_canonical_regression_rebuild_interval_normalizes_and_reloads(fit, expected_interval): + raw = aisimulate_core.RustForwardPassPerfModel + request = { + "model": "test/model", + "system": "test", + "backend": "vllm", + "worker_type": "decode", + "estimation_mode": "fpm_regression", + "estimator_config": {"fpm_regression": {"fit": fit}}, + } + normalized = json.loads(raw.normalize_config(json.dumps(request))) + assert normalized["estimator_config"]["fpm_regression"]["fit"] == { + "kind": "standardized_nnls", + "singular_ridge_scale": 1e-9, + "rebuild_interval": expected_interval, + } + model = raw.best_available(json.dumps(normalized)) + resolved = json.loads(model.diagnostics())["provenance"]["config"] + assert resolved["estimator_config"] == normalized["estimator_config"] + restored = raw.best_available(json.dumps(resolved)) + assert json.loads(restored.diagnostics())["provenance"]["config"] == resolved + + +@pytest.mark.parametrize("entrypoint", ["normalize_config", "best_available"]) +@pytest.mark.parametrize("invalid", [0, -1, True, False, 1.5, 4096.0, "4096", [], {}]) +def test_raw_canonical_regression_rebuild_interval_rejects_invalid_values_with_path(entrypoint, invalid): + request = { + "model": "test/model", + "system": "test", + "backend": "vllm", + "worker_type": "decode", + "estimation_mode": "auto", + "fallback_policy": "allow", + "estimator_config": {"fpm_regression": {"fit": {"rebuild_interval": invalid}}}, + } + # Invalid configuration fails before automatic selection or fallback. + with pytest.raises(ValueError, match=r"estimator_config\.fpm_regression\.fit\.rebuild_interval"): + getattr(aisimulate_core.RustForwardPassPerfModel, entrypoint)(json.dumps(request)) + + +def test_legacy_options_inherit_rust_rebuild_default_without_a_new_flat_control(): + migrated = json.loads(aisimulate_core.RustForwardPassPerfModel.legacy_estimator_config("{}")) + assert migrated["fpm_regression"]["fit"]["rebuild_interval"] is None + assert migrated["fpm_regression"]["sampling"] == {"bins_per_axis": [4, 4], "max_observations": 64} + assert "rebuild_interval" not in sdk.ForwardPassPerfOptions.__dataclass_fields__ + assert "regression_rebuild_interval" not in sdk.ForwardPassPerfOptions.__dataclass_fields__ + assert ( + sdk.ForwardPassPerfModelConfig( + model="test/model", system="test", backend="vllm", worker_type="decode" + ).estimator_config + == {} + ) + + def test_raw_fpm_binding_validates_regression_weights() -> None: with pytest.raises(ValueError, match="regression_attention_kv_weight"): _raw_regression_model( diff --git a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py index 973975149e..3bdf44073a 100644 --- a/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py +++ b/python/aisimulate/tests/unit/sdk/test_rust_engine_step.py @@ -920,7 +920,8 @@ def iteration(kind: str, scale: int) -> list[dict[str, object]]: assert model.estimate_forward_pass_time_ms(iteration(cold_kind, 3)) is None -def test_forward_pass_constructor_passes_complete_typed_request(monkeypatch) -> None: +@pytest.mark.parametrize("rebuild_interval", [23, 4096, None]) +def test_forward_pass_constructor_passes_complete_typed_request(monkeypatch, rebuild_interval) -> None: import aisimulate_core calls = [] @@ -941,7 +942,10 @@ def best_available(config_json): "fallback_policy": "deny", "estimator_config": { "features": {"attention_kv_weight": 2.0, "prefill_attention_pair_weight": 3.0, "ffn_token_weight": 4.0}, - "fpm_regression": {"sampling": {"bins_per_axis": [4, 16], "max_observations": 128}}, + "fpm_regression": { + "sampling": {"bins_per_axis": [4, 16], "max_observations": 128}, + "fit": {"rebuild_interval": rebuild_interval}, + }, "correction": {"enabled": False}, }, } @@ -1018,6 +1022,111 @@ def test_forward_pass_config_preserves_existing_positional_arguments(tmp_path: P assert json.loads(json.dumps(pinned.to_dict()))["moe_kernel_source"] == source +@pytest.mark.parametrize("as_mapping", [False, True]) +@pytest.mark.parametrize( + ("fit", "expected_interval"), + [ + ({}, None), + ({"rebuild_interval": 17}, 17), + ({"rebuild_interval": 4096}, 4096), + ({"rebuild_interval": None}, None), + ], +) +def test_canonical_regression_rebuild_interval_survives_saved_config(as_mapping, fit, expected_interval): + from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + config = ForwardPassPerfModelConfig( + model="test/model", + system="test", + backend="vllm", + worker_type="decode", + estimation_mode="fpm_regression", + estimator_config={"fpm_regression": {"fit": fit}}, + ) + payload = config.to_dict() if as_mapping else config + model = RustForwardPassPerfModel.best_available(payload) + provenance = model.diagnostics()["provenance"] + resolved = provenance["config"] + assert resolved["estimator_config"]["fpm_regression"]["fit"]["rebuild_interval"] == expected_interval + assert resolved["estimator_config"]["fpm_regression"]["sampling"] == { + "bins_per_axis": [4, 4], + "max_observations": 64, + } + assert resolved["worker_type"] == "decode" + assert resolved["estimation_mode"] == "fpm_regression" + assert resolved["fallback_policy"] == "deny" + assert config.estimator_config == {"fpm_regression": {"fit": fit}} + if not fit: + # Default expansion belongs to Rust, not the Python request object. + assert "rebuild_interval" not in config.estimator_config["fpm_regression"]["fit"] + + saved = json.loads(json.dumps(resolved)) + restored = RustForwardPassPerfModel.best_available(saved) + assert restored.diagnostics()["provenance"]["config"] == resolved + assert restored.regression_store_diagnostics() == [ + {"workload_kind": "pure_decode", "ready": False, "retained_observations": 0} + ] + + +@pytest.mark.parametrize("invalid", [0, -1, True, False, 1.5, 4096.0, "4096", [], {}, float("nan"), float("inf")]) +def test_canonical_regression_rebuild_interval_rejects_invalid_facade_values_with_path(invalid): + from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + config = ForwardPassPerfModelConfig( + model="test/model", + system="test", + backend="vllm", + worker_type="decode", + estimation_mode="auto", + fallback_policy="allow", + estimator_config={"fpm_regression": {"fit": {"rebuild_interval": invalid}}}, + ) + # The dataclass facade must preserve invalid values for Rust to reject; + # neither Python defaults nor estimator fallback may hide the error. + with pytest.raises(ValueError, match=r"estimator_config\.fpm_regression\.fit\.rebuild_interval"): + RustForwardPassPerfModel.best_available(config) + + +@pytest.mark.parametrize("fit", [{}, {"rebuild_interval": 1}, {"rebuild_interval": 7}, {"rebuild_interval": None}]) +def test_canonical_regression_updates_after_evictions_with_default_or_custom_rebuild_schedule(fit): + from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel + + config = ForwardPassPerfModelConfig( + model="test/model", + system="test", + backend="vllm", + worker_type="decode", + estimation_mode="fpm_regression", + estimator_config={ + "fpm_regression": { + "sampling": {"bins_per_axis": [2, 2], "max_observations": 8}, + "fit": fit, + } + }, + ) + model = RustForwardPassPerfModel.best_available(config) + + def sample(index): + batch = index % 7 + 1 + kv = index * index * 13 + 11 + # Hand-specified affine surface in Decode's KV and batch features. + latency_ms = 2.0 + 0.001 * kv + 0.1 * batch + return { + "version": 1, + "wall_time": latency_ms / 1000, + "scheduled_requests": {"num_decode_requests": batch, "sum_decode_kv_tokens": kv}, + } + + query = sample(11) + assert model.estimate_forward_pass_time_ms(query) is None + for index in range(1, 81): + model.tune_with_fpms(sample(index)) + assert model.regression_store_diagnostics() == [ + {"workload_kind": "pure_decode", "ready": True, "retained_observations": 8} + ] + assert model.estimate_forward_pass_time_ms(query) == pytest.approx(query["wall_time"] * 1000, rel=1e-8) + + def _supported_fpm_config() -> dict[str, object]: return { "schema_version": 1,