From 00a8fd22bea784f68a42c78a57544666eacca8e3 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 20 Oct 2026 07:37:53 +0000 Subject: [PATCH] [Rust Frontend] Cap request token histogram buckets at max_model_len The request token-count histograms (prompt, generation, max_num_generation, params_max_tokens, prefill_kv_computed) used a 1-2-5 series fixed at 131072 plus an extra bound at 131072. Python builds `request_tokens` buckets with `build_1_2_5_buckets(max_model_len)` and no extra bound. Set the cap from the engine-reported max_model_len when `Llm` is created, through a shared histogram constructor on the lazily built registry. Co-authored-by: Claude Opus 5.5 Signed-off-by: Bugen Zhao --- rust/src/llm/src/lib.rs | 2 + rust/src/llm/src/request_metrics.rs | 2 + rust/src/metrics/src/request.rs | 116 ++++++++++++++++++++++------ 3 files changed, 95 insertions(+), 25 deletions(-) diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index 1e15f11c12e3..44dbef90f0ae 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -5,6 +5,7 @@ use std::num::NonZeroU32; use tracing::Span; use vllm_engine_core_client::EngineCoreClient; +use vllm_metrics::METRICS; mod error; mod inflight; @@ -45,6 +46,7 @@ impl Llm { /// Create a new minimal LLM facade from an already connected engine-core /// client. pub fn new(client: EngineCoreClient) -> Self { + METRICS.request.set_max_model_len(client.max_model_len()); Self { client, randomize_request_id: true, diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index eebf0596b0d2..271c5b411750 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -406,11 +406,13 @@ pub fn current_unix_timestamp_secs() -> f64 { mod tests { use vllm_engine_core_client::protocol::output::{EngineCoreEvent, EngineCoreEventType}; use vllm_engine_core_client::protocol::stats::PrefillStats; + use vllm_metrics::METRICS; use super::{RequestMetricsTracker, RequestTimestamps, diff_or_zero}; #[test] fn tracker_updates_timing_state_across_prefill_decode_and_finish() { + METRICS.request.set_max_model_len(4096); let mut tracker = RequestMetricsTracker::new("model".to_string(), 2, 100.0, 64, Some(128), 1, true); diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index 46260cf58397..c24d6a6b29d1 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::sync::{Arc, OnceLock}; + use prometheus_client::encoding::EncodeLabelSet; -use prometheus_client::metrics::family::Family; +use prometheus_client::metrics::family::{Family, MetricConstructor}; use prometheus_client::registry::Registry; use crate::{EngineLabels, Histogram, HistogramFamily, U64Counter}; @@ -25,20 +27,19 @@ const ITERATION_TOKENS_BUCKETS: [f64; 13] = [ 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, ]; +/// Build a 1-2-5 series of buckets capped at `max_value`, like Python's +/// `build_1_2_5_buckets`: `100` gives `[1, 2, 5, 10, 20, 50, 100]`. fn build_1_2_5_buckets(max_value: u32) -> Vec { + let max_value = u64::from(max_value); let mut buckets = Vec::new(); let mut exponent = 0; loop { - for mantissa in [1_u32, 2, 5] { - let value = mantissa * 10_u32.pow(exponent); - if value <= max_value { - buckets.push(value as f64); - } else { - if buckets.last().copied() != Some(max_value as f64) { - buckets.push(max_value as f64); - } + for mantissa in [1_u64, 2, 5] { + let value = mantissa * 10_u64.pow(exponent); + if value > max_value { return buckets; } + buckets.push(value as f64); } exponent += 1; } @@ -60,11 +61,6 @@ fn request_latency_histogram() -> Histogram { Histogram::new(REQUEST_LATENCY_BUCKETS.iter().copied()) } -fn request_token_count_histogram() -> Histogram { - // TODO: determine max value based on `max_model_len`. - Histogram::new(build_1_2_5_buckets(131_072)) -} - fn request_params_n_histogram() -> Histogram { Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied()) } @@ -93,6 +89,28 @@ pub struct PromptTokenSourceLabels { pub(crate) type FinishedReasonCounterFamily = Family; pub(crate) type PromptTokenSourceCounterFamily = Family; +pub(crate) type RequestTokensHistogramFamily = + Family; + +/// Constructor of the request token-count histograms, which share Python's +/// `request_tokens` buckets: a 1-2-5 series capped at `max_model_len`. +/// +/// The registry is built before any engine is connected, so the buckets are +/// filled in later by [`RequestMetrics::set_max_model_len`]. +#[derive(Clone, Debug, Default)] +pub struct RequestTokensHistogram { + buckets: Arc>>, +} + +impl MetricConstructor for RequestTokensHistogram { + fn new_metric(&self) -> Histogram { + let buckets = self + .buckets + .get() + .expect("max_model_len must be set before recording request token metrics"); + Histogram::new(buckets.iter().copied()) + } +} /// Request and output-batch Prometheus families. pub struct RequestMetrics { @@ -110,13 +128,13 @@ pub struct RequestMetrics { // Request lifecycle counters and histograms. pub request_success: FinishedReasonCounterFamily, - pub request_prompt_tokens: HistogramFamily, - pub request_generation_tokens: HistogramFamily, - pub request_max_num_generation_tokens: HistogramFamily, - pub request_params_max_tokens: HistogramFamily, + pub request_prompt_tokens: RequestTokensHistogramFamily, + pub request_generation_tokens: RequestTokensHistogramFamily, + pub request_max_num_generation_tokens: RequestTokensHistogramFamily, + pub request_params_max_tokens: RequestTokensHistogramFamily, pub request_params_n: HistogramFamily, pub request_num_preemptions: HistogramFamily, - pub request_prefill_kv_computed_tokens: HistogramFamily, + pub request_prefill_kv_computed_tokens: RequestTokensHistogramFamily, pub time_to_first_token_seconds: HistogramFamily, pub inter_token_latency_seconds: HistogramFamily, pub e2e_request_latency_seconds: HistogramFamily, @@ -125,11 +143,16 @@ pub struct RequestMetrics { pub request_decode_time_seconds: HistogramFamily, pub request_inference_time_seconds: HistogramFamily, pub request_time_per_output_token_seconds: HistogramFamily, + + /// Constructor shared by the request token-count histogram families. + request_tokens_histogram: RequestTokensHistogram, } impl RequestMetrics { /// Register the request-oriented metric families into the shared registry. pub(crate) fn register(registry: &mut Registry) -> Self { + let request_tokens_histogram = RequestTokensHistogram::default(); + let iteration_tokens_total = Family::new_with_constructor(iteration_tokens_histogram as fn() -> Histogram); registry.register( @@ -189,8 +212,7 @@ impl RequestMetrics { request_success.clone(), ); - let request_prompt_tokens = - Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); + let request_prompt_tokens = Family::new_with_constructor(request_tokens_histogram.clone()); registry.register( "vllm:request_prompt_tokens", "Number of prefill tokens processed.", @@ -198,7 +220,7 @@ impl RequestMetrics { ); let request_generation_tokens = - Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); + Family::new_with_constructor(request_tokens_histogram.clone()); registry.register( "vllm:request_generation_tokens", "Number of generation tokens processed.", @@ -206,7 +228,7 @@ impl RequestMetrics { ); let request_max_num_generation_tokens = - Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); + Family::new_with_constructor(request_tokens_histogram.clone()); registry.register( "vllm:request_max_num_generation_tokens", "Histogram of maximum number of requested generation tokens.", @@ -214,7 +236,7 @@ impl RequestMetrics { ); let request_params_max_tokens = - Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); + Family::new_with_constructor(request_tokens_histogram.clone()); registry.register( "vllm:request_params_max_tokens", "Histogram of the max_tokens request parameter.", @@ -238,7 +260,7 @@ impl RequestMetrics { ); let request_prefill_kv_computed_tokens = - Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); + Family::new_with_constructor(request_tokens_histogram.clone()); registry.register( "vllm:request_prefill_kv_computed_tokens", "Histogram of new KV tokens computed during prefill (excluding cached tokens).", @@ -334,6 +356,50 @@ impl RequestMetrics { request_decode_time_seconds, request_inference_time_seconds, request_time_per_output_token_seconds, + request_tokens_histogram, } } + + /// Cap the request token-count histogram buckets at the engine-reported + /// `max_model_len`. + /// + /// Must be called before the first request is recorded. The registry is + /// process-global and serves one model, so only the first call takes + /// effect. + pub fn set_max_model_len(&self, max_model_len: u32) { + let _ = self.request_tokens_histogram.buckets.set(build_1_2_5_buckets(max_model_len)); + } +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use itertools::Itertools; + use prometheus_client::encoding::text::encode; + + use super::*; + + /// Request token-count histograms use Python's 1-2-5 buckets capped at + /// `max_model_len`, with no extra bound at `max_model_len` itself. + #[test] + fn request_token_buckets_are_capped_at_max_model_len() { + let mut registry = Registry::default(); + let metrics = RequestMetrics::register(&mut registry); + metrics.set_max_model_len(4096); + metrics.request_prompt_tokens.get_or_create_owned(&EngineLabels { + model_name: "model".to_string(), + engine: 0, + }); + + let mut output = String::new(); + encode(&mut output, ®istry).unwrap(); + let bounds = output + .lines() + .filter(|line| line.starts_with("vllm:request_prompt_tokens_bucket")) + .filter_map(|line| line.split_once("le=\"")?.1.split_once('"')) + .map(|(bound, _)| bound) + .join(" "); + expect!["1.0 2.0 5.0 10.0 20.0 50.0 100.0 200.0 500.0 1000.0 2000.0 +Inf"] + .assert_eq(&bounds); + } }