Skip to content

Commit 1b65928

Browse files
committed
[Metrics][Spec Decoding] Add per-request acceptance rate Prometheus histogram
Track speculative decoding acceptance at the finished-request level. Per-request draft and accepted token counts are carried through EngineCoreOutput into IterationStats and FinishedRequestStats. Expose vllm:request_spec_decode_acceptance_rate as a Prometheus histogram observing each finished request's accepted/draft ratio. The existing aggregate speculative decoding counters are preserved. The scheduler now builds aggregate and request-level spec decode stats together so invalid draft token handling stays consistent. Add focused metrics stats tests for request-level histogram observation and per-request spec decode stat accumulation. Signed-off-by: jeff.ye <jeff.ye@novita.ai>
1 parent 3b8fc3f commit 1b65928

6 files changed

Lines changed: 220 additions & 15 deletions

File tree

tests/v1/metrics/test_stats.py

Lines changed: 114 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,33 @@
11
# SPDX-License-Identifier: Apache-2.0
22
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3-
from vllm.v1.engine import FinishReason
3+
from unittest.mock import MagicMock
4+
5+
from vllm.config import SpeculativeConfig
6+
from vllm.v1.engine import EngineCoreOutput, FinishReason
47
from vllm.v1.metrics.stats import (
58
IterationStats,
69
PrefillStats,
710
PromptTokenStats,
811
RequestStateStats,
912
)
13+
from vllm.v1.spec_decode.metrics import (
14+
SpecDecodeRequestStats,
15+
SpecDecodingProm,
16+
)
17+
18+
19+
class MockCounter:
20+
def __init__(self, *args, **kwargs):
21+
self.children = {}
22+
23+
def labels(self, *labelvalues):
24+
child = MagicMock()
25+
self.children[labelvalues] = child
26+
return child
27+
28+
29+
class MockHistogram(MockCounter):
30+
pass
1031

1132

1233
def test_iteration_stats_repr():
@@ -122,6 +143,98 @@ def test_prefill_kv_computed_edge_cases():
122143
assert finished_req2.request_id == "test-req-004"
123144

124145

146+
def test_finished_request_includes_spec_decode_stats():
147+
iteration_stats = IterationStats()
148+
req_stats = RequestStateStats(arrival_time=0.0)
149+
req_stats.scheduled_ts = 0.1
150+
req_stats.first_token_ts = 0.5
151+
req_stats.last_token_ts = 1.0
152+
153+
for output in (
154+
EngineCoreOutput(
155+
request_id="test-req-005",
156+
new_token_ids=[1, 2],
157+
spec_decode_stats=SpecDecodeRequestStats(
158+
num_draft_tokens=4,
159+
num_accepted_tokens=1,
160+
),
161+
),
162+
EngineCoreOutput(
163+
request_id="test-req-005",
164+
new_token_ids=[3, 4, 5],
165+
spec_decode_stats=SpecDecodeRequestStats(
166+
num_draft_tokens=4,
167+
num_accepted_tokens=2,
168+
),
169+
),
170+
):
171+
iteration_stats.update_from_output(
172+
output=output,
173+
engine_core_timestamp=2.0,
174+
is_prefilling=False,
175+
req_stats=req_stats,
176+
lora_states=MagicMock(),
177+
lora_name=None,
178+
)
179+
180+
iteration_stats.update_from_finished_request(
181+
finish_reason=FinishReason.STOP,
182+
request_id="test-req-005",
183+
num_prompt_tokens=2000,
184+
max_tokens_param=300,
185+
req_stats=req_stats,
186+
)
187+
188+
finished_req = iteration_stats.finished_requests[0]
189+
assert finished_req.spec_decode_stats is not None
190+
assert finished_req.spec_decode_stats.num_draft_tokens == 8
191+
assert finished_req.spec_decode_stats.num_accepted_tokens == 3
192+
assert finished_req.spec_decode_stats.acceptance_rate == 3 / 8
193+
194+
195+
def test_spec_decode_request_stats_merge():
196+
stats = SpecDecodeRequestStats(
197+
num_draft_tokens=4,
198+
num_accepted_tokens=1,
199+
).merge(
200+
SpecDecodeRequestStats(
201+
num_draft_tokens=4,
202+
num_accepted_tokens=2,
203+
)
204+
)
205+
206+
assert stats.num_draft_tokens == 8
207+
assert stats.num_accepted_tokens == 3
208+
assert stats.acceptance_rate == 3 / 8
209+
assert stats.merge(None) is stats
210+
211+
212+
def test_request_spec_decode_acceptance_rate_observed():
213+
class MockSpecDecodingProm(SpecDecodingProm):
214+
_counter_cls = MockCounter
215+
_histogram_cls = MockHistogram
216+
217+
prom = MockSpecDecodingProm(
218+
speculative_config=SpeculativeConfig(
219+
method="ngram",
220+
num_speculative_tokens=2,
221+
),
222+
labelnames=["model_name", "engine"],
223+
per_engine_labelvalues={0: ["test-model", "0"]},
224+
)
225+
226+
prom.observe_finished_request(
227+
spec_decode_stats=SpecDecodeRequestStats(
228+
num_draft_tokens=200,
229+
num_accepted_tokens=10,
230+
),
231+
)
232+
233+
prom.histogram_request_spec_decode_acceptance_rate[
234+
0
235+
].observe.assert_called_once_with(0.05)
236+
237+
125238
def test_prompt_token_stats_all_computed():
126239
"""Test all tokens computed locally, no caching."""
127240
stats = PromptTokenStats()

