From ab0234f87f1dfe7deaa9b648f7c06449a06d48a0 Mon Sep 17 00:00:00 2001 From: yuntaonie <1571859588@qq.com> Date: Sat, 22 Aug 2026 18:31:26 +0800 Subject: [PATCH 1/2] feat(observability): report real prefix-cache query/hit counters in /metrics Thread prefix-cache query/hit counters from the qwen3 scheduler through SchedulerMetrics into the vLLM SchedulerStats.prefix_cache_stats surface, so Prometheus /metrics no longer reads zeros for prefix cache hit rate. - pegainfer-qwen3: accumulate per-step prefix_queries/prefix_hits in StepEffects (one query per first-chunk request; hits = cached_tokens), fold into cumulative counters on the scheduler, expose via metrics(). - pegainfer-frontend: add prefix_cache_queries/hits to SchedulerMetrics and map them to PrefixCacheStats in the vLLM bridge. This complements the cached_tokens usage path (TokenEvent::Scheduled) that upstream already landed for #603; it covers the /metrics consumer only. Verified: pegainfer-frontend --lib tests pass (65); cargo clippy and cargo check --workspace --lib clean. qwen3 A100 e2e blocked locally by a rdma-mummy-sys bindgen environment issue unrelated to this change. --- pegainfer-frontend/src/engine/metrics.rs | 6 ++++++ pegainfer-frontend/src/vllm/bridge.rs | 10 ++++++++++ pegainfer-frontend/src/vllm/bridge/tests.rs | 2 ++ pegainfer-qwen3/src/frontend_adapter.rs | 15 +++++++++++++++ pegainfer-qwen3/src/scheduler/effects.rs | 10 ++++++++++ pegainfer-qwen3/src/scheduler/resolve.rs | 8 ++++++++ 6 files changed, 51 insertions(+) diff --git a/pegainfer-frontend/src/engine/metrics.rs b/pegainfer-frontend/src/engine/metrics.rs index 8fead9060..773b041bd 100644 --- a/pegainfer-frontend/src/engine/metrics.rs +++ b/pegainfer-frontend/src/engine/metrics.rs @@ -26,6 +26,12 @@ 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 queries (one per request that reached its first + /// prefill chunk). Monotonic; mapped to vLLM `PrefixCacheStats.queries`. + pub prefix_cache_queries: u64, + /// Cumulative prefix-cache hits, in tokens. Monotonic; mapped to vLLM + /// `PrefixCacheStats.hits`. + pub prefix_cache_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 0744e75d2..90cc5e0c3 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; @@ -597,6 +599,14 @@ pub(crate) fn scheduler_stats_from(snapshot: &SchedulerMetrics) -> SchedulerStat } else { snapshot.kv_used_blocks as f64 / snapshot.kv_total_blocks as f64 }, + prefix_cache_stats: PrefixCacheStats { + base: BaseCacheStats { + queries: snapshot.prefix_cache_queries, + hits: snapshot.prefix_cache_hits, + ..BaseCacheStats::default() + }, + ..PrefixCacheStats::default() + }, ..SchedulerStats::default() } } diff --git a/pegainfer-frontend/src/vllm/bridge/tests.rs b/pegainfer-frontend/src/vllm/bridge/tests.rs index 0926609f7..1de286d1c 100644 --- a/pegainfer-frontend/src/vllm/bridge/tests.rs +++ b/pegainfer-frontend/src/vllm/bridge/tests.rs @@ -570,6 +570,8 @@ async fn load_snapshots_become_stats_only_batches() { num_running_reqs: 2, num_waiting_reqs: 1, spec_decode: None, + prefix_cache_queries: 0, + prefix_cache_hits: 0, }); let (output_tx, mut output_rx) = mpsc::unbounded_channel(); let shutdown = CancellationToken::new(); diff --git a/pegainfer-qwen3/src/frontend_adapter.rs b/pegainfer-qwen3/src/frontend_adapter.rs index 3954e3644..e5d2b1450 100644 --- a/pegainfer-qwen3/src/frontend_adapter.rs +++ b/pegainfer-qwen3/src/frontend_adapter.rs @@ -241,6 +241,12 @@ pub(crate) struct Qwen3Scheduler { /// it cannot run against an adapter set the command was about to change. pending_control: VecDeque, post_control_deferred: Vec, + /// Cumulative prefix-cache queries (one per admitted request that reached + /// its first prefill chunk). Monotonic; reported verbatim in `SchedulerMetrics`. + prefix_cache_queries: u64, + /// Cumulative prefix-cache hits, in tokens (sum of cached prefix lengths). + /// Monotonic; reported verbatim in `SchedulerMetrics`. + prefix_cache_hits: u64, } impl Qwen3Scheduler { @@ -265,6 +271,8 @@ impl Qwen3Scheduler { lora_rx, pending_control: VecDeque::new(), post_control_deferred: Vec::new(), + prefix_cache_queries: 0, + prefix_cache_hits: 0, } } @@ -342,6 +350,11 @@ impl Qwen3Scheduler { // terminal rides the committed step, which the driver ships after // publishing metrics — the finishing batch's send-time stats then // read the drained occupancy instead of racing the publish. + // Fold this step's prefix-cache counters into the cumulative totals + // before any effect is dropped, so retries/re-queues never lose a count. + self.prefix_cache_queries += effects.prefix_queries; + self.prefix_cache_hits += effects.prefix_hits; + let mut finishes: Vec<(RequestId, FinishReason)> = Vec::new(); for cached in effects.cached { @@ -801,6 +814,8 @@ impl Scheduler for Qwen3Scheduler { + self.loading.len() + self.post_control_deferred.len()) as u64, spec_decode: self.executor.spec_decode_counters(), + prefix_cache_queries: self.prefix_cache_queries, + prefix_cache_hits: self.prefix_cache_hits, } } } diff --git a/pegainfer-qwen3/src/scheduler/effects.rs b/pegainfer-qwen3/src/scheduler/effects.rs index 8e37a70e3..9b5d1d97d 100644 --- a/pegainfer-qwen3/src/scheduler/effects.rs +++ b/pegainfer-qwen3/src/scheduler/effects.rs @@ -87,6 +87,14 @@ pub(crate) struct StepEffects { pub(crate) prompt_echoes: Vec, pub(crate) pending: Vec, pub(crate) decode: Vec, + /// Prefix-cache queries counted this step (one per request whose first + /// prefill chunk was resolved — that is where the cache is consulted). + /// Cumulative counters live on the scheduler; this is the per-step delta. + pub(crate) prefix_queries: u64, + /// Prefix-cache hits counted this step, in tokens (the cached prefix + /// length of each first-chunk request). Cumulative counters live on the + /// scheduler; this is the per-step delta. + pub(crate) prefix_hits: u64, } impl StepEffects { @@ -96,6 +104,8 @@ impl StepEffects { prompt_echoes: Vec::new(), pending: Vec::new(), decode: Vec::new(), + prefix_queries: 0, + prefix_hits: 0, } } } diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index 12356210f..1e0c7bb75 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -27,12 +27,16 @@ pub(crate) fn resolve_step( prompt_echoes: Vec::new(), pending: Vec::new(), decode: resolve_decode_outputs(executor, active, &result.requests), + prefix_queries: 0, + prefix_hits: 0, }, ExecutionArtifacts::SpeculativeDecode { verify } => StepEffects { cached: Vec::new(), prompt_echoes: Vec::new(), pending: Vec::new(), decode: resolve_speculative_outputs(executor, active, &verify.requests), + prefix_queries: 0, + prefix_hits: 0, }, ExecutionArtifacts::Unified { pending, result } => { let mut effects = resolve_prefill_outputs(executor, pending, result.prefill_requests); @@ -107,6 +111,10 @@ fn resolve_prefill_outputs( request_id: req.request_id, cached_tokens: result.cached_tokens, }); + // One cache query per request (the cache is consulted once, on the + // first chunk); the cached prefix length is the hit count, in tokens. + effects.prefix_queries += 1; + effects.prefix_hits += result.cached_tokens as u64; } if !result.completed { From b988a322fc274e318b4f00e579c91816c9bb0214 Mon Sep 17 00:00:00 2001 From: yuntaonie <1571859588@qq.com> Date: Mon, 24 Aug 2026 22:12:45 +0800 Subject: [PATCH 2/2] fix(observability): report prefix-cache counters as token-granular per-send deltas MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two review bugs on the prefix-cache `/metrics` surface, both fixed: 1. Unit mismatch (queries vs hits). Previously `prefix_queries += 1` (a request) while `prefix_hits += cached_tokens` (a token count), so `hit_rate = hits/queries` could exceed 100%. Now both are TOKEN-granular, matching vLLM's `PrefixCacheStats`: `prefix_queries` counts the prompt tokens looked up in the cache and `prefix_hits` counts the cached tokens. Because cached <= prompt, `hits <= queries` and the rate stays in [0, 1]. Guarded by the existing `prefill_pos == 0` check so each request is counted exactly once (no double counting across chunked prefill). 2. Cumulative overcount in Prometheus. The scheduler holds running totals, but `dispatch_step` / `publish_scheduler_stats` shipped that running total on *every* token batch, and the frontend adds each `SchedulerStats` value into its `prefix_cache_*_total` counters — so a cached request re-added the whole history on every subsequent batch until restart. Now the bridge ships per-send DELTAS (cur - last), mirroring the existing spec-decode path: `prefix_cache_delta()` in bridge.rs, with `last_prefix_*` state in both the legacy `publish_scheduler_stats` loop and the stepped bridge (AtomicU64, since `dispatch_step` takes `&self`). Adds a `FakeExecutor` prefix-hit hook, a multi-batch/multi-scrape qwen3 scheduler test (`prefix_cache_metrics_stable_across_batches_and_scrapes`, token-granular assertions), and a frontend test (`prefix_cache_stats_are_per_interval_deltas_not_running_totals`) that proves the bridge ships the interval delta, not the running total. Signed-off-by: yuntaonie <1571859588@qq.com> --- pegainfer-frontend/src/engine/metrics.rs | 7 +- pegainfer-frontend/src/vllm/bridge.rs | 43 +++++-- pegainfer-frontend/src/vllm/bridge/stepped.rs | 52 +++++++- pegainfer-frontend/src/vllm/bridge/tests.rs | 56 +++++++++ pegainfer-frontend/src/vllm/mod.rs | 2 + pegainfer-qwen3/src/frontend_adapter.rs | 8 +- pegainfer-qwen3/src/frontend_adapter/tests.rs | 116 ++++++++++++++++++ pegainfer-qwen3/src/scheduler/effects.rs | 15 ++- pegainfer-qwen3/src/scheduler/resolve.rs | 17 ++- pegainfer-qwen3/src/scheduler/test_support.rs | 20 ++- 10 files changed, 309 insertions(+), 27 deletions(-) diff --git a/pegainfer-frontend/src/engine/metrics.rs b/pegainfer-frontend/src/engine/metrics.rs index 773b041bd..1656522e0 100644 --- a/pegainfer-frontend/src/engine/metrics.rs +++ b/pegainfer-frontend/src/engine/metrics.rs @@ -29,8 +29,11 @@ pub struct SchedulerMetrics { /// Cumulative prefix-cache queries (one per request that reached its first /// prefill chunk). Monotonic; mapped to vLLM `PrefixCacheStats.queries`. pub prefix_cache_queries: u64, - /// Cumulative prefix-cache hits, in tokens. Monotonic; mapped to vLLM - /// `PrefixCacheStats.hits`. + /// Cumulative prefix-cache hits, token-granularity (the total number of + /// queried prompt tokens already cached). Same unit as + /// `prefix_cache_queries`, so `hit_rate = hits/queries` stays in [0, 1]. + /// Monotonic; mapped to vLLM `PrefixCacheStats.hits` as a per-send delta + /// (see `scheduler_stats_from`), never as the running total. pub prefix_cache_hits: u64, } diff --git a/pegainfer-frontend/src/vllm/bridge.rs b/pegainfer-frontend/src/vllm/bridge.rs index 90cc5e0c3..b4597512b 100644 --- a/pegainfer-frontend/src/vllm/bridge.rs +++ b/pegainfer-frontend/src/vllm/bridge.rs @@ -37,7 +37,6 @@ 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; @@ -590,6 +589,14 @@ fn stop_sentinel_id(eos_token_id: Option, stop_token_ids: &[u32]) -> Option /// vLLM `SchedulerStats` view of a load snapshot — what the frontend's /// Prometheus gauges (`scheduler_running`, `scheduler_waiting`, /// `kv_cache_usage`) and DP load balancer consume. +/// +/// `prefix_cache_stats` is left at zero here on purpose: the running totals in +/// `SchedulerMetrics` must NOT be shipped as-is, because the frontend's +/// Prometheus logger increments its `prefix_cache_*_total` counters by the +/// value of *every* `SchedulerStats` it receives. The bridge therefore exports +/// per-send **deltas** (see [`prefix_cache_delta`]) so a cached request does +/// not re-add the whole history on each subsequent token batch. Callers fill +/// `prefix_cache_stats` from `prefix_cache_delta` after this call. pub(crate) fn scheduler_stats_from(snapshot: &SchedulerMetrics) -> SchedulerStats { SchedulerStats { num_running_reqs: snapshot.num_running_reqs, @@ -599,18 +606,28 @@ pub(crate) fn scheduler_stats_from(snapshot: &SchedulerMetrics) -> SchedulerStat } else { snapshot.kv_used_blocks as f64 / snapshot.kv_total_blocks as f64 }, - prefix_cache_stats: PrefixCacheStats { - base: BaseCacheStats { - queries: snapshot.prefix_cache_queries, - hits: snapshot.prefix_cache_hits, - ..BaseCacheStats::default() - }, - ..PrefixCacheStats::default() - }, ..SchedulerStats::default() } } +/// Per-interval prefix-cache delta from a cumulative snapshot and the last +/// snapshot we shipped, in the wire shape the frontend increments its +/// `prefix_cache_queries_total` / `prefix_cache_hits_total` counters by. The +/// scheduler holds running totals; the wire carries deltas so Prometheus does +/// not re-add the running total on every token batch (see the P2 review on the +/// prefix-cache counters). `last` is updated by the caller to `cur` after use. +pub(crate) fn prefix_cache_delta( + last_queries: u64, + last_hits: u64, + cur: &SchedulerMetrics, +) -> BaseCacheStats { + BaseCacheStats { + queries: cur.prefix_cache_queries.saturating_sub(last_queries), + hits: cur.prefix_cache_hits.saturating_sub(last_hits), + ..BaseCacheStats::default() + } +} + /// Per-interval spec-decode delta from two cumulative snapshots, in the wire /// shape the frontend increments its `vllm:spec_decode_*_total` counters by (see /// [`SpecDecodeCounters`] for why the transport carries totals and the wire @@ -647,6 +664,11 @@ async fn publish_scheduler_stats( shutdown: CancellationToken, ) -> Result<()> { let mut last_spec = SpecDecodeCounters::default(); + // Last prefix-cache totals we shipped, so each interval carries only the + // delta (the frontend adds every SchedulerStats value into its *_total + // counters — shipping the running total would re-count history each batch). + let mut last_prefix_q = 0u64; + let mut last_prefix_h = 0u64; loop { let snapshot = *load_rx.borrow_and_update(); let spec_decoding_stats = if let Some(cur) = &snapshot.spec_decode { @@ -662,6 +684,9 @@ async fn publish_scheduler_stats( None }; let mut stats = scheduler_stats_from(&snapshot); + stats.prefix_cache_stats.base = prefix_cache_delta(last_prefix_q, last_prefix_h, &snapshot); + last_prefix_q = snapshot.prefix_cache_queries; + last_prefix_h = snapshot.prefix_cache_hits; stats.spec_decoding_stats = spec_decoding_stats; let outputs = RequestBatchOutputs { engine_index, diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 663484106..a73acc491 100644 --- a/pegainfer-frontend/src/vllm/bridge/stepped.rs +++ b/pegainfer-frontend/src/vllm/bridge/stepped.rs @@ -40,6 +40,7 @@ use super::BridgeLink; use super::connect_link; use super::engine_output; use super::now_secs_f64; +use super::prefix_cache_delta; use super::scheduler_stats_from; use super::send_outputs; use super::send_terminal_output; @@ -68,6 +69,13 @@ pub(crate) struct SteppedEngineBridge { pub(crate) max_model_len: u32, pub(crate) engine_index: u32, pub(crate) data_parallel_size: u32, + /// Last prefix-cache totals we shipped, so each token batch carries only + /// the delta. The frontend adds every `SchedulerStats` value into its + /// `prefix_cache_*_total` counters, so shipping the running total would + /// re-count history on every subsequent batch (see the P2 review on the + /// prefix-cache counters). `AtomicU64` because `dispatch_step` takes `&self`. + pub(crate) last_prefix_queries: std::sync::atomic::AtomicU64, + pub(crate) last_prefix_hits: std::sync::atomic::AtomicU64, } impl SteppedEngineBridge { @@ -94,11 +102,31 @@ impl SteppedEngineBridge { &shutdown, ) .await?; + // Seed the gauges before any traffic. Ship the delta since last send + // (zero here, the first interval) so the frontend's *_total counters + // accumulate deltas, never the running total. + let snapshot = self.scheduler.metrics(); + let mut seed_stats = scheduler_stats_from(&snapshot); + let last_q = self + .last_prefix_queries + .load(std::sync::atomic::Ordering::Relaxed); + let last_h = self + .last_prefix_hits + .load(std::sync::atomic::Ordering::Relaxed); + seed_stats.prefix_cache_stats.base = prefix_cache_delta(last_q, last_h, &snapshot); + self.last_prefix_queries.store( + snapshot.prefix_cache_queries, + std::sync::atomic::Ordering::Relaxed, + ); + self.last_prefix_hits.store( + snapshot.prefix_cache_hits, + std::sync::atomic::Ordering::Relaxed, + ); send_outputs( &output_tx, RequestBatchOutputs { engine_index: self.engine_index, - scheduler_stats: Some(Box::new(scheduler_stats_from(&self.scheduler.metrics()))), + scheduler_stats: Some(Box::new(seed_stats)), timestamp: now_secs_f64(), ..Default::default() } @@ -215,13 +243,33 @@ impl SteppedEngineBridge { // load before committing the step), so the batch carries stats that // match its own tokens — a finishing batch reports the drained state // and the gauges settle instead of freezing at the last busy value. + // Ship only the delta since the last send so the frontend's + // `prefix_cache_*_total` counters accumulate increments, not the whole + // running total on every batch. + let snapshot = self.scheduler.metrics(); + let mut batch_stats = scheduler_stats_from(&snapshot); + let last_q = self + .last_prefix_queries + .load(std::sync::atomic::Ordering::Relaxed); + let last_h = self + .last_prefix_hits + .load(std::sync::atomic::Ordering::Relaxed); + batch_stats.prefix_cache_stats.base = prefix_cache_delta(last_q, last_h, &snapshot); + self.last_prefix_queries.store( + snapshot.prefix_cache_queries, + std::sync::atomic::Ordering::Relaxed, + ); + self.last_prefix_hits.store( + snapshot.prefix_cache_hits, + std::sync::atomic::Ordering::Relaxed, + ); send_outputs( output_tx, RequestBatchOutputs { engine_index: self.engine_index, outputs, finished_requests: (!finished_requests.is_empty()).then_some(finished_requests), - scheduler_stats: Some(Box::new(scheduler_stats_from(&self.scheduler.metrics()))), + scheduler_stats: Some(Box::new(batch_stats)), timestamp: now_secs_f64(), } .into(), diff --git a/pegainfer-frontend/src/vllm/bridge/tests.rs b/pegainfer-frontend/src/vllm/bridge/tests.rs index 1de286d1c..e39747a30 100644 --- a/pegainfer-frontend/src/vllm/bridge/tests.rs +++ b/pegainfer-frontend/src/vllm/bridge/tests.rs @@ -681,3 +681,59 @@ async fn spec_stats_are_per_interval_deltas_that_skip_idle_intervals() { .expect("stats task exits on shutdown") .expect("stats publisher shuts down cleanly"); } + +/// The scheduler holds running prefix-cache totals; the bridge must ship +/// per-send DELTAS (like spec-decode), never the running total. Otherwise the +/// frontend's `prefix_cache_*_total` counters re-add the whole history on every +/// token batch and overcount until restart. Regression test for the P2 review +/// on the prefix-cache counters. +#[tokio::test] +async fn prefix_cache_stats_are_per_interval_deltas_not_running_totals() { + let (load_tx, load_rx) = tokio::sync::watch::channel(SchedulerMetrics { + prefix_cache_queries: 100, + prefix_cache_hits: 37, + ..SchedulerMetrics::default() + }); + 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(), + )); + + // First publish diffs against zero: carries the whole cumulative. + let first = next_scheduler_stats(&mut output_rx).await; + let pc = first.prefix_cache_stats.base; + assert_eq!(pc.queries, 100, "first interval diffs against zero"); + assert_eq!(pc.hits, 37); + + // The running total doubles; the sent value must be the DELTA (100, 37), + // not the new running total (200, 74). This is the overcount bug. + load_tx.send_replace(SchedulerMetrics { + prefix_cache_queries: 200, + prefix_cache_hits: 74, + ..SchedulerMetrics::default() + }); + let second = next_scheduler_stats(&mut output_rx).await; + let pc = second.prefix_cache_stats.base; + assert_eq!(pc.queries, 100, "second interval ships the delta, not 200"); + assert_eq!(pc.hits, 37, "second interval ships the delta, not 74"); + + // No further change: an idle interval ships a zero delta (not a re-add). + load_tx.send_replace(SchedulerMetrics { + prefix_cache_queries: 200, + prefix_cache_hits: 74, + ..SchedulerMetrics::default() + }); + let idle = next_scheduler_stats(&mut output_rx).await; + let pc = idle.prefix_cache_stats.base; + assert_eq!(pc.queries, 0, "unchanged interval ships a zero delta"); + assert_eq!(pc.hits, 0); + + shutdown.cancel(); + task.await + .expect("stats task exits on shutdown") + .expect("stats publisher shuts down cleanly"); +} diff --git a/pegainfer-frontend/src/vllm/mod.rs b/pegainfer-frontend/src/vllm/mod.rs index a2f696c97..b889d61b9 100644 --- a/pegainfer-frontend/src/vllm/mod.rs +++ b/pegainfer-frontend/src/vllm/mod.rs @@ -294,6 +294,8 @@ where max_model_len, engine_index: engine_index as u32, data_parallel_size, + last_prefix_queries: std::sync::atomic::AtomicU64::new(0), + last_prefix_hits: std::sync::atomic::AtomicU64::new(0), }; let shutdown = bridge_shutdown.clone(); bridges.spawn(async move { (engine_index, bridge.run(shutdown).await) }); diff --git a/pegainfer-qwen3/src/frontend_adapter.rs b/pegainfer-qwen3/src/frontend_adapter.rs index e5d2b1450..640b5f72d 100644 --- a/pegainfer-qwen3/src/frontend_adapter.rs +++ b/pegainfer-qwen3/src/frontend_adapter.rs @@ -244,8 +244,12 @@ pub(crate) struct Qwen3Scheduler { /// Cumulative prefix-cache queries (one per admitted request that reached /// its first prefill chunk). Monotonic; reported verbatim in `SchedulerMetrics`. prefix_cache_queries: u64, - /// Cumulative prefix-cache hits, in tokens (sum of cached prefix lengths). - /// Monotonic; reported verbatim in `SchedulerMetrics`. + /// Cumulative prefix-cache hits, token-granularity (the total number of + /// queried prompt tokens that were already cached, summed across requests). + /// Same unit as `prefix_cache_queries`, so `hit_rate = hits/queries` ∈ [0, 1]. + /// Monotonic; the bridge exports per-send deltas of this total (see + /// `scheduler_stats_from` / the stepped bridge) so Prometheus does not + /// re-add the running total on every token batch. prefix_cache_hits: u64, } diff --git a/pegainfer-qwen3/src/frontend_adapter/tests.rs b/pegainfer-qwen3/src/frontend_adapter/tests.rs index 35441bc7d..f1e442324 100644 --- a/pegainfer-qwen3/src/frontend_adapter/tests.rs +++ b/pegainfer-qwen3/src/frontend_adapter/tests.rs @@ -371,3 +371,119 @@ fn lora_control_waits_until_scheduler_idle() { .expect_err("adapter load should be a stub error"); assert!(matches!(error, LoraControlError::Failed(_))); } + +/// E2E (no GPU) for the prefix-cache `/metrics` counters: drive a real +/// `Qwen3Scheduler` through the engine contract with a fake executor that +/// reports a cached prefix on every request's first chunk, then scrape +/// `SchedulerMetrics` across multiple batches and scrapes. +/// +/// Guards against the two bugs from review: +/// * unit mismatch — `prefix_queries`/`prefix_hits` are both TOKEN-granular +/// (queried prompt tokens / cached tokens), matching vLLM's +/// `PrefixCacheStats`, so `hit_rate = hits/queries` stays in [0, 1]; +/// * cumulative-counter double counting — every request is counted exactly +/// once (on its first prefill chunk), and the bridge exports per-send +/// deltas of the running total (see `prefix_cache_delta`), so repeated +/// scrapes at the same instant are identical and the frontend does not +/// re-add history on every token batch. +#[test] +fn prefix_cache_metrics_stable_across_batches_and_scrapes() { + const BATCHES: u64 = 4; + const PER_BATCH: u64 = 3; + const TOTAL: u64 = BATCHES * PER_BATCH; + const PROMPT_TOKENS: u64 = 64; // tokens looked up per request + const HIT_TOKENS: u64 = 37; // simulated cached prefix length (<= prompt) + + // A fake KV cache that reports a 37-token hit on the first chunk of every + // request. With the fix, queries count the 64 prompt tokens queried and + // hits count the 37 tokens already cached — both token-granular. + let dropped = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(64, Arc::clone(&dropped)).with_prefix_hit(HIT_TOKENS as usize); + let (partition, _lora, mut steps) = launch(executor, false); + + let mut controls = Vec::new(); + let mut scrapes: Vec<(u64, u64, u64)> = Vec::new(); // (batch, queries, hits) + + for batch in 0..BATCHES { + for _ in 0..PER_BATCH { + controls.push(partition.handle.submit(request(PROMPT_TOKENS as usize, 4))); + } + // Wait until this batch's prefills have been counted. + let target = (batch + 1) * PER_BATCH * PROMPT_TOKENS; + assert!( + wait_until(Duration::from_secs(2), || partition + .handle + .metrics() + .prefix_cache_queries + >= target), + "batch {batch} prefix queries never reached {target}" + ); + let m = partition.handle.metrics(); + eprintln!( + "[scrape] after batch {batch}: prefix_cache_queries={} prefix_cache_hits={}", + m.prefix_cache_queries, m.prefix_cache_hits + ); + scrapes.push((batch, m.prefix_cache_queries, m.prefix_cache_hits)); + } + + // Drain the step stream so the driver thread can exit cleanly. + for c in &controls { + let _ = steps.collect_terminal(c.id()); + } + + // Final stable snapshot: repeated scrapes at the same instant are identical. + let a = partition.handle.metrics(); + let b = partition.handle.metrics(); + let c = partition.handle.metrics(); + assert_eq!(a.prefix_cache_queries, b.prefix_cache_queries); + assert_eq!(a.prefix_cache_hits, b.prefix_cache_hits); + assert_eq!(b.prefix_cache_queries, c.prefix_cache_queries); + assert_eq!(b.prefix_cache_hits, c.prefix_cache_hits); + eprintln!( + "[scrape] final (x3 identical): prefix_cache_queries={} prefix_cache_hits={}", + a.prefix_cache_queries, a.prefix_cache_hits + ); + + // Token-granular correctness per vLLM PrefixCacheStats: every request + // queries PROMPT_TOKENS and hits HIT_TOKENS, so the running totals are + // TOTAL * those, and hit_rate = HIT_TOKENS / PROMPT_TOKENS in [0, 1]. + let expected_q = TOTAL * PROMPT_TOKENS; + let expected_h = TOTAL * HIT_TOKENS; + assert_eq!(a.prefix_cache_queries, expected_q); + assert_eq!(a.prefix_cache_hits, expected_h); + assert!( + a.prefix_cache_hits <= a.prefix_cache_queries, + "hits (cached tokens) must not exceed queries (queried tokens)" + ); + let hit_rate = a.prefix_cache_hits as f64 / a.prefix_cache_queries as f64; + eprintln!( + "[rate] prefix hit_rate={:.3} (== {}/{})", + hit_rate, HIT_TOKENS, PROMPT_TOKENS + ); + assert!((hit_rate - HIT_TOKENS as f64 / PROMPT_TOKENS as f64).abs() < 1e-9); + + // Each batch contributed a stable, non-zero delta of exactly + // PER_BATCH * PROMPT_TOKENS (queries) and PER_BATCH * HIT_TOKENS (hits). + let mut prev_q = 0u64; + let mut prev_h = 0u64; + for (batch, q, h) in scrapes { + let dq = q - prev_q; + let dh = h - prev_h; + eprintln!( + "[delta] batch {batch}: +queries={dq} +hits={dh} (hit_rate={:.3})", + h as f64 / q.max(1) as f64 + ); + assert_eq!( + dq, + PER_BATCH * PROMPT_TOKENS, + "batch {batch} added exactly PER_BATCH*PROMPT_TOKENS queries" + ); + assert_eq!( + dh, + PER_BATCH * HIT_TOKENS, + "batch {batch} added exactly PER_BATCH*HIT_TOKENS hits" + ); + prev_q = q; + prev_h = h; + } +} diff --git a/pegainfer-qwen3/src/scheduler/effects.rs b/pegainfer-qwen3/src/scheduler/effects.rs index 9b5d1d97d..e2f89a506 100644 --- a/pegainfer-qwen3/src/scheduler/effects.rs +++ b/pegainfer-qwen3/src/scheduler/effects.rs @@ -87,13 +87,16 @@ pub(crate) struct StepEffects { pub(crate) prompt_echoes: Vec, pub(crate) pending: Vec, pub(crate) decode: Vec, - /// Prefix-cache queries counted this step (one per request whose first - /// prefill chunk was resolved — that is where the cache is consulted). - /// Cumulative counters live on the scheduler; this is the per-step delta. - pub(crate) prefix_queries: u64, - /// Prefix-cache hits counted this step, in tokens (the cached prefix - /// length of each first-chunk request). Cumulative counters live on the + /// Prefix-cache queries counted this step, token-granularity: the number of + /// prompt tokens looked up in the cache (summed over requests whose first + /// prefill chunk was resolved this step — that is where the cache is + /// consulted). Same unit as `prefix_hits`. Cumulative counters live on the /// scheduler; this is the per-step delta. + pub(crate) prefix_queries: u64, + /// Prefix-cache hits counted this step, token-granularity: the number of + /// queried prompt tokens that were already cached (`cached_tokens`). Same + /// unit as `prefix_queries`, so `hit_rate = hits/queries` stays in [0, 1]. + /// Cumulative counters live on the scheduler; this is the per-step delta. pub(crate) prefix_hits: u64, } diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index 1e0c7bb75..8cea807cd 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -104,16 +104,23 @@ fn resolve_prefill_outputs( // release builds too. assert_eq!(req.request_id, result.request_id); - // Report the prefix-cache hit count on the request's first chunk only - // — that is where it is determined. Later chunks must not re-report. + // Report the prefix-cache counters on the request's first chunk only — + // that is where they are determined. Later chunks must not re-report. + // + // Both counters are TOKEN-granularity, matching vLLM's `PrefixCacheStats` + // (the frontend compares `hits / queries` as hit tokens / queried tokens): + // * `prefix_queries` = the number of prompt tokens this request looked + // up in the cache (the whole prompt is consulted once, on chunk 0). + // * `prefix_hits` = the number of those tokens already cached + // (`cached_tokens`). + // Because `cached_tokens <= prompt_tokens`, `hits <= queries` holds and + // the hit rate stays in [0, 1] with no impossible >100% rates. if req.prefill_pos == 0 { effects.cached.push(CachedTokensEffect { request_id: req.request_id, cached_tokens: result.cached_tokens, }); - // One cache query per request (the cache is consulted once, on the - // first chunk); the cached prefix length is the hit count, in tokens. - effects.prefix_queries += 1; + effects.prefix_queries += req.prompt_tokens.len() as u64; effects.prefix_hits += result.cached_tokens as u64; } diff --git a/pegainfer-qwen3/src/scheduler/test_support.rs b/pegainfer-qwen3/src/scheduler/test_support.rs index 16d19323a..98f6b833d 100644 --- a/pegainfer-qwen3/src/scheduler/test_support.rs +++ b/pegainfer-qwen3/src/scheduler/test_support.rs @@ -39,6 +39,10 @@ pub(crate) struct FakeExecutor { pub(crate) dropped: Arc>>, pub(crate) prefetch_offers: Arc>>, stop_token: Option, + // When > 0, the first prefill chunk of every request reports this many + // `cached_tokens` (a simulated prefix-cache hit). Drives the prefix-cache + // query/hit counters without a real GPU KV cache. + prefix_hit_tokens: usize, } impl FakeExecutor { @@ -56,9 +60,17 @@ impl FakeExecutor { dropped, prefetch_offers: Arc::new(Mutex::new(Vec::new())), stop_token: None, + prefix_hit_tokens: 0, } } + /// Simulate a prefix-cache hit on every request's first prefill chunk by + /// reporting `tokens` cached tokens. + pub(crate) fn with_prefix_hit(mut self, tokens: usize) -> Self { + self.prefix_hit_tokens = tokens; + self + } + pub(crate) fn with_stop_token(mut self, token: u32) -> Self { self.stop_token = Some(token); self @@ -101,7 +113,13 @@ impl FakeExecutor { first_token: 100 + req.request_id.raw() as u32, first_token_logprob: None, prompt_logprobs: None, - cached_tokens: 0, + // A simulated prefix-cache hit is reported only on the request's + // first chunk (start == 0); later chunks carry no cached prefix. + cached_tokens: if start == 0 { + self.prefix_hit_tokens + } else { + 0 + }, completed, prefill_pos, }