[Rust Frontend] Add iteration token histogram (#56990)
Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Sep 15, 2026 at 20:12 UTC
78eaaad2c7a7f8c7864d9c9c72772e3f9f2669d0
5 files changed
+90
-14
rust/src/engine-core-client/src/client/imp.rs
+25
-1
@@ -19,7 +19,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti
19
use crate::client::stream::EngineCoreStreamOutput;
20
use crate::client::{AbortCause, AbortRequest};
21
use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
22
-use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
22
+use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder};
23
use crate::protocol::encode_msgpack;
24
use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
25
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
@@ -47,6 +47,7 @@ pub(crate) struct ClientInner {
47
/// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
48
msgpack_zero_copy_threshold: usize,
49
scheduler_stats_recorder: SchedulerStatsRecorder,
50
+ iteration_metrics: BTreeMap<u32, IterationMetricHandles>,
51
request_reg: Mutex<RequestRegistry>,
52
utility_reg: Mutex<UtilityRegistry>,
53
health_error: ArcSwapOption<Error>,
@@ -64,12 +65,21 @@ impl ClientInner {
65
) -> Self {
66
let scheduler_stats_recorder =
67
SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines);
68
+ let iteration_metrics = engines
69
+ .iter()
70
+ .filter_map(|engine| {
71
+ let engine = engine.engine_id.engine_index()?;
72
+ let handles = IterationMetricHandles::new(&METRICS.request, &model_name, engine);
73
+ Some((engine, handles))
74
+ })
75
+ .collect();
76
Self {
77
input_send,
78
handle,
79
model_name,
80
msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
81
scheduler_stats_recorder,
82
+ iteration_metrics,
83
request_reg: Mutex::new(RequestRegistry::new(engines)),
84
utility_reg: Mutex::new(UtilityRegistry::default()),
85
health_error: ArcSwapOption::empty(),
@@ -415,6 +425,8 @@ pub(crate) async fn run_output_dispatcher_loop(
425
426
match outputs {
427
EngineCoreOutputs::RequestBatch(batch) => {
428
+ let has_outputs = !batch.outputs.is_empty();
429
+ let mut iteration_tokens = 0_u64;
430
let senders = inner.take_senders_for_outputs(&batch.outputs);
431
for (output, sender) in batch.outputs.into_iter().zip(senders) {
432
let request_id = output.request_id.clone();
@@ -423,6 +435,12 @@ pub(crate) async fn run_output_dispatcher_loop(
435
continue;
436
};
437
438
+ iteration_tokens += output.new_token_ids.len() as u64;
439
+ if let Some(prefill) = &output.prefill_stats {
440
+ // The engine emits prefill_stats once, on the first output.
441
+ iteration_tokens += u64::from(prefill.num_computed_tokens);
442
+ }
443
+
444
let wrapped_output = EngineCoreStreamOutput {
445
engine_index: batch.engine_index,
446
timestamp: batch.timestamp,
@@ -433,6 +451,12 @@ pub(crate) async fn run_output_dispatcher_loop(
451
}
452
}
453
454
+ if has_outputs
455
+ && let Some(handles) = inner.iteration_metrics.get(&batch.engine_index)
456
+ {
457
+ handles.iteration_tokens_total.observe(iteration_tokens as f64);
458
+ }
459
+
460
// The sender for normally-finished requests should have already been removed
461
// from the registry when their final output was dispatched
462
// above. This serves as a safety net to capture any
rust/src/engine-core-client/src/metrics.rs
+19
-2
@@ -8,8 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
8
use vllm_metrics::{
9
EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames,
10
LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily,
11
- MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
12
- WaitingReasonLabels,
11
+ MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics,
12
+ U64Counter, U64Gauge, WaitingReasonLabels,
13
};
14
15
use crate::protocol::stats::{
@@ -20,6 +20,23 @@ use crate::transport::ConnectedEngine;
20
const WAITING_REASON_CAPACITY: &str = "capacity";
21
const WAITING_REASON_DEFERRED: &str = "deferred";
22
23
+/// Cached output-batch metric handles for one model and engine index.
24
+pub(crate) struct IterationMetricHandles {
25
+ pub iteration_tokens_total: HistogramMetric,
26
+}
27
+
28
+impl IterationMetricHandles {
29
+ pub(crate) fn new(metrics: &RequestMetrics, model_name: &str, engine: u32) -> Self {
30
+ let labels = EngineLabels {
31
+ model_name: model_name.to_string(),
32
+ engine,
33
+ };
34
+ Self {
35
+ iteration_tokens_total: metrics.iteration_tokens_total.get_or_create_owned(&labels),
36
+ }
37
+ }
38
+}
39
+
40
/// Cached scheduler-stats metric handles for all engines connected to one
41
/// frontend client.
42
pub(crate) struct SchedulerStatsRecorder {
rust/src/metrics/src/request.rs
+21
-6
@@ -21,6 +21,9 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [
21
480.0, 960.0, 1920.0, 7680.0,
22
];
23
const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0];
24
+const ITERATION_TOKENS_BUCKETS: [f64; 13] = [
25
+ 1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0,
26
+];
27
28
fn build_1_2_5_buckets(max_value: u32) -> Vec<f64> {
29
let mut buckets = Vec::new();
@@ -66,6 +69,10 @@ fn request_params_n_histogram() -> Histogram {
69
Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied())
70
}
71
72
+fn iteration_tokens_histogram() -> Histogram {
73
+ Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied())
74
+}
75
+
76
#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
77
pub struct FinishedReasonLabels {
78
pub model_name: String,
@@ -83,8 +90,12 @@ pub struct PromptTokenSourceLabels {
90
pub(crate) type FinishedReasonCounterFamily = Family<FinishedReasonLabels, U64Counter>;
91
pub(crate) type PromptTokenSourceCounterFamily = Family<PromptTokenSourceLabels, U64Counter>;
92
86
-/// Request-lifecycle Prometheus families exported from the `llm` layer.
93
+/// Request and output-batch Prometheus families.
94
pub struct RequestMetrics {
95
+ /// Computed prompt and generated tokens per nonempty engine output batch,
96
+ /// recorded by the engine-core client for requests tracked by this frontend.
97
+ pub iteration_tokens_total: HistogramFamily,
98
+
99
// Request-derived counters.
100
pub num_preemptions: Family<EngineLabels, U64Counter>,
101
pub prompt_tokens: Family<EngineLabels, U64Counter>,
@@ -92,11 +103,6 @@ pub struct RequestMetrics {
103
pub prompt_tokens_cached: Family<EngineLabels, U64Counter>,
104
pub generation_tokens: Family<EngineLabels, U64Counter>,
105
95
- // We intentionally don't support iteration-level histograms for now, since it seems to make
96
- // more sense if the engine maintains these metrics and frontend simply forwards.
97
- //
98
- // pub iteration_tokens_total: HistogramFamily,
99
-
106
// Request lifecycle counters and histograms.
107
pub request_success: FinishedReasonCounterFamily,
108
pub request_prompt_tokens: HistogramFamily,
@@ -118,6 +124,14 @@ pub struct RequestMetrics {
124
impl RequestMetrics {
125
/// Register the request-oriented metric families into the shared registry.
126
pub(crate) fn register(registry: &mut Registry) -> Self {
127
+ let iteration_tokens_total =
128
+ Family::new_with_constructor(iteration_tokens_histogram as fn() -> Histogram);
129
+ registry.register(
130
+ "vllm:iteration_tokens_total",
131
+ "Histogram of computed prompt and generated tokens per nonempty engine output batch.",
132
+ iteration_tokens_total.clone(),
133
+ );
134
+
135
// Request-derived counters.
136
let num_preemptions = Family::default();
137
registry.register(
@@ -276,6 +290,7 @@ impl RequestMetrics {
290
);
291
292
Self {
293
+ iteration_tokens_total,
294
num_preemptions,
295
prompt_tokens,
296
prompt_tokens_by_source,
rust/src/mock-engine/src/engine.rs
+14
-5
@@ -19,6 +19,7 @@ use vllm_engine_core_client::protocol::output::{
19
UtilityCallOutput,
20
};
21
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
22
+use vllm_engine_core_client::protocol::stats::PrefillStats;
23
use vllm_engine_core_client::protocol::utility::{
24
EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope,
25
};
@@ -187,6 +188,11 @@ impl ActiveRequest {
188
189
/// Advance this request by one mock engine step.
190
fn step(&mut self, opt: &Opt) -> EngineCoreOutput {
191
+ let prefill_stats = (self.generated == 0).then(|| PrefillStats {
192
+ num_prompt_tokens: self.prompt_len as u32,
193
+ num_computed_tokens: self.prompt_len as u32,
194
+ ..Default::default()
195
+ });
196
let remaining = self.max_tokens - self.generated;
197
let chunk_len = remaining.min(opt.output_token_chunk_size);
198
let mut new_token_ids = Vec::with_capacity(chunk_len);
@@ -196,11 +202,14 @@ impl ActiveRequest {
202
self.generated += chunk_len;
203
204
let finished = self.generated >= self.max_tokens;
199
- request_output(
200
- self.request_id.clone(),
201
- new_token_ids,
202
- finished.then_some(EngineCoreFinishReason::Length),
203
- )
205
+ EngineCoreOutput {
206
+ prefill_stats,
207
+ ..request_output(
208
+ self.request_id.clone(),
209
+ new_token_ids,
210
+ finished.then_some(EngineCoreFinishReason::Length),
211
+ )
212
+ }
213
}
214
}
215
rust/src/mock-engine/src/tests.rs
+11
@@ -10,6 +10,7 @@ use tokio_util::sync::CancellationToken;
10
use vllm_engine_core_client::protocol::output::EngineCoreFinishReason;
11
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
12
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
13
+use vllm_engine_core_client::protocol::stats::PrefillStats;
14
use vllm_engine_core_client::test_utils::IpcNamespace;
15
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
16
@@ -119,12 +120,22 @@ async fn chunk_size_one_outputs_one_token_per_update() {
120
let first = stream.next().await.expect("first").expect("first ok");
121
assert_eq!(first.new_token_ids.len(), 1);
122
assert_eq!(first.finish_reason, None);
123
+ assert_eq!(
124
+ first.prefill_stats,
125
+ Some(PrefillStats {
126
+ num_prompt_tokens: 3,
127
+ num_computed_tokens: 3,
128
+ ..Default::default()
129
+ })
130
+ );
131
let second = stream.next().await.expect("second").expect("second ok");
132
assert_eq!(second.new_token_ids.len(), 1);
133
assert_eq!(second.finish_reason, None);
134
+ assert!(second.prefill_stats.is_none());
135
let third = stream.next().await.expect("third").expect("third ok");
136
assert_eq!(third.new_token_ids.len(), 1);
137
assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length));
138
+ assert!(third.prefill_stats.is_none());
139
assert!(stream.next().await.is_none());
140
141
shutdown_mock(client, shutdown, task).await;