vllm/v1/core/sched/scheduler.py

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@
5656
from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats
5757
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
5858
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
59-
from vllm.v1.spec_decode.metrics import SpecDecodingStats
59+
from vllm.v1.spec_decode.metrics import SpecDecodeRequestStats, SpecDecodingStats
6060
from vllm.v1.structured_output import StructuredOutputManager
6161
from vllm.v1.utils import record_function_or_nullcontext
6262

@@ -1473,13 +1473,23 @@ def update_from_output(
14731473
scheduled_spec_token_ids = (
14741474
scheduler_output.scheduled_spec_decode_tokens.get(req_id)
14751475
)
1476+
spec_decode_request_stats: SpecDecodeRequestStats | None = None
14761477
if scheduled_spec_token_ids and (
14771478
generated_token_ids or self.num_sampled_tokens_per_step == 0
14781479
):
14791480
num_draft_tokens = len(scheduled_spec_token_ids)
14801481
num_sampled = self.num_sampled_tokens_per_step
14811482
num_accepted = max(len(generated_token_ids) - num_sampled, 0)
14821483
num_rejected = num_draft_tokens - num_accepted
1484+
spec_decoding_stats, spec_decode_request_stats = (
1485+
self.make_spec_decoding_stats(
1486+
spec_decoding_stats,
1487+
num_draft_tokens=num_draft_tokens,
1488+
num_accepted_tokens=num_accepted,
1489+
num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
1490+
request_id=req_id,
1491+
)
1492+
)
14831493
# num_computed_tokens represents the number of tokens
14841494
# processed in the current step, considering scheduled
14851495
# tokens and rejections. If some tokens are rejected,
@@ -1491,14 +1501,6 @@ def update_from_output(
14911501
# the scheduled spec tokens count and so is similarly adjusted.
14921502
if request.num_output_placeholders > 0:
14931503
request.num_output_placeholders -= num_rejected
1494-
spec_decoding_stats = self.make_spec_decoding_stats(
1495-
spec_decoding_stats,
1496-
num_draft_tokens=num_draft_tokens,
1497-
num_accepted_tokens=num_accepted,
1498-
num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
1499-
request_id=req_id,
1500-
)
1501-
15021504
# Free encoder inputs only after the step has actually executed.
15031505
if request.has_encoder_inputs:
15041506
self._free_encoder_inputs(request)
@@ -1627,6 +1629,7 @@ def update_from_output(
16271629
trace_headers=request.trace_headers,
16281630
routed_experts=routed_experts,
16291631
num_nans_in_logits=request.num_nans_in_logits,
1632+
spec_decode_stats=spec_decode_request_stats,
16301633
)
16311634
)
16321635
else:
@@ -2157,17 +2160,24 @@ def make_spec_decoding_stats(
21572160
num_accepted_tokens: int,
21582161
num_invalid_spec_tokens: dict[str, int] | None,
21592162
request_id: str,
2160-
) -> SpecDecodingStats | None:
2163+
) -> tuple[SpecDecodingStats | None, SpecDecodeRequestStats | None]:
21612164
if not self.log_stats or not num_draft_tokens:
2162-
return None
2165+
return None, None
21632166
if spec_decoding_stats is None:
21642167
spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens)
21652168
if num_invalid_spec_tokens:
21662169
num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0)
21672170
spec_decoding_stats.observe_draft(
2168-
num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens
2171+
num_draft_tokens=num_draft_tokens,
2172+
num_accepted_tokens=num_accepted_tokens,
2173+
)
2174+
return (
2175+
spec_decoding_stats,
2176+
SpecDecodeRequestStats(
2177+
num_draft_tokens=num_draft_tokens,
2178+
num_accepted_tokens=num_accepted_tokens,
2179+
),
21692180
)
2170-
return spec_decoding_stats
21712181

