A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 22:45 UTC
[Metrics] Add labeled prompt token metrics for P/D disaggregation (#33290)
Add labeled Prometheus metrics to distinguish where prompt tokens come from in P/D disaggregated deployments. In P/D disaggregation, decode instances receive KV cache from prefill instances. Currently, decode reports inflated prompt throughput because it counts all prompt tokens as "computed", even though most were transferred. This PR adds labeled metrics so users can understand actual compute work vs transferred work: vllm:prompt_tokens_by_source_total{source="local_compute"} # Tokens prefilled locally vllm:prompt_tokens_by_source_total{source="external_kv_transfer"} # Tokens received via KV transfer vllm:prompt_tokens_by_source_total{source="local_cache_hit"} # Tokens from local prefix cache vllm:prompt_tokens_cached_total # Total cached (local + external, -1 when all Signed-off-by: Zhanqiu Hu <zh338@cornell.edu>
zhanqiuhu committed
Feb 4, 2026 at 02:46 UTC
4403e3ed4c880365978d1716e5f3a8dbe6a6af31
5 files changed
+227
-5
tests/v1/metrics/test_stats.py
+103
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from vllm.v1.engine import FinishReason
4
-from vllm.v1.metrics.stats import IterationStats, RequestStateStats
4
+from vllm.v1.metrics.stats import IterationStats, PromptTokenStats, RequestStateStats
5
6
7
def test_iteration_stats_repr():
@@ -107,3 +107,105 @@ def test_prefill_kv_computed_edge_cases():
107
finished_req2.num_cached_tokens, 0
108
)
109
assert prefill_kv_computed2 == 0 # All cached, nothing computed
110
+
111
+
112
+def test_prompt_token_stats_all_computed():
113
+ """Test all tokens computed locally, no caching."""
114
+ stats = PromptTokenStats()
115
+
116
+ # Case 1: No caching (All tokens computed locally)
117
+ stats.update_from_output(
118
+ num_cached_tokens=0,
119
+ num_external_computed_tokens=0,
120
+ prompt_len=1000,
121
+ )
122
+
123
+ assert stats.computed == 1000
124
+ assert stats.local_cache_hit == 0
125
+ assert stats.external_kv_transfer == 0
126
+ assert stats.total == 1000
127
+
128
+
129
+def test_prompt_token_stats_partial_local_cache():
130
+ """Test partial local prefix cache hit."""
131
+ stats = PromptTokenStats()
132
+
133
+ # Case 2: Partial local cache
134
+ stats.update_from_output(
135
+ num_cached_tokens=300,
136
+ num_external_computed_tokens=0,
137
+ prompt_len=1000,
138
+ )
139
+
140
+ assert stats.computed == 700
141
+ assert stats.local_cache_hit == 300
142
+ assert stats.external_kv_transfer == 0
143
+
144
+
145
+def test_prompt_token_stats_partial_external_transfer():
146
+ """Test partial external KV transfer."""
147
+ stats = PromptTokenStats()
148
+
149
+ # Case 3: Partial external transfer
150
+ stats.update_from_output(
151
+ num_cached_tokens=500,
152
+ num_external_computed_tokens=500,
153
+ prompt_len=1000,
154
+ )
155
+
156
+ assert stats.computed == 500
157
+ assert stats.local_cache_hit == 0
158
+ assert stats.external_kv_transfer == 500
159
+
160
+
161
+def test_prompt_token_stats_mixed_sources():
162
+ """Test mix of local cache and external transfer."""
163
+ stats = PromptTokenStats()
164
+
165
+ # Case 4: Mixed sources
166
+ stats.update_from_output(
167
+ num_cached_tokens=600,
168
+ num_external_computed_tokens=200,
169
+ prompt_len=1000,
170
+ )
171
+
172
+ assert stats.computed == 400
173
+ assert stats.local_cache_hit == 400
174
+ assert stats.external_kv_transfer == 200
175
+
176
+
177
+def test_prompt_token_stats_full_local_cache_recompute():
178
+ """Test full local cache triggers last token recomputation.
179
+
180
+ When all tokens are cached, the scheduler reduces num_cached_tokens by 1
181
+ to force the model to recompute the last token.
182
+ """
183
+ stats = PromptTokenStats()
184
+
185
+ # Case 5: Full local cache (999 cached after reduction, 1 recomputed)
186
+ stats.update_from_output(
187
+ num_cached_tokens=999,
188
+ num_external_computed_tokens=0,
189
+ prompt_len=1000,
190
+ )
191
+
192
+ assert stats.computed == 1
193
+ assert stats.local_cache_hit == 1000
194
+ assert stats.recomputed_tokens == 1
195
+
196
+
197
+def test_prompt_token_stats_full_external_transfer_recompute():
198
+ """Test full external transfer triggers last token recomputation."""
199
+ stats = PromptTokenStats()
200
+
201
+ # Case 6: Full external transfer (999 cached after reduction, 1 recomputed)
202
+ stats.update_from_output(
203
+ num_cached_tokens=999,
204
+ num_external_computed_tokens=1000,
205
+ prompt_len=1000,
206
+ )
207
+
208
+ assert stats.computed == 1
209
+ assert stats.local_cache_hit == 0
210
+ assert stats.external_kv_transfer == 1000
211
+ assert stats.recomputed_tokens == 1
vllm/v1/core/sched/scheduler.py
+1
@@ -1378,6 +1378,7 @@ class Scheduler(SchedulerInterface):
1378
kv_transfer_params=kv_transfer_params,
1379
trace_headers=request.trace_headers,
1380
num_cached_tokens=request.num_cached_tokens,
1381
+ num_external_computed_tokens=request.num_external_computed_tokens,
1382
routed_experts=routed_experts,
1383
num_nans_in_logits=request.num_nans_in_logits,
1384
)
vllm/v1/engine/__init__.py
+3
-1
@@ -139,8 +139,10 @@ class EngineCoreOutput(
139
kv_transfer_params: dict[str, Any] | None = None
140
141
trace_headers: Mapping[str, str] | None = None
142
- # The number of tokens with prefix cache hits.
142
+ # The number of tokens with prefix cache hits (local + external).
143
num_cached_tokens: int = 0
144
+ # The number of tokens computed remotely (original count from connector).
145
+ num_external_computed_tokens: int = 0
146
routed_experts: np.ndarray | None = None
147
# The number of NaNs in logits.
148
# A value greater than 0 indicates that the output is corrupted.
vllm/v1/metrics/loggers.py
+46
-1
@@ -25,6 +25,7 @@ from vllm.v1.metrics.stats import (
25
CachingMetrics,
26
IterationStats,
27
MultiModalCacheStats,
28
+ PromptTokenStats,
29
SchedulerStats,
30
)
31
from vllm.v1.spec_decode.metrics import SpecDecodingLogging, SpecDecodingProm
@@ -136,7 +137,8 @@ class LoggingStatLogger(StatLoggerBase):
137
138
def _track_iteration_stats(self, iteration_stats: IterationStats):
139
# Save tracked stats for token counters.
139
- self.num_prompt_tokens += iteration_stats.num_prompt_tokens
140
+ # Use computed tokens for prompt throughput (excludes cached/transferred)
141
+ self.num_prompt_tokens += iteration_stats.prompt_token_stats.computed
142
self.num_generation_tokens += iteration_stats.num_generation_tokens
143
self.num_corrupted_reqs += iteration_stats.num_corrupted_reqs
144
self.num_preemptions += iteration_stats.num_preempted_reqs
@@ -590,6 +592,41 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
592
counter_prompt_tokens, engine_indexes, model_name
593
)
594
595
+ # Labeled prompt token counters by source
596
+ counter_prompt_tokens_by_source = self._counter_cls(
597
+ name="vllm:prompt_tokens_by_source",
598
+ documentation="Number of prompt tokens by source.",
599
+ labelnames=labelnames + ["source"],
600
+ )
601
+ self.counter_prompt_tokens_by_source: dict[str, dict[int, Counter]] = {}
602
+ for source in PromptTokenStats.ALL_SOURCES:
603
+ self.counter_prompt_tokens_by_source[source] = {
604
+ idx: counter_prompt_tokens_by_source.labels(
605
+ model_name, str(idx), source
606
+ )
607
+ for idx in engine_indexes
608
+ }
609
+
610
+ # Cached prompt tokens counter
611
+ counter_prompt_tokens_cached = self._counter_cls(
612
+ name="vllm:prompt_tokens_cached",
613
+ documentation="Number of cached prompt tokens (local + external).",
614
+ labelnames=labelnames,
615
+ )
616
+ self.counter_prompt_tokens_cached = make_per_engine(
617
+ counter_prompt_tokens_cached, engine_indexes, model_name
618
+ )
619
+
620
+ # Recomputed tokens (last token recomputed when entire prompt is cached)
621
+ counter_prompt_tokens_recomputed = self._counter_cls(
622
+ name="vllm:prompt_tokens_recomputed",
623
+ documentation="Number of cached tokens recomputed for forward pass.",
624
+ labelnames=labelnames,
625
+ )
626
+ self.counter_prompt_tokens_recomputed = make_per_engine(
627
+ counter_prompt_tokens_recomputed, engine_indexes, model_name
628
+ )
629
+
630
counter_generation_tokens = self._counter_cls(
631
name="vllm:generation_tokens",
632
documentation="Number of generation tokens processed.",
@@ -1070,6 +1107,14 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
1107
iteration_stats.num_preempted_reqs
1108
)
1109
self.counter_prompt_tokens[engine_idx].inc(iteration_stats.num_prompt_tokens)
1110
+ # Labeled prompt token counters by source
1111
+ pts = iteration_stats.prompt_token_stats
1112
+ for source in PromptTokenStats.ALL_SOURCES:
1113
+ self.counter_prompt_tokens_by_source[source][engine_idx].inc(
1114
+ pts.get_by_source(source)
1115
+ )
1116
+ self.counter_prompt_tokens_cached[engine_idx].inc(pts.cached_tokens)
1117
+ self.counter_prompt_tokens_recomputed[engine_idx].inc(pts.recomputed_tokens)
1118
self.counter_generation_tokens[engine_idx].inc(
1119
iteration_stats.num_generation_tokens
1120
)
vllm/v1/metrics/stats.py
+74
-2
@@ -231,13 +231,76 @@ class FinishedRequestStats:
231
num_cached_tokens: int = 0
232
233
234
+@dataclass
235
+class PromptTokenStats:
236
+ """Breakdown of prompt tokens by source.
237
+
238
+ Fields:
239
+ computed: Tokens prefilled locally (actual compute work).
240
+ local_cache_hit: Tokens from local prefix cache.
241
+ external_kv_transfer: Tokens from external KV transfer.
242
+ cached_tokens: Tokens skipped during prefill (from scheduler).
243
+ recomputed_tokens: Cached tokens that were recomputed (see below).
244
+ total: Total prompt tokens.
245
+
246
+ Invariants:
247
+ computed + local_cache_hit + external_kv_transfer - recomputed_tokens = total
248
+ local_cache_hit + external_kv_transfer - recomputed_tokens = cached_tokens
249
+ """
250
+
251
+ ALL_SOURCES: tuple[str, ...] = (
252
+ "local_compute",
253
+ "local_cache_hit",
254
+ "external_kv_transfer",
255
+ )
256
+
257
+ computed: int = 0
258
+ local_cache_hit: int = 0
259
+ external_kv_transfer: int = 0
260
+ cached_tokens: int = 0
261
+ recomputed_tokens: int = 0
262
+ total: int = 0
263
+
264
+ def update_from_output(
265
+ self,
266
+ num_cached_tokens: int,
267
+ num_external_computed_tokens: int,
268
+ prompt_len: int,
269
+ ) -> None:
270
+ """Update stats from a prefill output."""
271
+ # When all tokens are cached, the scheduler reduces num_cached_tokens
272
+ # by 1 to force the model to recompute the last token, since the model
273
+ # needs at least one input token to run a forward pass.
274
+ recomputed = 1 if (num_cached_tokens + 1 == prompt_len) else 0
275
+
276
+ self.computed += prompt_len - num_cached_tokens
277
+ self.external_kv_transfer += num_external_computed_tokens
278
+ self.local_cache_hit += (
279
+ num_cached_tokens + recomputed - num_external_computed_tokens
280
+ )
281
+ self.cached_tokens += num_cached_tokens
282
+ self.recomputed_tokens += recomputed
283
+ self.total += prompt_len
284
+
285
+ def get_by_source(self, source: str) -> int:
286
+ """Get token count by source label."""
287
+ source_map = {
288
+ "local_compute": self.computed,
289
+ "local_cache_hit": self.local_cache_hit,
290
+ "external_kv_transfer": self.external_kv_transfer,
291
+ }
292
+ if source not in source_map:
293
+ raise ValueError(f"Unknown source: {source}")
294
+ return source_map[source]
295
+
296
+
297
class IterationStats:
298
"""Stats associated with a single set of EngineCoreOutputs."""
299
300
def __init__(self):
301
self.iteration_timestamp = time.time()
302
self.num_generation_tokens = 0
240
- self.num_prompt_tokens = 0
303
+ self.prompt_token_stats = PromptTokenStats()
304
self.num_preempted_reqs = 0
305
self.finished_requests: list[FinishedRequestStats] = []
306
self.max_num_generation_tokens_iter: list[int] = []
@@ -250,6 +313,11 @@ class IterationStats:
313
field_to_value_str = ", ".join(f"{k}={v}" for k, v in vars(self).items())
314
return f"{self.__class__.__name__}({field_to_value_str})"
315
316
+ @property
317
+ def num_prompt_tokens(self) -> int:
318
+ """Total prompt tokens (for backward compatibility)."""
319
+ return self.prompt_token_stats.total
320
+
321
def _time_since(self, start: float) -> float:
322
"""Calculate an interval relative to this iteration's timestamp."""
323
return self.iteration_timestamp - start
@@ -268,7 +336,11 @@ class IterationStats:
336
337
self.num_generation_tokens += num_new_generation_tokens
338
if is_prefilling:
271
- self.num_prompt_tokens += prompt_len
339
+ self.prompt_token_stats.update_from_output(
340
+ num_cached_tokens=output.num_cached_tokens,
341
+ num_external_computed_tokens=output.num_external_computed_tokens,
342
+ prompt_len=prompt_len,
343
+ )
344
345
first_token_latency = self._time_since(req_stats.arrival_time)
346
self.time_to_first_tokens_iter.append(first_token_latency)