From a71174f257b3b37fa2a93f133baab120409a7f29 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Fri, 11 Sep 2026 12:52:39 +0000 Subject: [PATCH] feat(frontend): expose cumulative prefix cache metrics through vllm Signed-off-by: JinYan Su <751080330@qq.com> --- .../subsystems/frontend/prometheus-metrics.md | 26 ++++++- pegainfer-frontend/src/engine/metrics.rs | 15 +++- pegainfer-frontend/src/vllm/bridge.rs | 35 ++++++++- pegainfer-frontend/src/vllm/bridge/stepped.rs | 78 +++++++++++++++---- pegainfer-frontend/src/vllm/bridge/tests.rs | 70 +++++++++++++++++ pegainfer-gemma4/src/engine.rs | 1 + pegainfer-glm52/src/scheduler/load.rs | 1 + pegainfer-qwen3/src/frontend_adapter.rs | 1 + pegainfer-qwen35/src/scheduler/mod.rs | 2 + pegainfer-qwen35/src/scheduler/tests.rs | 1 + 10 files changed, 205 insertions(+), 25 deletions(-) diff --git a/docs/subsystems/frontend/prometheus-metrics.md b/docs/subsystems/frontend/prometheus-metrics.md index e17df01c8..b996f0e3c 100644 --- a/docs/subsystems/frontend/prometheus-metrics.md +++ b/docs/subsystems/frontend/prometheus-metrics.md @@ -1,8 +1,8 @@ # Prometheus /metrics via the vLLM frontend -**TL;DR:** `/metrics` exposes request histograms for every model and engine gauges for schedulers that publish `LoadSnapshot`: Qwen3 and Qwen3.5 use one logical engine, while GLM5.2 EP8/DP8 uses eight rank-local engines and GLM5.2 TP8 uses one logical engine. The bridge forwards each partition's stats under the same identity the vLLM frontend uses for least-load routing. +**TL;DR:** `/metrics` exposes cumulative prefix-cache counters for schedulers that populate `SchedulerMetrics::prefix_cache`, plus request histograms for every model and engine gauges for schedulers that publish `LoadSnapshot`: Qwen3 and Qwen3.5 use one logical engine, while GLM5.2 EP8/DP8 uses eight rank-local engines and GLM5.2 TP8 uses one logical engine. The bridge forwards each partition's stats under the same identity the vLLM frontend uses for least-load routing. -Last touched: 2026-07 +Last touched: 2026-09 ## How the numbers flow @@ -18,9 +18,27 @@ Measured cost is noise in both covered configurations: - Qwen3 TPOT: 10.6387 ms (main) vs 10.6395 ms (metrics branch) over 828 tokens. - GLM5.2 EP8, three-run median at concurrency 64: 1268.58 vs 1264.82 output tok/s (-0.30%); TPOT p50 41.76 vs 41.35 ms. +## Prefix-cache counters + +Schedulers can publish lifetime `PrefixCacheCounters { requests, queries, hits }` +in `SchedulerMetrics::prefix_cache`. Count each successfully admitted request +once: queries are all prompt tokens, hits are the tokens actually reused, +including prefixes restored from local host offload. Waiting/retries do not +count again; periodic logging must not reset these totals. + +Both bridge paths convert cumulative snapshots into per-send deltas for +`vllm:prefix_cache_queries_total` and `vllm:prefix_cache_hits_total`. Coalesced +snapshots retain increments; repeated snapshots add zero. The stepped bridge +also sends a stats-only batch when a cache lookup has no output tokens. +Schedulers that leave these counters at their default still report zero. + +The token hit ratio is the rate of hits divided by the rate of queries, +aggregated over the same model/engine labels. With no queries the ratio is +undefined. External connector cache metrics are separate and remain unwired. + ## What deliberately reads zero (state at capture time) -- `prefix_cache_queries/hits` and the by-reason waiting split (`reason="deferred"` is driven by a skipped-request counter we don't report; all waiting shows as `reason="capacity"`). +- The by-reason waiting split (`reason="deferred"` is driven by a skipped-request counter we don't report; all waiting shows as `reason="capacity"`). - Spec-decode counters, per-GPU FLOPs/bytes estimates, KV-block residency histograms, cudagraph stats — the bridge sends `SchedulerStats::default()` for these fields. - Every model crate whose scheduler doesn't publish a `LoadSnapshot` watch (currently deepseek and kimi) gets path 1 only; its engine gauges are absent, not lying-zero — the bridge skips the stats task for that partition when no watch exists. @@ -28,4 +46,4 @@ Measured cost is noise in both covered configurations: Qwen3.5 single-GPU live RTX 5090 validation confirmed that running and KV gauges rise during generation, waiting rises under batch-slot pressure, and all three return to zero after drain and recovery. The commands and metric samples are recorded in [Qwen3.5 Scheduler LoadSnapshot](../../models/qwen35/load-snapshot.md#validation-boundary). TP uses the same scheduler publication path but was not part of that live run. -Next, wire the DeepSeek-V2-Lite and Kimi-K2 schedulers using the same recipe, and report real prefix-cache query/hit counters instead of zeros. A future partitioned model must expose its logical scheduler partitions instead of averaging them behind engine 0. +Next, wire the DeepSeek-V2-Lite and Kimi-K2 schedulers using the same recipe, and populate prefix-cache counters in model schedulers that still leave them at zero. A future partitioned model must expose its logical scheduler partitions instead of averaging them behind engine 0. diff --git a/pegainfer-frontend/src/engine/metrics.rs b/pegainfer-frontend/src/engine/metrics.rs index 8fead9060..72b0f6a16 100644 --- a/pegainfer-frontend/src/engine/metrics.rs +++ b/pegainfer-frontend/src/engine/metrics.rs @@ -1,7 +1,6 @@ //! What the scheduler republishes about itself: [`SchedulerMetrics`], the //! per-iteration snapshot of occupancy gauges plus whatever richer counters a -//! model line serves (today: cumulative speculative-decode acceptance, when a -//! draft model is loaded). +//! model line serves (prefix-cache reuse and speculative-decode acceptance). use std::error::Error; use std::fmt; @@ -26,6 +25,18 @@ pub struct SchedulerMetrics { pub num_waiting_reqs: u64, /// Cumulative spec-decode counters, or `None` when no draft model is loaded. pub spec_decode: Option, + /// Cumulative prefix-cache lookups; count each admitted request once. + pub prefix_cache: PrefixCacheCounters, +} + +/// Lifetime prefix-cache totals, independent of periodic logging windows. +/// Queries count all prompt tokens and hits count tokens actually reused, +/// including locally offloaded prefixes. Admission retries must not count again. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct PrefixCacheCounters { + pub requests: u64, + pub queries: u64, + pub hits: u64, } /// Upper bound on a drafter's `K`, fixing the width of diff --git a/pegainfer-frontend/src/vllm/bridge.rs b/pegainfer-frontend/src/vllm/bridge.rs index 0094702c8..b1d54c64f 100644 --- a/pegainfer-frontend/src/vllm/bridge.rs +++ b/pegainfer-frontend/src/vllm/bridge.rs @@ -35,7 +35,9 @@ use vllm_engine_core_client::protocol::output::StopReason; use vllm_engine_core_client::protocol::output::UtilityCallOutput; use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::request::EngineCoreRequestType; +use vllm_engine_core_client::protocol::stats::BaseCacheStats; use vllm_engine_core_client::protocol::stats::PrefillStats; +use vllm_engine_core_client::protocol::stats::PrefixCacheStats; use vllm_engine_core_client::protocol::stats::SchedulerStats; use vllm_engine_core_client::protocol::stats::SpecDecodingStats; use vllm_engine_core_client::protocol::utility::UtilityCallId; @@ -53,6 +55,7 @@ use zeromq::util::PeerIdentity; use crate::engine::EngineHandle; use crate::engine::FinishReason; use crate::engine::GenerateRequest; +use crate::engine::PrefixCacheCounters; use crate::engine::RequestAbortReason; use crate::engine::RequestTag; use crate::engine::SchedulerMetrics; @@ -650,6 +653,33 @@ impl SpecDecodeTracker { } } +/// One baseline per bridge: coalesced snapshots retain all increments, and +/// repeated snapshots contribute zero to the frontend's Prometheus counters. +#[derive(Default)] +pub(crate) struct SchedulerStatsTracker { + spec: SpecDecodeTracker, + prefix: PrefixCacheCounters, +} + +impl SchedulerStatsTracker { + pub(crate) fn interval(&mut self, snapshot: &SchedulerMetrics) -> SchedulerStats { + let cur = snapshot.prefix_cache; + let mut stats = scheduler_stats_from(snapshot); + stats.spec_decoding_stats = self.spec.interval(snapshot); + stats.prefix_cache_stats = PrefixCacheStats { + base: BaseCacheStats { + requests: cur.requests.saturating_sub(self.prefix.requests), + queries: cur.queries.saturating_sub(self.prefix.queries), + hits: cur.hits.saturating_sub(self.prefix.hits), + ..Default::default() + }, + ..Default::default() + }; + self.prefix = cur; + stats + } +} + /// Forward every scheduler load snapshot as a stats-only output batch; the /// frontend records it into the shared Prometheus registry. Sends the current /// snapshot up front so the gauges initialize before the first step, then one @@ -662,11 +692,10 @@ async fn publish_scheduler_stats( output_tx: mpsc::UnboundedSender, shutdown: CancellationToken, ) -> Result<()> { - let mut spec = SpecDecodeTracker::default(); + let mut tracker = SchedulerStatsTracker::default(); loop { let snapshot = *load_rx.borrow_and_update(); - let mut stats = scheduler_stats_from(&snapshot); - stats.spec_decoding_stats = spec.interval(&snapshot); + let stats = tracker.interval(&snapshot); let outputs = RequestBatchOutputs { engine_index, scheduler_stats: Some(Box::new(stats)), diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 48179e9cc..0924ad543 100644 --- a/pegainfer-frontend/src/vllm/bridge/stepped.rs +++ b/pegainfer-frontend/src/vllm/bridge/stepped.rs @@ -38,11 +38,10 @@ use zeromq::ZmqMessage; use zeromq::prelude::SocketRecv; use super::BridgeLink; -use super::SpecDecodeTracker; +use super::SchedulerStatsTracker; use super::connect_link; use super::engine_output; use super::now_secs_f64; -use super::scheduler_stats_from; use super::send_outputs; use super::send_terminal_output; use super::send_utility_response; @@ -80,7 +79,7 @@ impl SteppedEngineBridge { .scheduler .take_steps() .context("partition step stream already taken")?; - let mut spec = SpecDecodeTracker::default(); + let mut tracker = SchedulerStatsTracker::default(); // Stats are pull-at-send: no push task, the load cell is read when a // batch goes out (and once here, so the frontend's gauges initialize // before any traffic). An idle engine publishes nothing. @@ -103,7 +102,7 @@ impl SteppedEngineBridge { &output_tx, RequestBatchOutputs { engine_index: self.engine_index, - scheduler_stats: Some(Box::new(self.stats(&mut spec))), + scheduler_stats: Some(Box::new(self.stats(&mut tracker))), timestamp: now_secs_f64(), ..Default::default() } @@ -146,7 +145,7 @@ impl SteppedEngineBridge { &anchor, &mut streams, &mut names, - &mut spec, + &mut tracker, &output_tx, ) { break Err(error).context("failed to dispatch local engine step"); @@ -182,13 +181,10 @@ impl SteppedEngineBridge { run_result } - /// Stats for an outgoing batch; the spec delta runs from the last batch + /// Stats for an outgoing batch; counter deltas run from the last batch /// stamped, not the last step run. - fn stats(&self, spec: &mut SpecDecodeTracker) -> SchedulerStats { - let snapshot = self.scheduler.metrics(); - let mut stats = scheduler_stats_from(&snapshot); - stats.spec_decoding_stats = spec.interval(&snapshot); - stats + fn stats(&self, tracker: &mut SchedulerStatsTracker) -> SchedulerStats { + tracker.interval(&self.scheduler.metrics()) } fn dispatch_step( @@ -197,7 +193,7 @@ impl SteppedEngineBridge { anchor: &UnixAnchor, streams: &mut HashMap, names: &mut HashMap, - spec: &mut SpecDecodeTracker, + tracker: &mut SchedulerStatsTracker, output_tx: &tokio::sync::mpsc::UnboundedSender< vllm_engine_core_client::protocol::output::EngineCoreOutputs, >, @@ -225,10 +221,10 @@ impl SteppedEngineBridge { } if outputs.is_empty() { - // A drafted step with no batch to ride would strand its increment + // A cache lookup or draft with no batch to ride would strand its increment // until the next batch, which may never come. - let stats = self.stats(spec); - if stats.spec_decoding_stats.is_some() { + let stats = self.stats(tracker); + if stats.spec_decoding_stats.is_some() || stats.prefix_cache_stats.base.requests > 0 { send_outputs( output_tx, RequestBatchOutputs { @@ -252,7 +248,7 @@ impl SteppedEngineBridge { engine_index: self.engine_index, outputs, finished_requests: (!finished_requests.is_empty()).then_some(finished_requests), - scheduler_stats: Some(Box::new(self.stats(spec))), + scheduler_stats: Some(Box::new(self.stats(tracker))), timestamp: now_secs_f64(), } .into(), @@ -698,6 +694,56 @@ mod tests { } } + #[test] + fn prefix_only_step_is_sent_once_without_output_tokens() { + use vllm_engine_core_client::protocol::output::EngineCoreOutputs; + + use crate::engine::PrefixCacheCounters; + use crate::engine::SchedulerMetrics; + + let (handle, backend) = scheduler_pair(); + let bridge = bridge(handle); + backend.metrics.publish(&SchedulerMetrics { + prefix_cache: PrefixCacheCounters { + requests: 1, + queries: 100, + hits: 60, + }, + ..Default::default() + }); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let mut tracker = SchedulerStatsTracker::default(); + let mut streams = HashMap::new(); + let mut names = HashMap::new(); + let anchor = UnixAnchor::now(); + for _ in 0..2 { + bridge + .dispatch_step( + StepOutputs::default(), + &anchor, + &mut streams, + &mut names, + &mut tracker, + &tx, + ) + .unwrap(); + } + let EngineCoreOutputs::RequestBatch(batch) = rx.try_recv().unwrap() else { + panic!("expected stats batch"); + }; + let stats = batch + .scheduler_stats + .unwrap() + .prefix_cache_stats + .base + .clone(); + assert_eq!((stats.requests, stats.queries, stats.hits), (1, 100, 60)); + assert!( + rx.try_recv().is_err(), + "unchanged totals must not send another batch" + ); + } + fn wire_request(completion: Option, prompt: Option) -> EngineCoreRequest { let mut params = EngineCoreSamplingParams::for_test(); params.logprobs = completion; diff --git a/pegainfer-frontend/src/vllm/bridge/tests.rs b/pegainfer-frontend/src/vllm/bridge/tests.rs index 05ec7346b..51bf00330 100644 --- a/pegainfer-frontend/src/vllm/bridge/tests.rs +++ b/pegainfer-frontend/src/vllm/bridge/tests.rs @@ -571,6 +571,7 @@ async fn load_snapshots_become_stats_only_batches() { kv_total_blocks: 100, num_running_reqs: 2, num_waiting_reqs: 1, + prefix_cache: crate::engine::PrefixCacheCounters::default(), spec_decode: None, }); let (output_tx, mut output_rx) = mpsc::unbounded_channel(); @@ -681,3 +682,72 @@ async fn spec_stats_are_per_interval_deltas_that_skip_idle_intervals() { .expect("stats task exits on shutdown") .expect("stats publisher shuts down cleanly"); } + +/// Every possible coalescing of four admissions must preserve the same totals. +/// Reading a snapshot twice (or changing only gauges) must never count twice. +#[test] +fn prefix_stats_preserve_totals_across_all_coalescings() { + let admissions = [(10, 0), (20, 10), (30, 20), (40, 0)]; + for mask in 0..8 { + let mut tracker = SchedulerStatsTracker::default(); + let mut snapshot = SchedulerMetrics::default(); + let mut observed = (0, 0, 0); + for (i, (queries, hits)) in admissions.into_iter().enumerate() { + snapshot.prefix_cache.requests += 1; + snapshot.prefix_cache.queries += queries; + snapshot.prefix_cache.hits += hits; + if i == 3 || mask & (1 << i) != 0 { + let stats = tracker.interval(&snapshot).prefix_cache_stats.base; + observed.0 += stats.requests; + observed.1 += stats.queries; + observed.2 += stats.hits; + snapshot.num_running_reqs += 1; + let repeated = tracker.interval(&snapshot).prefix_cache_stats.base; + assert_eq!( + (repeated.requests, repeated.queries, repeated.hits), + (0, 0, 0) + ); + } + } + assert_eq!(observed, (4, 100, 30), "coalescing mask {mask}"); + } +} + +#[tokio::test] +async fn prefix_stats_are_forwarded_by_the_watch_publisher() { + let snapshot = SchedulerMetrics { + prefix_cache: PrefixCacheCounters { + requests: 2, + queries: 100, + hits: 40, + }, + ..Default::default() + }; + let (load_tx, load_rx) = tokio::sync::watch::channel(snapshot); + let (output_tx, mut output_rx) = mpsc::unbounded_channel(); + let shutdown = CancellationToken::new(); + let task = tokio::spawn(publish_scheduler_stats( + 0, + load_rx, + output_tx, + shutdown.clone(), + )); + let first = next_scheduler_stats(&mut output_rx) + .await + .prefix_cache_stats + .base + .clone(); + assert_eq!((first.requests, first.queries, first.hits), (2, 100, 40)); + load_tx.send_replace(snapshot); + let repeated = next_scheduler_stats(&mut output_rx) + .await + .prefix_cache_stats + .base + .clone(); + assert_eq!( + (repeated.requests, repeated.queries, repeated.hits), + (0, 0, 0) + ); + shutdown.cancel(); + task.await.unwrap().unwrap(); +} diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 92191c532..a762250fa 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -2316,6 +2316,7 @@ impl Scheduler for Gemma4Scheduler { kv_total_blocks: (local_total + global_total) as u64, num_running_reqs: (self.active.len() + walkers + lane_inflight) as u64, num_waiting_reqs: self.pending.len() as u64, + prefix_cache: pegainfer_frontend::engine::PrefixCacheCounters::default(), spec_decode: None, } } diff --git a/pegainfer-glm52/src/scheduler/load.rs b/pegainfer-glm52/src/scheduler/load.rs index 83f026ccb..7e31294d7 100644 --- a/pegainfer-glm52/src/scheduler/load.rs +++ b/pegainfer-glm52/src/scheduler/load.rs @@ -32,6 +32,7 @@ pub(super) fn publish_load( // the frontend's placement signal never undercounts a rank // mid-resolve and never counts a request twice. num_waiting_reqs: (pending.len() + resolving) as u64, + prefix_cache: pegainfer_frontend::engine::PrefixCacheCounters::default(), spec_decode: None, }); } diff --git a/pegainfer-qwen3/src/frontend_adapter.rs b/pegainfer-qwen3/src/frontend_adapter.rs index e02125ea5..febbf514f 100644 --- a/pegainfer-qwen3/src/frontend_adapter.rs +++ b/pegainfer-qwen3/src/frontend_adapter.rs @@ -801,6 +801,7 @@ impl Scheduler for Qwen3Scheduler { num_waiting_reqs: (self.deferred.len() + self.loading.len() + self.post_control_deferred.len()) as u64, + prefix_cache: pegainfer_frontend::engine::PrefixCacheCounters::default(), spec_decode: self.executor.spec_decode_counters(), } } diff --git a/pegainfer-qwen35/src/scheduler/mod.rs b/pegainfer-qwen35/src/scheduler/mod.rs index ff726d1b4..d5121fd22 100644 --- a/pegainfer-qwen35/src/scheduler/mod.rs +++ b/pegainfer-qwen35/src/scheduler/mod.rs @@ -542,6 +542,7 @@ fn publish_load( kv_total_blocks, num_running_reqs, num_waiting_reqs, + prefix_cache: pegainfer_frontend::engine::PrefixCacheCounters::default(), spec_decode: None, }); } @@ -602,6 +603,7 @@ fn terminal_scheduler_shutdown( kv_total_blocks, num_running_reqs: 0, num_waiting_reqs: 0, + prefix_cache: pegainfer_frontend::engine::PrefixCacheCounters::default(), spec_decode: None, }); } diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 9bc6f1e3f..7c58351e7 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -719,6 +719,7 @@ fn terminal_shutdown_closes_drains_and_errors_every_owner_once() { kv_total_blocks: 64, num_running_reqs: 9, num_waiting_reqs: 9, + prefix_cache: pegainfer_frontend::engine::PrefixCacheCounters::default(), spec_decode: None, }); terminal_scheduler_shutdown(