21722182
def shutdown(self) -> None:
21732183
logger.debug_once("[shutdown] Scheduler: start")

vllm/v1/engine/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from vllm.v1.metrics.stats import PrefillStats, SchedulerStats
1919
from vllm.v1.outputs import LogprobsLists, LogprobsTensors
2020
from vllm.v1.serial_utils import UtilityResult
21+
from vllm.v1.spec_decode.metrics import SpecDecodeRequestStats
2122

2223
# Type for pause_generation mode parameter.
2324
# - "abort": Abort all in-flight requests immediately (default).
@@ -197,6 +198,7 @@ class EngineCoreOutput(
197198
# The number of NaNs in logits.
198199
# A value greater than 0 indicates that the output is corrupted.
199200
num_nans_in_logits: int = 0
201+
spec_decode_stats: SpecDecodeRequestStats | None = None
200202

201203
@property
202204
def finished(self) -> bool:

vllm/v1/metrics/loggers.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1221,6 +1221,11 @@ def record(
12211221
finished_request.max_tokens_param
12221222
)
12231223

1224+
self.spec_decoding_prom.observe_finished_request(
1225+
spec_decode_stats=finished_request.spec_decode_stats,
1226+
engine_idx=engine_idx,
1227+
)
1228+
12241229
def record_sleep_state(self, sleep: int = 0, level: int = 0):
12251230
awake = 1
12261231
discard_all = 0

vllm/v1/metrics/stats.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import vllm.envs as envs
1010
from vllm.compilation.cuda_graph import CUDAGraphStat
1111
from vllm.v1.metrics.perf import PerfStats
12-
from vllm.v1.spec_decode.metrics import SpecDecodingStats
12+
from vllm.v1.spec_decode.metrics import SpecDecodeRequestStats, SpecDecodingStats
1313

1414
if TYPE_CHECKING:
1515
from vllm.v1.engine import EngineCoreEvent, EngineCoreOutput, FinishReason
@@ -218,6 +218,7 @@ class RequestStateStats:
218218

219219
# Track if this request is corrupted (NaNs in logits)
220220
is_corrupted: bool = False
221+
spec_decode_stats: SpecDecodeRequestStats | None = None
221222

222223

223224
@dataclass
@@ -237,6 +238,7 @@ class FinishedRequestStats:
237238
mean_time_per_output_token: float = 0.0
238239
is_corrupted: bool = False
239240
num_cached_tokens: int = 0
241+
spec_decode_stats: SpecDecodeRequestStats | None = None
240242

241243

242244
@dataclass
@@ -371,6 +373,10 @@ def update_from_output(
371373
req_stats.first_token_latency = first_token_latency
372374

373375
req_stats.num_generation_tokens += num_new_generation_tokens
376+
if output.spec_decode_stats is not None:
377+
req_stats.spec_decode_stats = output.spec_decode_stats.merge(
378+
req_stats.spec_decode_stats
379+
)
374380

375381
# Track if this request is corrupted (only check once per request)
376382
# Early exit if already marked as corrupted to avoid redundant checks
@@ -472,6 +478,7 @@ def update_from_finished_request(
472478
mean_time_per_output_token=mean_time_per_output_token,
473479
is_corrupted=req_stats.is_corrupted,
474480
num_cached_tokens=num_cached_tokens,
481+
spec_decode_stats=req_stats.spec_decode_stats,
475482
)
476483
self.finished_requests.append(finished_req)
477484

vllm/v1/spec_decode/metrics.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import time
55
from dataclasses import dataclass, field
6+
from typing import NamedTuple
67

78
import numpy as np
89
import prometheus_client
@@ -14,6 +15,27 @@
1415
logger = init_logger(__name__)
1516

1617

18+
class SpecDecodeRequestStats(NamedTuple):
19+
"""Request-level speculative decoding stats."""
20+
21+
num_accepted_tokens: int
22+
num_draft_tokens: int
23+
24+
@property
25+
def acceptance_rate(self) -> float:
26+
if self.num_draft_tokens == 0:
27+
return float("nan")
28+
return self.num_accepted_tokens / self.num_draft_tokens
29+
30+
def merge(self, other: "SpecDecodeRequestStats | None") -> "SpecDecodeRequestStats":
31+
if other is None:
32+
return self
33+
return SpecDecodeRequestStats(
34+
num_accepted_tokens=(self.num_accepted_tokens + other.num_accepted_tokens),
35+
num_draft_tokens=self.num_draft_tokens + other.num_draft_tokens,
36+
)
37+
38+
1739
@dataclass
1840
class SpecDecodingStats:
1941
"""Per-step iteration decoding stats from scheduler.
@@ -192,6 +214,7 @@ class SpecDecodingProm:
192214
"""
193215

194216
_counter_cls = prometheus_client.Counter
217+
_histogram_cls = prometheus_client.Histogram
195218

196219
def __init__(
197220
self,
@@ -208,6 +231,34 @@ def __init__(
208231
if not self.spec_decoding_enabled:
209232
return
210233

234+
if not is_diffusion:
235+
histogram_request_acceptance_rate = self._histogram_cls(
236+
name="vllm:request_spec_decode_acceptance_rate",
237+
documentation=(
238+
"Histogram of request-level speculative decoding draft "
239+
"acceptance rates for finished requests."
240+
),
241+
buckets=[
242+
0.0,
243+
0.1,
244+
0.2,
245+
0.3,
246+
0.4,
247+
0.5,
248+
0.6,
249+
0.7,
250+
0.8,
251+
0.9,
252+
1.0,
253+
],
254+
labelnames=labelnames,
255+
)
256+
self.histogram_request_spec_decode_acceptance_rate = (
257+
create_metric_per_engine(
258+
histogram_request_acceptance_rate, per_engine_labelvalues
259+
)
260+
)
261+
211262
if is_diffusion:
212263
counter_specs = [
213264
("vllm:diffusion_num_denoising_steps", "Number of denoising steps."),
@@ -275,3 +326,20 @@ def observe(self, spec_decoding_stats: SpecDecodingStats, engine_idx: int = 0):
275326
self.counter_spec_decode_num_accepted_tokens_per_pos.get(engine_idx, [])
276327
):
277328
counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos])
329+
330+
def observe_finished_request(
331+
self,
332+
spec_decode_stats: SpecDecodeRequestStats | None,
333+
engine_idx: int = 0,
334+
) -> None:
335+
if (
336+
not self.spec_decoding_enabled
337+
or self.is_diffusion
338+
or spec_decode_stats is None
339+
or spec_decode_stats.num_draft_tokens <= 0
340+
):
341+
return
342+
343+
self.histogram_request_spec_decode_acceptance_rate[engine_idx].observe(
344+
spec_decode_stats.acceptance_rate
345+
)

0 commit comments

Comments
 (0)