aisimulate_core/perfmodel/fpm/
options.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Tuning controls for the forward-pass perf model.
5
6use serde::{Deserialize, Serialize};
7
8use crate::AicError;
9
10use super::samples::integer_sqrt;
11
12pub(crate) const DEFAULT_MAX_OBSERVATIONS: usize = 64;
13pub(crate) const DEFAULT_MIN_OBSERVATIONS: usize = 5;
14pub(crate) const DEFAULT_MIN_FASTER_CORRECTION_FACTOR: f64 = 0.5;
15pub(crate) const DEFAULT_MAX_SLOWER_CORRECTION_FACTOR: f64 = 2.0;
16pub(crate) const DEFAULT_BUCKET_COUNT: usize = 16;
17pub(crate) const DEFAULT_MAX_NUM_TOKENS: u32 = 8192;
18pub(crate) const DEFAULT_MAX_BATCH_SIZE: u32 = 512;
19pub(crate) const DEFAULT_MAX_KV_TOKENS: u32 = 2_000_000;
20
21/// In-memory tuning controls for `ForwardPassPerfModel`.
22///
23/// The defaults retain a bounded sliding sample set, wait for enough
24/// observations before predicting from learned data, bucket observations by
25/// workload kind, and bound native correction factors to `[0.5, 2.0]`.
26#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
27pub struct ForwardPassPerfOptions {
28    /// Maximum retained observations across all buckets for each inferred workload kind.
29    #[serde(default = "default_max_observations")]
30    pub max_observations: usize,
31    /// Minimum retained observations required before a regression fit or native
32    /// correction is used for an inferred workload kind.
33    #[serde(default = "default_min_observations")]
34    pub min_observations: usize,
35    /// Optional absolute lower bound on native correction factors for
36    /// observations faster than the native estimate.
37    ///
38    /// Values must be finite, greater than `0.0`, and at most `1.0`. Defaults
39    /// to `0.5`, limiting learned speedups to `2x`. Setting this to `1.0`
40    /// disables faster corrections; setting it to `None` removes the lower
41    /// bound. Regression fallback does not use this option.
42    #[serde(default = "default_min_faster_correction_factor")]
43    pub min_faster_correction_factor: Option<f64>,
44    /// Optional absolute upper bound on native correction factors for
45    /// observations slower than the native estimate.
46    ///
47    /// Values must be finite and at least `1.0`. Defaults to `2.0`, limiting
48    /// learned slowdowns to `2x`. Setting this to `1.0` disables slower
49    /// corrections; setting it to `None` removes the upper bound. Regression
50    /// fallback does not use this option.
51    #[serde(default = "default_max_slower_correction_factor")]
52    pub max_slower_correction_factor: Option<f64>,
53    /// Target bucket count for workload-specific sample retirement and correction lookup.
54    #[serde(default = "default_bucket_count")]
55    pub bucket_count: usize,
56    /// Upper bound for the `sum_prefill_tokens` correction axis.
57    ///
58    /// Used by prefill and mixed/agg workload kinds. The lower bound is always `0`.
59    #[serde(default = "default_max_num_tokens")]
60    pub max_num_tokens: u32,
61    /// Upper bound for the `num_decode_requests` correction axis.
62    ///
63    /// Used by the decode workload kind. The lower bound is always `0`.
64    #[serde(default = "default_max_batch_size")]
65    pub max_batch_size: u32,
66    /// Upper bound for the `sum_decode_kv_tokens` correction axis.
67    ///
68    /// Used by decode and mixed/agg workload kinds. The lower bound is always `0`.
69    #[serde(default = "default_max_kv_tokens")]
70    pub max_kv_tokens: u32,
71}
72
73impl Default for ForwardPassPerfOptions {
74    fn default() -> Self {
75        Self {
76            max_observations: DEFAULT_MAX_OBSERVATIONS,
77            min_observations: DEFAULT_MIN_OBSERVATIONS,
78            min_faster_correction_factor: default_min_faster_correction_factor(),
79            max_slower_correction_factor: default_max_slower_correction_factor(),
80            bucket_count: DEFAULT_BUCKET_COUNT,
81            max_num_tokens: DEFAULT_MAX_NUM_TOKENS,
82            max_batch_size: DEFAULT_MAX_BATCH_SIZE,
83            max_kv_tokens: DEFAULT_MAX_KV_TOKENS,
84        }
85    }
86}
87
88pub(crate) fn validate_options(options: &ForwardPassPerfOptions) -> Result<(), AicError> {
89    if options.max_observations == 0 {
90        return Err(invalid_perf_options("max_observations must be >= 1"));
91    }
92    if options.min_observations == 0 {
93        return Err(invalid_perf_options("min_observations must be >= 1"));
94    }
95    if let Some(min_faster_correction_factor) = options.min_faster_correction_factor {
96        if !min_faster_correction_factor.is_finite()
97            || min_faster_correction_factor <= 0.0
98            || min_faster_correction_factor > 1.0
99        {
100            return Err(invalid_perf_options(
101                "min_faster_correction_factor must be finite and in (0.0, 1.0]",
102            ));
103        }
104    }
105    if let Some(max_slower_correction_factor) = options.max_slower_correction_factor {
106        if !max_slower_correction_factor.is_finite() || max_slower_correction_factor < 1.0 {
107            return Err(invalid_perf_options(
108                "max_slower_correction_factor must be finite and >= 1.0",
109            ));
110        }
111    }
112    if options.bucket_count == 0 {
113        return Err(invalid_perf_options("bucket_count must be >= 1"));
114    }
115    if options.max_num_tokens == 0 {
116        return Err(invalid_perf_options("max_num_tokens must be >= 1"));
117    }
118    if options.max_batch_size == 0 {
119        return Err(invalid_perf_options("max_batch_size must be >= 1"));
120    }
121    if options.max_kv_tokens == 0 {
122        return Err(invalid_perf_options("max_kv_tokens must be >= 1"));
123    }
124    if options.min_observations > options.max_observations {
125        return Err(invalid_perf_options(
126            "min_observations must be <= max_observations",
127        ));
128    }
129    let sqrt = integer_sqrt(options.bucket_count);
130    if sqrt * sqrt != options.bucket_count {
131        return Err(invalid_perf_options(
132            "bucket_count must be a perfect square",
133        ));
134    }
135    Ok(())
136}
137
138fn invalid_perf_options(message: &str) -> AicError {
139    AicError::InvalidEngineConfig(format!("invalid forward pass perf options: {message}"))
140}
141
142fn default_max_observations() -> usize {
143    DEFAULT_MAX_OBSERVATIONS
144}
145
146fn default_min_observations() -> usize {
147    DEFAULT_MIN_OBSERVATIONS
148}
149
150fn default_min_faster_correction_factor() -> Option<f64> {
151    Some(DEFAULT_MIN_FASTER_CORRECTION_FACTOR)
152}
153
154fn default_max_slower_correction_factor() -> Option<f64> {
155    Some(DEFAULT_MAX_SLOWER_CORRECTION_FACTOR)
156}
157
158fn default_bucket_count() -> usize {
159    DEFAULT_BUCKET_COUNT
160}
161
162fn default_max_num_tokens() -> u32 {
163    DEFAULT_MAX_NUM_TOKENS
164}
165
166fn default_max_batch_size() -> u32 {
167    DEFAULT_MAX_BATCH_SIZE
168}
169
170fn default_max_kv_tokens() -> u32 {
171    DEFAULT_MAX_KV_TOKENS
172}
aisimulate_core/perfmodel/fpm/options.rs

aisimulate_core/perfmodel/fpm/
options.rs
