[Perf][Rust Frontend] Make histogram observations lock-free (#58574)
Co-authored-by: jthomson04 <jwillthomson19@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Sep 25, 2026 at 10:25 UTC
055f150df539401e8d218da2aec24c9551f42a11
8 files changed
+181
-7
rust/Cargo.lock
+1
@@ -5903,6 +5903,7 @@ dependencies = [
5903
name = "vllm-metrics"
5904
version = "0.1.0"
5905
dependencies = [
5906
+ "expect-test",
5907
"itertools 0.14.0",
5908
"prometheus-client",
5909
]
rust/clippy.toml
new
+3
@@ -0,0 +1,3 @@
1
+disallowed-types = [
2
+ { path = "prometheus_client::metrics::histogram::Histogram", reason = "it locks on every observation; use `vllm_metrics::Histogram`" },
3
+]
rust/src/metrics/Cargo.toml
+3
@@ -8,5 +8,8 @@ license.workspace = true
8
itertools.workspace = true
9
prometheus-client.workspace = true
10
11
+[dev-dependencies]
12
+expect-test.workspace = true
13
+
14
[lints]
15
workspace = true
rust/src/metrics/src/api_server.rs
+1
-2
@@ -3,10 +3,9 @@
3
4
use prometheus_client::encoding::EncodeLabelSet;
5
use prometheus_client::metrics::family::Family;
6
-use prometheus_client::metrics::histogram::Histogram;
6
use prometheus_client::registry::Registry;
7
9
-use crate::U64Counter;
8
+use crate::{Histogram, U64Counter};
9
10
const HTTP_REQUEST_DURATION_BUCKETS: [f64; 3] = [0.1, 0.5, 1.0];
11
const HTTP_REQUEST_DURATION_HIGHR_BUCKETS: [f64; 21] = [
rust/src/metrics/src/histogram.rs
new
+169
@@ -0,0 +1,169 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+use std::fmt;
5
+use std::iter::once;
6
+use std::sync::Arc;
7
+use std::sync::atomic::{AtomicU64, Ordering};
8
+
9
+use prometheus_client::encoding::{EncodeMetric, MetricEncoder, NoLabelSet};
10
+use prometheus_client::metrics::{MetricType, TypedMetric};
11
+
12
+/// Prometheus histogram with lock-free observations.
13
+///
14
+/// `prometheus_client::metrics::histogram::Histogram` takes an exclusive lock for every
15
+/// observation, which serializes all Tokio workers on hot histograms such as inter-token
16
+/// latency. This histogram updates atomic cells instead, so observers never block each other
17
+/// or a concurrent scrape.
18
+///
19
+/// A scrape may see an observation's bucket count before its sum. `_count` is derived from the
20
+/// bucket counts, so it always equals the `+Inf` bucket.
21
+#[derive(Clone, Debug)]
22
+pub struct Histogram {
23
+ inner: Arc<Inner>,
24
+}
25
+
26
+#[derive(Debug)]
27
+struct Inner {
28
+ /// Upper bounds, ending with `f64::MAX`, which is encoded as `+Inf`.
29
+ upper_bounds: Box<[f64]>,
30
+ /// Non-cumulative observation count of each bucket.
31
+ bucket_counts: Box<[AtomicU64]>,
32
+ /// Bits of the `f64` sum of all observations.
33
+ sum_bits: AtomicU64,
34
+}
35
+
36
+impl Histogram {
37
+ /// Create a histogram with the given finite bucket upper bounds, in ascending order.
38
+ pub fn new(buckets: impl IntoIterator<Item = f64>) -> Self {
39
+ let upper_bounds: Box<[f64]> = buckets.into_iter().chain(once(f64::MAX)).collect();
40
+ let bucket_counts = upper_bounds.iter().map(|_| AtomicU64::new(0)).collect();
41
+ Self {
42
+ inner: Arc::new(Inner {
43
+ upper_bounds,
44
+ bucket_counts,
45
+ sum_bits: AtomicU64::new(0.0_f64.to_bits()),
46
+ }),
47
+ }
48
+ }
49
+
50
+ pub fn observe(&self, value: f64) {
51
+ let inner = &*self.inner;
52
+ // Unlike `prometheus_client`, which counts `+inf` and NaN in `_count` but in no bucket,
53
+ // put them in the `+Inf` bucket as the Go client does.
54
+ let index = inner
55
+ .upper_bounds
56
+ .iter()
57
+ .position(|bound| value <= *bound)
58
+ .unwrap_or(inner.upper_bounds.len() - 1);
59
+ inner.bucket_counts[index].fetch_add(1, Ordering::Relaxed);
60
+ // The closure always returns `Some`, so the update cannot fail.
61
+ let _ = inner.sum_bits.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |bits| {
62
+ Some((f64::from_bits(bits) + value).to_bits())
63
+ });
64
+ }
65
+}
66
+
67
+impl TypedMetric for Histogram {
68
+ const TYPE: MetricType = MetricType::Histogram;
69
+}
70
+
71
+impl EncodeMetric for Histogram {
72
+ fn encode(&self, mut encoder: MetricEncoder) -> Result<(), fmt::Error> {
73
+ let inner = &*self.inner;
74
+ let buckets: Vec<(f64, u64)> = inner
75
+ .upper_bounds
76
+ .iter()
77
+ .zip(&inner.bucket_counts)
78
+ .map(|(bound, count)| (*bound, count.load(Ordering::Relaxed)))
79
+ .collect();
80
+ let count = buckets.iter().map(|(_, count)| count).sum();
81
+ let sum = f64::from_bits(inner.sum_bits.load(Ordering::Relaxed));
82
+ encoder.encode_histogram::<NoLabelSet>(sum, count, &buckets, None)
83
+ }
84
+
85
+ fn metric_type(&self) -> MetricType {
86
+ Self::TYPE
87
+ }
88
+}
89
+
90
+#[cfg(test)]
91
+mod tests {
92
+ use expect_test::expect;
93
+ use prometheus_client::encoding::text::encode;
94
+ use prometheus_client::registry::Registry;
95
+
96
+ use super::*;
97
+
98
+ const BUCKETS: [f64; 3] = [0.1, 1.0, 10.0];
99
+
100
+ fn render(metric: impl EncodeMetric + fmt::Debug + Send + Sync + 'static) -> String {
101
+ let mut registry = Registry::default();
102
+ registry.register("latency", "Latency", metric);
103
+ let mut output = String::new();
104
+ encode(&mut output, ®istry).unwrap();
105
+ output
106
+ }
107
+
108
+ #[test]
109
+ #[expect(
110
+ clippy::disallowed_types,
111
+ reason = "compare against the replaced implementation"
112
+ )]
113
+ fn finite_observations_match_prometheus_client() {
114
+ let baseline = prometheus_client::metrics::histogram::Histogram::new(BUCKETS);
115
+ let histogram = Histogram::new(BUCKETS);
116
+ let mut values = vec![-1.0, 0.0, 100.0, f64::MAX];
117
+ for bound in BUCKETS {
118
+ values.extend([bound.next_down(), bound, bound.next_up()]);
119
+ }
120
+ for value in values {
121
+ baseline.observe(value);
122
+ histogram.observe(value);
123
+ }
124
+ assert_eq!(render(histogram), render(baseline));
125
+ }
126
+
127
+ #[test]
128
+ fn concurrent_observations_are_not_lost() {
129
+ let histogram = Histogram::new(BUCKETS);
130
+ std::thread::scope(|scope| {
131
+ // Dyadic values keep the sum exact in any addition order.
132
+ for value in [0.0625, 0.5, 4.0, 64.0] {
133
+ let histogram = histogram.clone();
134
+ scope.spawn(move || (0..10_000).for_each(|_| histogram.observe(value)));
135
+ }
136
+ });
137
+ expect![[r##"
138
+ # HELP latency Latency.
139
+ # TYPE latency histogram
140
+ latency_sum 685625.0
141
+ latency_count 40000
142
+ latency_bucket{le="0.1"} 10000
143
+ latency_bucket{le="1.0"} 20000
144
+ latency_bucket{le="10.0"} 30000
145
+ latency_bucket{le="+Inf"} 40000
146
+ # EOF
147
+ "##]]
148
+ .assert_eq(&render(histogram));
149
+ }
150
+
151
+ #[test]
152
+ fn non_finite_observations_count_in_inf_bucket() {
153
+ let histogram = Histogram::new(BUCKETS);
154
+ histogram.observe(f64::INFINITY);
155
+ histogram.observe(f64::NAN);
156
+ expect![[r##"
157
+ # HELP latency Latency.
158
+ # TYPE latency histogram
159
+ latency_sum NaN
160
+ latency_count 2
161
+ latency_bucket{le="0.1"} 0
162
+ latency_bucket{le="1.0"} 0
163
+ latency_bucket{le="10.0"} 0
164
+ latency_bucket{le="+Inf"} 2
165
+ # EOF
166
+ "##]]
167
+ .assert_eq(&render(histogram));
168
+ }
169
+}
rust/src/metrics/src/lib.rs
+2
-1
@@ -9,14 +9,15 @@ use prometheus_client::encoding::text::encode;
9
use prometheus_client::metrics::counter::Counter;
10
pub use prometheus_client::metrics::family::Family;
11
use prometheus_client::metrics::gauge::Gauge;
12
-use prometheus_client::metrics::histogram::Histogram;
12
use prometheus_client::registry::Registry;
13
14
mod api_server;
15
+mod histogram;
16
mod request;
17
mod scheduler;
18
19
pub use api_server::*;
20
+pub use histogram::Histogram;
21
pub use request::*;
22
pub use scheduler::*;
23
rust/src/metrics/src/request.rs
+1
-2
@@ -3,10 +3,9 @@
3
4
use prometheus_client::encoding::EncodeLabelSet;
5
use prometheus_client::metrics::family::Family;
6
-use prometheus_client::metrics::histogram::Histogram;
6
use prometheus_client::registry::Registry;
7
9
-use crate::{EngineLabels, HistogramFamily, U64Counter};
8
+use crate::{EngineLabels, Histogram, HistogramFamily, U64Counter};
9
10
const TTFT_BUCKETS: [f64; 22] = [
11
0.001, 0.005, 0.01, 0.02, 0.04, 0.06, 0.08, 0.1, 0.25, 0.5, 0.75, 1.0, 2.5, 5.0, 7.5, 10.0,
rust/src/metrics/src/scheduler.rs
+1
-2
@@ -7,10 +7,9 @@ use std::sync::{Arc, Mutex};
7
use itertools::Itertools as _;
8
use prometheus_client::encoding::{EncodeLabelSet, EncodeLabelValue, LabelValueEncoder};
9
use prometheus_client::metrics::family::Family;
10
-use prometheus_client::metrics::histogram::Histogram;
10
use prometheus_client::registry::Registry;
11
13
-use crate::{F64Gauge, HistogramFamily, U64Counter, U64Gauge};
12
+use crate::{F64Gauge, Histogram, HistogramFamily, U64Counter, U64Gauge};
13
14
const KV_CACHE_RESIDENCY_BUCKETS: [f64; 21] = [
15
0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 60.0,