aisimulate_core/perfmodel/fpm/
metrics.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Forward-pass metrics telemetry types.
5//!
6//! These mirror Dynamo's scheduler `ForwardPassMetrics` v1 wire schema. AIC owns
7//! this Rust copy so it does not depend on Dynamo crates. The forward-pass perf
8//! model consumes the scheduled portion as estimator input and uses `wall_time`
9//! as the tuning target.
10
11use serde::{Deserialize, Serialize};
12
13use crate::AicError;
14
15/// Wire-schema version for [`ForwardPassMetrics`] (Dynamo's scheduler
16/// telemetry). The FPM input is rejected with [`AicError::UnsupportedSchemaVersion`]
17/// when `version` does not match.
18pub const FPM_VERSION: u32 = 1;
19
20/// Metrics for requests scheduled in one forward-pass iteration.
21///
22/// This mirrors Dynamo ForwardPassMetrics v1 scheduled request telemetry.
23/// AIC owns this Rust copy so AIC does not depend on Dynamo crates.
24#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
25pub struct ScheduledRequestMetrics {
26    /// Number of prefill requests, including new requests and chunked-prefill
27    /// continuations.
28    #[serde(default)]
29    pub num_prefill_requests: u32,
30    /// Total tokens freshly computed for prefill requests in this iteration.
31    #[serde(default)]
32    pub sum_prefill_tokens: u32,
33    /// Population variance of total prompt lengths across prefill requests.
34    #[serde(default)]
35    pub var_prefill_length: f64,
36    /// Total KV tokens read for prefill requests, including prefix cache hits
37    /// and previously computed chunks.
38    #[serde(default)]
39    pub sum_prefill_kv_tokens: u32,
40    /// Number of decode requests.
41    #[serde(default)]
42    pub num_decode_requests: u32,
43    /// Total KV context length across decode requests.
44    #[serde(default)]
45    pub sum_decode_kv_tokens: u32,
46    /// Population variance of KV context lengths across decode requests.
47    #[serde(default)]
48    pub var_decode_kv_tokens: f64,
49}
50
51/// Metrics for requests queued but not scheduled in one iteration.
52#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
53pub struct QueuedRequestMetrics {
54    #[serde(default)]
55    pub num_prefill_requests: u32,
56    #[serde(default)]
57    pub sum_prefill_tokens: u32,
58    #[serde(default)]
59    pub var_prefill_length: f64,
60    #[serde(default)]
61    pub num_decode_requests: u32,
62    #[serde(default)]
63    pub sum_decode_kv_tokens: u32,
64    #[serde(default)]
65    pub var_decode_kv_tokens: f64,
66}
67
68/// Per-iteration forward-pass metrics.
69///
70/// In Dynamo this struct is telemetry emitted after an engine iteration. In
71/// AIC the scheduled portion is also the estimator input. `wall_time` and
72/// `queued_requests` are accepted for schema parity but ignored by the
73/// latency estimator (the forward-pass perf model uses `wall_time` as the
74/// tuning target, never as an estimation input).
75#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
76pub struct ForwardPassMetrics {
77    #[serde(default = "default_fpm_version")]
78    pub version: u32,
79    #[serde(default)]
80    pub worker_id: String,
81    #[serde(default)]
82    pub dp_rank: u32,
83    #[serde(default)]
84    pub counter_id: u64,
85    #[serde(default)]
86    pub wall_time: f64,
87    #[serde(default)]
88    pub scheduled_requests: ScheduledRequestMetrics,
89    #[serde(default)]
90    pub queued_requests: QueuedRequestMetrics,
91}
92
93impl Default for ForwardPassMetrics {
94    fn default() -> Self {
95        Self {
96            version: FPM_VERSION,
97            worker_id: String::new(),
98            dp_rank: 0,
99            counter_id: 0,
100            wall_time: 0.0,
101            scheduled_requests: ScheduledRequestMetrics::default(),
102            queued_requests: QueuedRequestMetrics::default(),
103        }
104    }
105}
106
107/// Validate one [`ForwardPassMetrics`] entry. Rejects mismatched schema
108/// versions and inconsistent scheduled-request sums (token sums without a
109/// matching request count). Shared by the forward-pass perf model and
110/// [`crate::perfmodel::engine::Engine::forward_pass_time_ms`].
111pub(crate) fn validate_forward_pass_metrics(metrics: &ForwardPassMetrics) -> Result<(), AicError> {
112    if metrics.version != FPM_VERSION {
113        return Err(AicError::UnsupportedSchemaVersion {
114            kind: "ForwardPassMetrics",
115            got: metrics.version,
116            expected: FPM_VERSION,
117        });
118    }
119    let scheduled = &metrics.scheduled_requests;
120    if scheduled.num_prefill_requests == 0
121        && (scheduled.sum_prefill_tokens > 0 || scheduled.sum_prefill_kv_tokens > 0)
122    {
123        return Err(AicError::InvalidForwardPassMetrics(
124            "prefill token sums require num_prefill_requests > 0".to_string(),
125        ));
126    }
127    if scheduled.num_decode_requests == 0 && scheduled.sum_decode_kv_tokens > 0 {
128        return Err(AicError::InvalidForwardPassMetrics(
129            "decode KV token sum requires num_decode_requests > 0".to_string(),
130        ));
131    }
132    Ok(())
133}
134
135fn default_fpm_version() -> u32 {
136    FPM_VERSION
137}
aisimulate_core/perfmodel/fpm/metrics.rs

aisimulate_core/perfmodel/fpm/
metrics.rs
