dynamo_mocker/common/
perf_model.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Performance model for timing simulations in the mocker.
5//!
6//! This module provides two timing models:
7//! 1. Polynomial: Hardcoded polynomial formulas (default, backward compatible)
8//! 2. Interpolated: Grid-based interpolation from profiler data (loaded from NPZ files)
9
10use anyhow::{Context, Result};
11use ndarray::{Array1, Array2};
12use ndarray_interp::InterpolateError;
13use ndarray_interp::interp1d::{Interp1DBuilder, Linear};
14use ndarray_interp::interp2d::{Bilinear, Interp2DBuilder};
15use std::path::Path;
16use std::sync::Arc;
17
18/// Trait to abstract over 1D interpolation for prefill timing
19pub trait PrefillInterpolator: Send + Sync {
20    fn interp(&self, x: f64) -> Result<f64, InterpolateError>;
21}
22
23/// Trait to abstract over 2D interpolation for decode timing
24pub trait DecodeInterpolator: Send + Sync {
25    fn interp(&self, x: f64, y: f64) -> Result<f64, InterpolateError>;
26}
27
28/// Callback trait for direct AIC SDK calls.
29/// Implementors call the Rust AIC core API.
30pub trait AicCallback: Send + Sync {
31    /// Predict prefill latency in ms.
32    /// Parameters: (batch_size, effective_isl, prefix)
33    fn predict_prefill(
34        &self,
35        batch_size: usize,
36        effective_isl: usize,
37        prefix: usize,
38    ) -> Result<f64>;
39
40    /// Predict decode (generation) latency in ms.
41    /// Parameters: (batch_size, isl, osl)
42    fn predict_decode(&self, batch_size: usize, isl: usize, osl: usize) -> Result<f64>;
43}
44
45/// Wrapper to implement PrefillInterpolator for the concrete Interp1D type
46struct PrefillInterp1D {
47    inner: ndarray_interp::interp1d::Interp1D<
48        ndarray::OwnedRepr<f64>,
49        ndarray::OwnedRepr<f64>,
50        ndarray::Ix1,
51        Linear,
52    >,
53}
54
55impl PrefillInterpolator for PrefillInterp1D {
56    fn interp(&self, x: f64) -> Result<f64, InterpolateError> {
57        self.inner.interp_scalar(x)
58    }
59}
60
61/// Wrapper to implement DecodeInterpolator for the concrete Interp2D type
62struct DecodeInterp2D {
63    inner: ndarray_interp::interp2d::Interp2D<
64        ndarray::OwnedRepr<f64>,
65        ndarray::OwnedRepr<f64>,
66        ndarray::OwnedRepr<f64>,
67        ndarray::Ix2,
68        Bilinear,
69    >,
70}
71
72impl DecodeInterpolator for DecodeInterp2D {
73    fn interp(&self, x: f64, y: f64) -> Result<f64, InterpolateError> {
74        self.inner.interp_scalar(x, y)
75    }
76}
77
78/// Performance model for predicting prefill and decode timing
79#[derive(Default)]
80pub enum PerfModel {
81    /// Default polynomial-based model using hardcoded formulas
82    #[default]
83    Polynomial,
84    /// Interpolation-based model using profiler data
85    /// Decode axes: (active_kv_tokens, context_length)
86    Interpolated {
87        prefill_interp: Arc<dyn PrefillInterpolator>,
88        decode_interp: Arc<dyn DecodeInterpolator>,
89    },
90    /// AI Configurator SDK calls through the configured callback.
91    /// Passes the reduced prefill inputs (batch_size, effective_isl, prefix).
92    Aiconfigurator { callback: Arc<dyn AicCallback> },
93}
94
95impl Clone for PerfModel {
96    fn clone(&self) -> Self {
97        match self {
98            PerfModel::Polynomial => PerfModel::Polynomial,
99            PerfModel::Interpolated {
100                prefill_interp,
101                decode_interp,
102            } => PerfModel::Interpolated {
103                prefill_interp: Arc::clone(prefill_interp),
104                decode_interp: Arc::clone(decode_interp),
105            },
106            PerfModel::Aiconfigurator { callback } => PerfModel::Aiconfigurator {
107                callback: Arc::clone(callback),
108            },
109        }
110    }
111}
112
113impl std::fmt::Debug for PerfModel {
114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115        match self {
116            PerfModel::Polynomial => write!(f, "PerfModel::Polynomial"),
117            PerfModel::Interpolated { .. } => write!(f, "PerfModel::Interpolated {{ .. }}"),
118            PerfModel::Aiconfigurator { .. } => write!(f, "PerfModel::Aiconfigurator"),
119        }
120    }
121}
122
123impl PerfModel {
124    /// Load performance model from NPZ file
125    ///
126    /// Expected arrays in NPZ file:
127    /// - prefill_isl: 1D array of input sequence lengths
128    /// - prefill_ttft_ms: 1D array of time to first token in milliseconds
129    /// - decode_active_kv_tokens: 1D array of active KV token counts
130    /// - decode_context_length: 1D array of context lengths
131    /// - decode_itl: 2D array of inter-token latencies in milliseconds
132    pub fn from_npz(path: &Path) -> Result<Self> {
133        use ndarray_npy::NpzReader;
134        use std::fs::File;
135
136        tracing::info!("Loading performance model from NPZ file: {:?}", path);
137
138        let file =
139            File::open(path).with_context(|| format!("Failed to open NPZ file: {:?}", path))?;
140
141        let mut npz = NpzReader::new(file)
142            .with_context(|| format!("Failed to create NPZ reader for: {:?}", path))?;
143
144        // Load prefill arrays
145        let prefill_isl: Array1<f64> = npz
146            .by_name("prefill_isl")
147            .with_context(|| "Failed to load prefill_isl from NPZ")?;
148        let prefill_ttft_ms: Array1<f64> = npz
149            .by_name("prefill_ttft_ms")
150            .with_context(|| "Failed to load prefill_ttft_ms from NPZ")?;
151
152        // Load decode arrays
153        let decode_active_kv_tokens: Array1<f64> = npz
154            .by_name("decode_active_kv_tokens")
155            .with_context(|| "Failed to load decode_active_kv_tokens from NPZ")?;
156        let decode_context_length: Array1<f64> = npz
157            .by_name("decode_context_length")
158            .with_context(|| "Failed to load decode_context_length from NPZ")?;
159        let decode_itl: Array2<f64> = npz
160            .by_name("decode_itl")
161            .with_context(|| "Failed to load decode_itl from NPZ")?;
162
163        // Validate dimensions
164        if prefill_isl.len() != prefill_ttft_ms.len() {
165            anyhow::bail!(
166                "Prefill array length mismatch: isl={}, ttft={}",
167                prefill_isl.len(),
168                prefill_ttft_ms.len()
169            );
170        }
171
172        if decode_itl.nrows() != decode_active_kv_tokens.len()
173            || decode_itl.ncols() != decode_context_length.len()
174        {
175            anyhow::bail!(
176                "Decode array dimension mismatch: itl shape=({}, {}), active_kv={}, context={}",
177                decode_itl.nrows(),
178                decode_itl.ncols(),
179                decode_active_kv_tokens.len(),
180                decode_context_length.len()
181            );
182        }
183
184        tracing::info!(
185            "Loaded performance model: prefill_points={}, decode_grid={}x{}",
186            prefill_isl.len(),
187            decode_itl.nrows(),
188            decode_itl.ncols()
189        );
190
191        // Build interpolators once during loading
192        let prefill_interp = Interp1DBuilder::new(prefill_ttft_ms)
193            .x(prefill_isl)
194            .strategy(Linear::new().extrapolate(true))
195            .build()
196            .with_context(|| "Failed to build prefill interpolator")?;
197
198        let decode_interp = Interp2DBuilder::new(decode_itl)
199            .x(decode_active_kv_tokens)
200            .y(decode_context_length)
201            .strategy(Bilinear::new().extrapolate(true))
202            .build()
203            .with_context(|| "Failed to build decode interpolator")?;
204
205        Ok(PerfModel::Interpolated {
206            prefill_interp: Arc::new(PrefillInterp1D {
207                inner: prefill_interp,
208            }),
209            decode_interp: Arc::new(DecodeInterp2D {
210                inner: decode_interp,
211            }),
212        })
213    }
214
215    /// Create an Aiconfigurator perf model from a callback.
216    pub fn from_aic_callback(callback: Arc<dyn AicCallback>) -> Self {
217        PerfModel::Aiconfigurator { callback }
218    }
219
220    /// Predict prefill time in milliseconds.
221    ///
222    /// Callers always pass all parameters; each variant uses what it needs:
223    /// - Polynomial/Interpolated: uses total new tokens across the batch
224    ///   (`batch_size * (isl - prefix)`), modeling GPU processing total tokens in parallel
225    /// - Aiconfigurator: passes (batch_size, isl - prefix, prefix) to the AIC SDK
226    pub fn predict_prefill_time(
227        &self,
228        batch_size: usize,
229        isl: usize,
230        prefix: usize,
231    ) -> Result<f64> {
232        let new_tokens_per_req = isl.saturating_sub(prefix);
233        if batch_size == 0 || new_tokens_per_req == 0 {
234            return Ok(0.0);
235        }
236        let time = match self {
237            PerfModel::Polynomial => polynomial_prefill_time(batch_size, new_tokens_per_req),
238            PerfModel::Interpolated { prefill_interp, .. } => {
239                let tokens = (batch_size * new_tokens_per_req) as f64;
240                prefill_interp.interp(tokens).unwrap_or(0.0)
241            }
242            PerfModel::Aiconfigurator { callback } => callback
243                .predict_prefill(batch_size, new_tokens_per_req, prefix)
244                .context("AIC prefill prediction failed")?,
245        };
246        Ok(time.max(0.0))
247    }
248
249    /// Predict decode time in milliseconds.
250    ///
251    /// Callers always pass all parameters; each variant uses what it needs:
252    /// - Polynomial: uses (active_kv_tokens, total_kv_tokens) as utilization
253    /// - Interpolated: uses (active_kv_tokens, context_length)
254    /// - Aiconfigurator: uses (batch_size, context_length)
255    pub fn predict_decode_time(
256        &self,
257        batch_size: usize,
258        active_kv_tokens: usize,
259        context_length: usize,
260        total_kv_tokens: usize,
261    ) -> Result<f64> {
262        if batch_size == 0 {
263            return Ok(0.0);
264        }
265        let time = match self {
266            PerfModel::Polynomial => polynomial_decode_time(active_kv_tokens, total_kv_tokens),
267            PerfModel::Interpolated { decode_interp, .. } => decode_interp
268                .interp(active_kv_tokens as f64, context_length as f64)
269                .unwrap_or(0.0),
270            PerfModel::Aiconfigurator { callback } => callback
271                .predict_decode(batch_size, context_length, 2)
272                .context("AIC decode prediction failed")?,
273        };
274        // Token-emitting decode steps should not collapse onto the same timestamp.
275        let result = time.max(1.0);
276        tracing::trace!(
277            "Decode time prediction: batch_size={batch_size}, active_kv_tokens={active_kv_tokens}, context_length={context_length}, time={result:.2}ms"
278        );
279        Ok(result)
280    }
281}
282
283fn polynomial_prefill_time(batch_size: usize, new_tokens_per_request: usize) -> f64 {
284    // Total tokens across the batch — GPU processes them in parallel.
285    let tokens = (batch_size * new_tokens_per_request) as f64;
286    4.209989e-07 * tokens.powi(2) + 1.518344e-02 * tokens + 1.650142e+01
287}
288
289fn polynomial_decode_time(active_kv_tokens: usize, total_kv_tokens: usize) -> f64 {
290    let active_perc = if total_kv_tokens > 0 {
291        active_kv_tokens as f64 / total_kv_tokens as f64
292    } else {
293        tracing::warn!("Total KV tokens is 0, using 1.0 as capacity");
294        1.0
295    };
296    -25.74 * active_perc.powi(2) + 54.01 * active_perc + 5.74
297}
298
299#[cfg(test)]
300mod tests {
301    use super::{AicCallback, PerfModel};
302    use std::sync::Arc;
303
304    struct EchoBatchCallback;
305
306    impl AicCallback for EchoBatchCallback {
307        fn predict_prefill(
308            &self,
309            batch_size: usize,
310            _effective_isl: usize,
311            _prefix: usize,
312        ) -> anyhow::Result<f64> {
313            Ok(batch_size as f64)
314        }
315
316        fn predict_decode(
317            &self,
318            batch_size: usize,
319            _isl: usize,
320            _osl: usize,
321        ) -> anyhow::Result<f64> {
322            Ok(batch_size as f64)
323        }
324    }
325
326    struct FailingCallback;
327
328    impl AicCallback for FailingCallback {
329        fn predict_prefill(
330            &self,
331            _batch_size: usize,
332            _effective_isl: usize,
333            _prefix: usize,
334        ) -> anyhow::Result<f64> {
335            anyhow::bail!("missing AIC prefill point")
336        }
337
338        fn predict_decode(
339            &self,
340            _batch_size: usize,
341            _isl: usize,
342            _osl: usize,
343        ) -> anyhow::Result<f64> {
344            anyhow::bail!("missing AIC decode point")
345        }
346    }
347
348    #[test]
349    fn fully_cached_prompt_skips_prefill() {
350        assert_eq!(
351            PerfModel::default()
352                .predict_prefill_time(1, 128, 128)
353                .unwrap(),
354            0.0
355        );
356    }
357
358    #[test]
359    fn aic_forwards_scheduler_local_batch() {
360        let model = PerfModel::from_aic_callback(Arc::new(EchoBatchCallback));
361
362        assert_eq!(model.predict_prefill_time(7, 128, 0).unwrap(), 7.0);
363        assert_eq!(model.predict_decode_time(9, 0, 128, 0).unwrap(), 9.0);
364    }
365
366    #[test]
367    fn aic_prefill_prediction_errors_propagate() {
368        let error = PerfModel::from_aic_callback(Arc::new(FailingCallback))
369            .predict_prefill_time(2, 128, 32)
370            .unwrap_err();
371
372        assert_eq!(error.to_string(), "AIC prefill prediction failed");
373        assert_eq!(error.root_cause().to_string(), "missing AIC prefill point");
374    }
375
376    #[test]
377    fn aic_decode_prediction_errors_propagate() {
378        let error = PerfModel::from_aic_callback(Arc::new(FailingCallback))
379            .predict_decode_time(2, 64, 128, 1024)
380            .unwrap_err();
381
382        assert_eq!(error.to_string(), "AIC decode prediction failed");
383        assert_eq!(error.root_cause().to_string(), "missing AIC decode point");
384    }
385}
dynamo_mocker/common/perf_model.rs

dynamo_mocker/common/
perf_model.rs
