[Rust Frontend] Improve scheduler stats logging parity (#47435)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Jul 2, 2026 at 22:55 UTC
ec0ffaacc8a6437681b5743cfdcc7841ada6b4d7
6 files changed
+521
-30
rust/src/engine-core-client/src/client.rs
+8
@@ -372,6 +372,14 @@ impl EngineCoreClient {
372
self.engines.len()
373
}
374
375
+ /// Return the engine-side indices connected to this client.
376
+ pub fn engine_indices(&self) -> Vec<u32> {
377
+ self.engines
378
+ .iter()
379
+ .map(|engine| engine.engine_id.engine_index().expect("engine id must encode as u16"))
380
+ .collect()
381
+ }
382
+
383
/// Return the engine identities of all engines connected to this client.
384
pub fn engine_identities(&self) -> Vec<&[u8]> {
385
self.engines.iter().map(|engine| &*engine.engine_id).collect()
rust/src/engine-core-client/src/metrics.rs
+13
@@ -84,6 +84,10 @@ pub(crate) fn record_scheduler_stats(
84
.spec_decode_num_accepted_tokens
85
.get_or_create(&labels)
86
.inc_by(spec_decoding_stats.num_accepted_tokens);
87
+ metrics.log_stats.get_or_create(&labels).observe_spec_decode(
88
+ spec_decoding_stats.num_drafts,
89
+ &spec_decoding_stats.num_accepted_tokens_per_pos,
90
+ );
91
92
for (position, accepted_tokens) in
93
spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate()
@@ -119,6 +123,15 @@ pub(crate) fn record_scheduler_stats(
123
.inc_by(perf_stats.num_write_bytes_per_gpu);
124
}
125
126
+ if let Some(cudagraph_stats) = &stats.cudagraph_stats {
127
+ metrics.log_stats.get_or_create(&labels).observe_cudagraph(
128
+ cudagraph_stats.num_unpadded_tokens,
129
+ cudagraph_stats.num_padded_tokens,
130
+ cudagraph_stats.num_paddings,
131
+ &cudagraph_stats.runtime_mode,
132
+ );
133
+ }
134
+
135
// Sampled KV-cache residency histograms.
136
if !stats.kv_cache_eviction_events.is_empty() {
137
let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels);
rust/src/engine-core-client/src/protocol/stats.rs
+2
-2
@@ -141,7 +141,7 @@ pub struct PerfStats {
141
/// Original Python definition:
142
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/compilation/cuda_graph.py#L28-L33>
143
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
144
-pub struct CudagraphStat {
144
+pub struct CudagraphStats {
145
/// Number of real tokens in the captured batch before padding.
146
pub num_unpadded_tokens: u64,
147
/// Number of padded tokens in the captured batch.
@@ -182,7 +182,7 @@ pub struct SchedulerStats {
182
/// Connector-specific KV transfer stats, kept opaque for now.
183
pub kv_connector_stats: Option<BTreeMap<String, OpaqueValue>>,
184
/// CUDA graph runtime stats when graph metrics are enabled.
185
- pub cudagraph_stats: Option<CudagraphStat>,
185
+ pub cudagraph_stats: Option<CudagraphStats>,
186
/// Estimated MFU/performance stats, when enabled.
187
pub perf_stats: Option<PerfStats>,
188
}
rust/src/llm/src/lib.rs
+1
-1
@@ -51,7 +51,7 @@ impl Llm {
51
if enabled {
52
let stats_logger = StatsLogger::start(
53
self.client.model_name().to_string(),
54
- self.client.engine_count(),
54
+ self.client.engine_indices(),
55
);
56
self.stats_logger = Some(stats_logger);
57
} else {
rust/src/llm/src/log_stats.rs
+378
-25
@@ -4,10 +4,12 @@ use std::time::{Duration, Instant};
4
use tokio_util::task::AbortOnDropHandle;
5
use tracing::{debug, info};
6
use vllm_metrics::{
7
- EngineLabels, F64Gauge, METRICS, PromptTokenSourceLabels, U64Counter, U64Gauge,
7
+ EngineLabels, F64Gauge, METRICS, PromptTokenSourceLabels, SchedulerLogStatsAccumulator,
8
+ SchedulerLogStatsInterval, U64Counter, U64Gauge, WaitingReasonLabels,
9
};
10
11
const LOG_STATS_INTERVAL: Duration = Duration::from_secs(10);
12
+const WAITING_REASON_DEFERRED: &str = "deferred";
13
14
/// Cached, cloned metric handles for one engine. Each clone shares the same
15
/// underlying `Arc<Atomic*>` as the prometheus `Family` entry, so reads go
@@ -18,20 +20,58 @@ struct EngineMetrics {
20
generation_tokens: U64Counter,
21
prefix_cache_queries: U64Counter,
22
prefix_cache_hits: U64Counter,
23
+ external_prefix_cache_queries: U64Counter,
24
+ external_prefix_cache_hits: U64Counter,
25
+ num_preemptions: U64Counter,
26
+ spec_decode_num_drafts: U64Counter,
27
+ spec_decode_num_draft_tokens: U64Counter,
28
+ spec_decode_num_accepted_tokens: U64Counter,
29
+ estimated_flops_per_gpu: U64Counter,
30
+ estimated_read_bytes_per_gpu: U64Counter,
31
+ estimated_write_bytes_per_gpu: U64Counter,
32
+ log_stats: SchedulerLogStatsAccumulator,
33
34
// Gauges for instantaneous scheduler state.
35
scheduler_running: U64Gauge,
36
scheduler_waiting: U64Gauge,
37
+ scheduler_deferred: U64Gauge,
38
kv_cache_usage: F64Gauge,
39
}
40
41
/// Accumulated snapshot values from the last logging interval, used to compute
42
/// deltas.
43
+#[derive(Default)]
44
struct CounterSnapshot {
45
prompt_tokens: u64,
46
generation_tokens: u64,
47
prefix_cache_queries: u64,
48
prefix_cache_hits: u64,
49
+ external_prefix_cache_queries: u64,
50
+ external_prefix_cache_hits: u64,
51
+ num_preemptions: u64,
52
+ spec_decode_num_drafts: u64,
53
+ spec_decode_num_draft_tokens: u64,
54
+ spec_decode_num_accepted_tokens: u64,
55
+ estimated_flops_per_gpu: u64,
56
+ estimated_read_bytes_per_gpu: u64,
57
+ estimated_write_bytes_per_gpu: u64,
58
+}
59
+
60
+/// Derived spec-decoding values for one logging interval.
61
+struct SpecDecodingLogStats {
62
+ mean_acceptance_length: f64,
63
+ accepted_throughput: f64,
64
+ draft_throughput: f64,
65
+ accepted_tokens: u64,
66
+ draft_tokens: u64,
67
+ per_position_acceptance_rates: Vec<f64>,
68
+ draft_acceptance_rate: f64,
69
+}
70
+
71
+/// Derived MFU values for one logging interval.
72
+struct MfuLogStats {
73
+ tflops_per_gpu: f64,
74
+ gbps_per_gpu: f64,
75
}
76
77
/// Periodic stats logger that mirrors Python vLLM's `LoggingStatLogger`.
@@ -46,18 +86,20 @@ pub(crate) struct StatsLogger {
86
87
impl StatsLogger {
88
/// Start the background stats logging task.
49
- pub(crate) fn start(model_name: String, engine_count: usize) -> Self {
89
+ pub(crate) fn start(model_name: String, engine_indices: Vec<u32>) -> Self {
90
let task = AbortOnDropHandle::new(tokio::spawn(async move {
51
- run_stats_logger(model_name, engine_count).await;
91
+ run_stats_logger(model_name, engine_indices).await;
92
}));
93
Self { _task: task }
94
}
95
}
96
97
/// Resolve and clone all metric handles once so the hot path is lock-free.
58
-fn resolve_engine_metrics(model_name: &str, engine_count: usize) -> Vec<EngineMetrics> {
98
+fn resolve_engine_metrics(model_name: &str, engine_indices: &[u32]) -> Vec<EngineMetrics> {
99
let m = &METRICS;
60
- (0..engine_count as u32)
100
+ engine_indices
101
+ .iter()
102
+ .copied()
103
.map(|engine| {
104
let el = EngineLabels {
105
model_name: model_name.to_string(),
@@ -68,6 +110,11 @@ fn resolve_engine_metrics(model_name: &str, engine_count: usize) -> Vec<EngineMe
110
engine,
111
source: "local_compute",
112
};
113
+ let deferred = WaitingReasonLabels {
114
+ model_name: model_name.to_string(),
115
+ engine,
116
+ reason: WAITING_REASON_DEFERRED,
117
+ };
118
EngineMetrics {
119
// Use "local_compute" source for prompt throughput (excludes
120
// cached/transferred tokens), matching Python's
@@ -76,16 +123,51 @@ fn resolve_engine_metrics(model_name: &str, engine_count: usize) -> Vec<EngineMe
123
generation_tokens: m.request.generation_tokens.get_or_create_owned(&el),
124
prefix_cache_queries: m.scheduler.prefix_cache_queries.get_or_create_owned(&el),
125
prefix_cache_hits: m.scheduler.prefix_cache_hits.get_or_create_owned(&el),
126
+ external_prefix_cache_queries: m
127
+ .scheduler
128
+ .external_prefix_cache_queries
129
+ .get_or_create_owned(&el),
130
+ external_prefix_cache_hits: m
131
+ .scheduler
132
+ .external_prefix_cache_hits
133
+ .get_or_create_owned(&el),
134
+ num_preemptions: m.request.num_preemptions.get_or_create_owned(&el),
135
+ spec_decode_num_drafts: m.scheduler.spec_decode_num_drafts.get_or_create_owned(&el),
136
+ spec_decode_num_draft_tokens: m
137
+ .scheduler
138
+ .spec_decode_num_draft_tokens
139
+ .get_or_create_owned(&el),
140
+ spec_decode_num_accepted_tokens: m
141
+ .scheduler
142
+ .spec_decode_num_accepted_tokens
143
+ .get_or_create_owned(&el),
144
+ estimated_flops_per_gpu: m
145
+ .scheduler
146
+ .estimated_flops_per_gpu
147
+ .get_or_create_owned(&el),
148
+ estimated_read_bytes_per_gpu: m
149
+ .scheduler
150
+ .estimated_read_bytes_per_gpu
151
+ .get_or_create_owned(&el),
152
+ estimated_write_bytes_per_gpu: m
153
+ .scheduler
154
+ .estimated_write_bytes_per_gpu
155
+ .get_or_create_owned(&el),
156
+ log_stats: m.scheduler.log_stats.get_or_create_owned(&el),
157
scheduler_running: m.scheduler.scheduler_running.get_or_create_owned(&el),
158
scheduler_waiting: m.scheduler.scheduler_waiting.get_or_create_owned(&el),
159
+ scheduler_deferred: m
160
+ .scheduler
161
+ .scheduler_waiting_by_reason
162
+ .get_or_create_owned(&deferred),
163
kv_cache_usage: m.scheduler.kv_cache_usage.get_or_create_owned(&el),
164
}
165
})
166
.collect()
167
}
168
87
-async fn run_stats_logger(model_name: String, engine_count: usize) {
88
- let engines = resolve_engine_metrics(&model_name, engine_count);
169
+async fn run_stats_logger(model_name: String, engine_indices: Vec<u32>) {
170
+ let engines = resolve_engine_metrics(&model_name, &engine_indices);
171
172
let mut interval = tokio::time::interval(LOG_STATS_INTERVAL);
173
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
@@ -109,6 +191,7 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
191
}
192
193
let curr = read_counters(&engines);
194
+ let raw_log_stats = drain_scheduler_log_stats(&engines);
195
196
let prompt_throughput =
197
curr.prompt_tokens.wrapping_sub(prev.prompt_tokens) as f64 / elapsed;
@@ -121,17 +204,37 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
204
&& last_prompt_throughput == 0.0
205
&& last_generation_throughput == 0.0;
206
207
+ /// Emit one stats line at DEBUG while idle and INFO while active.
208
+ macro_rules! log_stats_line {
209
+ ($($arg:tt)*) => {
210
+ if is_idle {
211
+ debug!($($arg)*);
212
+ } else {
213
+ info!($($arg)*);
214
+ }
215
+ };
216
+ }
217
+
218
// Read scheduler gauges (aggregate across engines).
219
let (num_running, num_waiting, kv_cache_usage) = read_scheduler_gauges(&engines);
220
+ let num_deferred = read_deferred_waiting(&engines);
221
+ let delta_preemptions = curr.num_preemptions.wrapping_sub(prev.num_preemptions);
222
223
// Compute prefix cache hit rate over this interval.
224
let delta_queries = curr.prefix_cache_queries.wrapping_sub(prev.prefix_cache_queries);
129
- let prefix_cache_hit_rate = if delta_queries > 0 {
130
- let delta_hits = curr.prefix_cache_hits.wrapping_sub(prev.prefix_cache_hits);
131
- delta_hits as f64 / delta_queries as f64 * 100.0
132
- } else {
133
- 0.0
134
- };
225
+ let delta_hits = curr.prefix_cache_hits.wrapping_sub(prev.prefix_cache_hits);
226
+ let prefix_cache_hit_rate = cache_hit_rate(delta_hits, delta_queries);
227
+
228
+ let delta_external_queries = curr
229
+ .external_prefix_cache_queries
230
+ .wrapping_sub(prev.external_prefix_cache_queries);
231
+ let delta_external_hits =
232
+ curr.external_prefix_cache_hits.wrapping_sub(prev.external_prefix_cache_hits);
233
+ let external_prefix_cache_hit_rate =
234
+ cache_hit_rate(delta_external_hits, delta_external_queries);
235
+ let spec_decoding_log_stats =
236
+ spec_decoding_log_stats(&curr, &prev, elapsed, &raw_log_stats);
237
+ let mfu_log_stats = mfu_log_stats(&curr, &prev, elapsed, engines.len());
238
239
// Build the log line.
240
msg.clear();
@@ -140,17 +243,70 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
243
"Avg prompt tput: {prompt_throughput:.1} toks/s, \
244
Avg generation tput: {generation_throughput:.1} toks/s, \
245
Reqs Running: {num_running}, \
143
- Waiting: {num_waiting}, \
144
- GPU KV cache used: {:.1}%, \
246
+ Waiting: {num_waiting}"
247
+ )
248
+ .unwrap();
249
+ if num_deferred > 0 {
250
+ write!(msg, ", Deferred: {num_deferred} reqs").unwrap();
251
+ }
252
+ if delta_preemptions > 0 {
253
+ write!(msg, ", Preemptions: {delta_preemptions}").unwrap();
254
+ }
255
+ write!(
256
+ msg,
257
+ ", GPU KV cache used: {:.1}%, \
258
Prefix cache hit rate: {prefix_cache_hit_rate:.1}%",
259
kv_cache_usage * 100.0,
260
)
261
.unwrap();
262
+ if delta_external_queries > 0 {
263
+ write!(
264
+ msg,
265
+ ", External prefix cache hit rate: {external_prefix_cache_hit_rate:.1}%"
266
+ )
267
+ .unwrap();
268
+ }
269
150
- if is_idle {
151
- debug!("{msg}");
152
- } else {
153
- info!("{msg}");
270
+ log_stats_line!("{msg}");
271
+
272
+ if let Some(spec_stats) = spec_decoding_log_stats {
273
+ msg.clear();
274
+ write!(
275
+ msg,
276
+ "SpecDecoding metrics: \
277
+ Mean acceptance length: {:.2}, \
278
+ Accepted throughput: {:.2} tokens/s, \
279
+ Drafted throughput: {:.2} tokens/s, \
280
+ Accepted: {} tokens, \
281
+ Drafted: {} tokens",
282
+ spec_stats.mean_acceptance_length,
283
+ spec_stats.accepted_throughput,
284
+ spec_stats.draft_throughput,
285
+ spec_stats.accepted_tokens,
286
+ spec_stats.draft_tokens,
287
+ )
288
+ .unwrap();
289
+ if !spec_stats.per_position_acceptance_rates.is_empty() {
290
+ msg.push_str(", Per-position acceptance rate: ");
291
+ format_position_rates(&mut msg, &spec_stats.per_position_acceptance_rates);
292
+ }
293
+ write!(
294
+ msg,
295
+ ", Avg Draft acceptance rate: {:.1}%",
296
+ spec_stats.draft_acceptance_rate,
297
+ )
298
+ .unwrap();
299
+ log_stats_line!("{msg}");
300
+ }
301
+
302
+ // TODO: Decide on best way to surface CUDAGraph interval samples.
303
+
304
+ if let Some(mfu_stats) = mfu_log_stats {
305
+ log_stats_line!(
306
+ "MFU: {:.1} TF/s/GPU {:.1} GB/s/GPU",
307
+ mfu_stats.tflops_per_gpu,
308
+ mfu_stats.gbps_per_gpu,
309
+ );
310
}
311
312
last_prompt_throughput = prompt_throughput;
@@ -162,17 +318,21 @@ async fn run_stats_logger(model_name: String, engine_count: usize) {
318
319
/// Read the current cumulative counter values for throughput computation.
320
fn read_counters(engines: &[EngineMetrics]) -> CounterSnapshot {
165
- let mut snap = CounterSnapshot {
166
- prompt_tokens: 0,
167
- generation_tokens: 0,
168
- prefix_cache_queries: 0,
169
- prefix_cache_hits: 0,
170
- };
321
+ let mut snap = CounterSnapshot::default();
322
for e in engines {
323
snap.prompt_tokens += e.prompt_tokens_computed.get();
324
snap.generation_tokens += e.generation_tokens.get();
325
snap.prefix_cache_queries += e.prefix_cache_queries.get();
326
snap.prefix_cache_hits += e.prefix_cache_hits.get();
327
+ snap.external_prefix_cache_queries += e.external_prefix_cache_queries.get();
328
+ snap.external_prefix_cache_hits += e.external_prefix_cache_hits.get();
329
+ snap.num_preemptions += e.num_preemptions.get();
330
+ snap.spec_decode_num_drafts += e.spec_decode_num_drafts.get();
331
+ snap.spec_decode_num_draft_tokens += e.spec_decode_num_draft_tokens.get();
332
+ snap.spec_decode_num_accepted_tokens += e.spec_decode_num_accepted_tokens.get();
333
+ snap.estimated_flops_per_gpu += e.estimated_flops_per_gpu.get();
334
+ snap.estimated_read_bytes_per_gpu += e.estimated_read_bytes_per_gpu.get();
335
+ snap.estimated_write_bytes_per_gpu += e.estimated_write_bytes_per_gpu.get();
336
}
337
snap
338
}
@@ -197,3 +357,196 @@ fn read_scheduler_gauges(engines: &[EngineMetrics]) -> (u64, u64, f64) {
357
358
(num_running, num_waiting, kv_cache_usage)
359
}
360
+
361
+/// Read deferred waiting requests across all engines.
362
+fn read_deferred_waiting(engines: &[EngineMetrics]) -> u64 {
363
+ engines.iter().map(|e| e.scheduler_deferred.get()).sum()
364
+}
365
+
366
+/// Return the cache hit rate as a percentage for a counter delta.
367
+fn cache_hit_rate(hits: u64, queries: u64) -> f64 {
368
+ if queries > 0 {
369
+ hits as f64 / queries as f64 * 100.0
370
+ } else {
371
+ 0.0
372
+ }
373
+}
374
+
375
+/// Compute aggregate spec-decoding stats for one logging interval.
376
+fn spec_decoding_log_stats(
377
+ curr: &CounterSnapshot,
378
+ prev: &CounterSnapshot,
379
+ elapsed: f64,
380
+ raw_log_stats: &SchedulerLogStatsInterval,
381
+) -> Option<SpecDecodingLogStats> {
382
+ let num_drafts = curr.spec_decode_num_drafts.wrapping_sub(prev.spec_decode_num_drafts);
383
+ if num_drafts == 0 {
384
+ return None;
385
+ }
386
+
387
+ let draft_tokens = curr
388
+ .spec_decode_num_draft_tokens
389
+ .wrapping_sub(prev.spec_decode_num_draft_tokens);
390
+ let accepted_tokens = curr
391
+ .spec_decode_num_accepted_tokens
392
+ .wrapping_sub(prev.spec_decode_num_accepted_tokens);
393
+
394
+ let (accepted_throughput, draft_throughput) = if elapsed > 0.0 {
395
+ (
396
+ accepted_tokens as f64 / elapsed,
397
+ draft_tokens as f64 / elapsed,
398
+ )
399
+ } else {
400
+ (0.0, 0.0)
401
+ };
402
+ let draft_acceptance_rate = if draft_tokens > 0 {
403
+ accepted_tokens as f64 / draft_tokens as f64 * 100.0
404
+ } else {
405
+ f64::NAN
406
+ };
407
+ let per_position_acceptance_rates = if raw_log_stats.spec_num_drafts > 0 {
408
+ raw_log_stats
409
+ .spec_accepted_tokens_per_pos
410
+ .iter()
411
+ .map(|accepted_tokens| *accepted_tokens as f64 / raw_log_stats.spec_num_drafts as f64)
412
+ .collect()
413
+ } else {
414
+ Vec::new()
415
+ };
416
+
417
+ Some(SpecDecodingLogStats {
418
+ mean_acceptance_length: 1.0 + accepted_tokens as f64 / num_drafts as f64,
419
+ accepted_throughput,
420
+ draft_throughput,
421
+ accepted_tokens,
422
+ draft_tokens,
423
+ per_position_acceptance_rates,
424
+ draft_acceptance_rate,
425
+ })
426
+}
427
+
428
+/// Compute average per-GPU MFU rates for one logging interval.
429
+fn mfu_log_stats(
430
+ curr: &CounterSnapshot,
431
+ prev: &CounterSnapshot,
432
+ elapsed: f64,
433
+ engine_count: usize,
434
+) -> Option<MfuLogStats> {
435
+ let flops = curr.estimated_flops_per_gpu.wrapping_sub(prev.estimated_flops_per_gpu);
436
+ let read_bytes = curr
437
+ .estimated_read_bytes_per_gpu
438
+ .wrapping_sub(prev.estimated_read_bytes_per_gpu);
439
+ let write_bytes = curr
440
+ .estimated_write_bytes_per_gpu
441
+ .wrapping_sub(prev.estimated_write_bytes_per_gpu);
442
+
443
+ if flops == 0 && read_bytes == 0 && write_bytes == 0 {
444
+ return None;
445
+ }
446
+
447
+ let denominator = elapsed * engine_count.max(1) as f64;
448
+ let (tflops_per_gpu, gbps_per_gpu) = if denominator > 0.0 {
449
+ (
450
+ flops as f64 / denominator / 1e12,
451
+ (read_bytes as f64 + write_bytes as f64) / denominator / 1e9,
452
+ )
453
+ } else {
454
+ (0.0, 0.0)
455
+ };
456
+
457
+ Some(MfuLogStats {
458
+ tflops_per_gpu,
459
+ gbps_per_gpu,
460
+ })
461
+}
462
+
463
+/// Drain raw scheduler DTO stats for the configured model and engines.
464
+fn drain_scheduler_log_stats(engines: &[EngineMetrics]) -> SchedulerLogStatsInterval {
465
+ let mut interval = SchedulerLogStatsInterval::default();
466
+ for engine in engines {
467
+ interval.merge(engine.log_stats.drain());
468
+ }
469
+ interval
470
+}
471
+
472
+/// Append spec-decoding per-position acceptance rates like Python's logger.
473
+fn format_position_rates(output: &mut String, rates: &[f64]) {
474
+ for (position, rate) in rates.iter().enumerate() {
475
+ if position > 0 {
476
+ output.push_str(", ");
477
+ }
478
+ write!(output, "{rate:.3}").unwrap();
479
+ }
480
+}
481
+
482
+#[cfg(test)]
483
+mod tests {
484
+ use super::*;
485
+
486
+ #[test]
487
+ fn cache_hit_rate_returns_percent_for_non_empty_queries() {
488
+ assert_eq!(cache_hit_rate(25, 100), 25.0);
489
+ assert_eq!(cache_hit_rate(0, 0), 0.0);
490
+ }
491
+
492
+ #[test]
493
+ fn spec_decoding_log_stats_uses_interval_deltas() {
494
+ let raw_log_stats = SchedulerLogStatsInterval {
495
+ spec_num_drafts: 4,
496
+ spec_accepted_tokens_per_pos: vec![4, 2, 1],
497
+ ..Default::default()
498
+ };
499
+ let prev = CounterSnapshot {
500
+ spec_decode_num_drafts: 10,
501
+ spec_decode_num_draft_tokens: 100,
502
+ spec_decode_num_accepted_tokens: 40,
503
+ ..Default::default()
504
+ };
505
+ let curr = CounterSnapshot {
506
+ spec_decode_num_drafts: 14,
507
+ spec_decode_num_draft_tokens: 120,
508
+ spec_decode_num_accepted_tokens: 52,
509
+ ..Default::default()
510
+ };
511
+
512
+ let stats = spec_decoding_log_stats(&curr, &prev, 2.0, &raw_log_stats).unwrap();
513
+
514
+ assert_eq!(stats.mean_acceptance_length, 4.0);
515
+ assert_eq!(stats.accepted_throughput, 6.0);
516
+ assert_eq!(stats.draft_throughput, 10.0);
517
+ assert_eq!(stats.accepted_tokens, 12);
518
+ assert_eq!(stats.draft_tokens, 20);
519
+ assert_eq!(stats.per_position_acceptance_rates, vec![1.0, 0.5, 0.25]);
520
+ assert_eq!(stats.draft_acceptance_rate, 60.0);
521
+ }
522
+
523
+ #[test]
524
+ fn mfu_log_stats_averages_per_gpu_across_engines() {
525
+ let prev = CounterSnapshot {
526
+ estimated_flops_per_gpu: 10,
527
+ estimated_read_bytes_per_gpu: 10,
528
+ estimated_write_bytes_per_gpu: 10,
529
+ ..Default::default()
530
+ };
531
+ let curr = CounterSnapshot {
532
+ estimated_flops_per_gpu: 4_000_000_000_010,
533
+ estimated_read_bytes_per_gpu: 2_000_000_010,
534
+ estimated_write_bytes_per_gpu: 2_000_000_010,
535
+ ..Default::default()
536
+ };
537
+
538
+ let stats = mfu_log_stats(&curr, &prev, 2.0, 2).unwrap();
539
+
540
+ assert_eq!(stats.tflops_per_gpu, 1.0);
541
+ assert_eq!(stats.gbps_per_gpu, 1.0);
542
+ }
543
+
544
+ #[test]
545
+ fn format_position_rates_uses_three_decimal_places() {
546
+ let mut output = String::new();
547
+
548
+ format_position_rates(&mut output, &[1.0, 0.5, 0.25]);
549
+
550
+ assert_eq!(output, "1.000, 0.500, 0.250");
551
+ }
552
+}
rust/src/metrics/src/scheduler.rs
+119
-2
@@ -1,4 +1,5 @@
1
-use std::collections::BTreeSet;
1
+use std::collections::{BTreeMap, BTreeSet};
2
+use std::sync::{Arc, Mutex};
3
4
use itertools::Itertools as _;
5
use prometheus_client::encoding::{EncodeLabelSet, EncodeLabelValue, LabelValueEncoder};
@@ -62,6 +63,90 @@ pub struct LoraInfoLabels {
63
pub waiting_lora_adapters: LoraAdapterNames,
64
}
65
66
+/// CUDA graph sample key used for periodic text-log aggregation.
67
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
68
+pub struct CudagraphLogKey {
69
+ pub num_unpadded_tokens: u64,
70
+ pub num_padded_tokens: u64,
71
+ pub num_paddings: u64,
72
+ pub runtime_mode: String,
73
+}
74
+
75
+/// Raw scheduler stats accumulated for one periodic text-log interval.
76
+#[derive(Default)]
77
+pub struct SchedulerLogStatsInterval {
78
+ pub spec_num_drafts: u64,
79
+ pub spec_accepted_tokens_per_pos: Vec<u64>,
80
+ pub cudagraph_counts: BTreeMap<CudagraphLogKey, u64>,
81
+}
82
+
83
+impl SchedulerLogStatsInterval {
84
+ /// Merge another drained interval into this one.
85
+ pub fn merge(&mut self, other: Self) {
86
+ self.spec_num_drafts += other.spec_num_drafts;
87
+
88
+ if self.spec_accepted_tokens_per_pos.len() < other.spec_accepted_tokens_per_pos.len() {
89
+ self.spec_accepted_tokens_per_pos
90
+ .resize(other.spec_accepted_tokens_per_pos.len(), 0);
91
+ }
92
+ for (position, accepted_tokens) in
93
+ other.spec_accepted_tokens_per_pos.into_iter().enumerate()
94
+ {
95
+ self.spec_accepted_tokens_per_pos[position] += accepted_tokens;
96
+ }
97
+
98
+ for (key, count) in other.cudagraph_counts {
99
+ *self.cudagraph_counts.entry(key).or_default() += count;
100
+ }
101
+ }
102
+}
103
+
104
+/// Internal, non-Prometheus accumulator for periodic text logs that need raw
105
+/// scheduler DTOs.
106
+#[derive(Clone, Default)]
107
+pub struct SchedulerLogStatsAccumulator {
108
+ inner: Arc<Mutex<SchedulerLogStatsInterval>>,
109
+}
110
+
111
+impl SchedulerLogStatsAccumulator {
112
+ /// Observe spec-decoding fields needed for per-position text-log rates.
113
+ pub fn observe_spec_decode(&self, num_drafts: u64, accepted_tokens_per_pos: &[u64]) {
114
+ let mut inner = self.inner.lock().expect("scheduler log stats accumulator poisoned");
115
+ inner.spec_num_drafts += num_drafts;
116
+
117
+ if inner.spec_accepted_tokens_per_pos.len() < accepted_tokens_per_pos.len() {
118
+ inner.spec_accepted_tokens_per_pos.resize(accepted_tokens_per_pos.len(), 0);
119
+ }
120
+ for (position, accepted_tokens) in accepted_tokens_per_pos.iter().copied().enumerate() {
121
+ inner.spec_accepted_tokens_per_pos[position] += accepted_tokens;
122
+ }
123
+ }
124
+
125
+ /// Observe one CUDA graph runtime sample for the interval table.
126
+ pub fn observe_cudagraph(
127
+ &self,
128
+ num_unpadded_tokens: u64,
129
+ num_padded_tokens: u64,
130
+ num_paddings: u64,
131
+ runtime_mode: &str,
132
+ ) {
133
+ let mut inner = self.inner.lock().expect("scheduler log stats accumulator poisoned");
134
+ let key = CudagraphLogKey {
135
+ num_unpadded_tokens,
136
+ num_padded_tokens,
137
+ num_paddings,
138
+ runtime_mode: runtime_mode.to_string(),
139
+ };
140
+ *inner.cudagraph_counts.entry(key).or_default() += 1;
141
+ }
142
+
143
+ /// Drain and reset the current text-log interval.
144
+ pub fn drain(&self) -> SchedulerLogStatsInterval {
145
+ let mut inner = self.inner.lock().expect("scheduler log stats accumulator poisoned");
146
+ std::mem::take(&mut *inner)
147
+ }
148
+}
149
+
150
/// Scheduler/batch-scoped Prometheus families exported from `SchedulerStats`.
151
pub struct SchedulerMetrics {
152
// Scheduler state gauges.
@@ -95,6 +180,9 @@ pub struct SchedulerMetrics {
180
pub kv_block_lifetime_seconds: HistogramFamily,
181
pub kv_block_idle_before_evict_seconds: HistogramFamily,
182
pub kv_block_reuse_gap_seconds: HistogramFamily,
183
+
184
+ /// Non-Prometheus interval accumulators for periodic text-log helpers.
185
+ pub log_stats: Family<EngineLabels, SchedulerLogStatsAccumulator>,
186
}
187
188
impl SchedulerMetrics {
@@ -265,13 +353,14 @@ impl SchedulerMetrics {
353
kv_block_lifetime_seconds,
354
kv_block_idle_before_evict_seconds,
355
kv_block_reuse_gap_seconds,
356
+ log_stats: Family::default(),
357
}
358
}
359
}
360
361
#[cfg(test)]
362
mod tests {
274
- use crate::{EngineLabels, Metrics};
363
+ use crate::{CudagraphLogKey, EngineLabels, Metrics, SchedulerLogStatsAccumulator};
364
365
#[test]
366
fn perf_counters_render_with_a_single_total_suffix() {
@@ -301,4 +390,32 @@ mod tests {
390
assert!(!rendered.contains("vllm:estimated_read_bytes_per_gpu_total_total"));
391
assert!(!rendered.contains("vllm:estimated_write_bytes_per_gpu_total_total"));
392
}
393
+
394
+ #[test]
395
+ fn log_stats_accumulator_drains_interval_data() {
396
+ let accumulator = SchedulerLogStatsAccumulator::default();
397
+
398
+ accumulator.observe_spec_decode(2, &[1, 2]);
399
+ accumulator.observe_spec_decode(3, &[3, 4, 5]);
400
+ accumulator.observe_cudagraph(8, 16, 8, "FULL");
401
+ accumulator.observe_cudagraph(8, 16, 8, "FULL");
402
+
403
+ let interval = accumulator.drain();
404
+
405
+ assert_eq!(interval.spec_num_drafts, 5);
406
+ assert_eq!(interval.spec_accepted_tokens_per_pos, vec![4, 6, 5]);
407
+ assert_eq!(
408
+ interval
409
+ .cudagraph_counts
410
+ .get(&CudagraphLogKey {
411
+ num_unpadded_tokens: 8,
412
+ num_padded_tokens: 16,
413
+ num_paddings: 8,
414
+ runtime_mode: "FULL".to_string(),
415
+ })
416
+ .copied(),
417
+ Some(2)
418
+ );
419
+ assert_eq!(accumulator.drain().spec_num_drafts, 0);
420
+ }
421
}