Skip to content

Commit e962167

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 2c764c0 commit e962167

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
@@ -57,7 +57,7 @@
5757
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
5858
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
5959
from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
60-
from vllm.v1.spec_decode.metrics import SpecDecodingStats
60+
from vllm.v1.spec_decode.metrics import SpecDecodeRequestStats, SpecDecodingStats
6161
from vllm.v1.structured_output import StructuredOutputManager
6262
from vllm.v1.utils import record_function_or_nullcontext
6363

@@ -1489,13 +1489,23 @@ def update_from_output(
14891489
scheduled_spec_token_ids = (
14901490
scheduler_output.scheduled_spec_decode_tokens.get(req_id)
14911491
)
1492+
spec_decode_request_stats: SpecDecodeRequestStats | None = None
14921493
if scheduled_spec_token_ids and (
14931494
generated_token_ids or self.num_sampled_tokens_per_step == 0
14941495
):
14951496
num_draft_tokens = len(scheduled_spec_token_ids)
14961497
num_sampled = self.num_sampled_tokens_per_step
14971498
num_accepted = max(len(generated_token_ids) - num_sampled, 0)
14981499
num_rejected = num_draft_tokens - num_accepted
1500+
spec_decoding_stats, spec_decode_request_stats = (
1501+
self.make_spec_decoding_stats(
1502+
spec_decoding_stats,
1503+
num_draft_tokens=num_draft_tokens,
1504+
num_accepted_tokens=num_accepted,
1505+
num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
1506+
request_id=req_id,
1507+
)
1508+
)
14991509
# num_computed_tokens represents the number of tokens
15001510
# processed in the current step, considering scheduled
15011511
# tokens and rejections. If some tokens are rejected,
@@ -1507,14 +1517,6 @@ def update_from_output(
15071517
# the scheduled spec tokens count and so is similarly adjusted.
15081518
if request.num_output_placeholders > 0:
15091519
request.num_output_placeholders -= num_rejected
1510-
spec_decoding_stats = self.make_spec_decoding_stats(
1511-
spec_decoding_stats,
1512-
num_draft_tokens=num_draft_tokens,
1513-
num_accepted_tokens=num_accepted,
1514-
num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
1515-
request_id=req_id,
1516-
)
1517-
15181520
# Free encoder inputs only after the step has actually executed.
15191521
if request.has_encoder_inputs:
15201522
self._free_encoder_inputs(request)
@@ -1643,6 +1645,7 @@ def update_from_output(
16431645
trace_headers=request.trace_headers,
16441646
routed_experts=routed_experts,
16451647
num_nans_in_logits=request.num_nans_in_logits,
1648+
spec_decode_stats=spec_decode_request_stats,
16461649
)
16471650
)
16481651
else:
@@ -2178,17 +2181,24 @@ def make_spec_decoding_stats(
21782181
num_accepted_tokens: int,
21792182
num_invalid_spec_tokens: dict[str, int] | None,
21802183
request_id: str,
2181-
) -> SpecDecodingStats | None:
2184+
) -> tuple[SpecDecodingStats | None, SpecDecodeRequestStats | None]:
21822185
if not self.log_stats or not num_draft_tokens:
2183-
return None
2186+
return None, None
21842187
if spec_decoding_stats is None:
21852188
spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens)
21862189
if num_invalid_spec_tokens:
21872190
num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0)
21882191
spec_decoding_stats.observe_draft(
2189-
num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens
2192+
num_draft_tokens=num_draft_tokens,
2193+
num_accepted_tokens=num_accepted_tokens,
2194+
)
2195+
return (
2196+
spec_decoding_stats,
2197+
SpecDecodeRequestStats(
2198+
num_draft_tokens=num_draft_tokens,
2199+
num_accepted_tokens=num_accepted_tokens,
2200+
),
21902201
)
2191-
return spec_decoding_stats
21922202

21932203
def shutdown(self) -> None:
21942204
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.
@@ -196,6 +218,7 @@ class SpecDecodingProm:
196218
"""
197219

198220
_counter_cls = prometheus_client.Counter
221+
_histogram_cls = prometheus_client.Histogram
199222

200223
def __init__(
201224
self,
@@ -212,6 +235,34 @@ def __init__(
212235
if not self.spec_decoding_enabled:
213236
return
214237

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

0 commit comments

Comments
 (0)