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 15, 13:30 UTC
[V1][Spec Decoding] Add num_drafts and num_accepted_tokens_per_position metrics (#16665)
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Mark McLoughlin committed
Apr 24, 2025 at 16:57 UTC
340d7b1b217794c1b89e13936c66b920fd8d842d
4 files changed
+158
-60
tests/v1/core/test_scheduler.py
+25
-14
@@ -6,7 +6,7 @@ import pytest
6
import torch
7
8
from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,
9
- SchedulerConfig, VllmConfig)
9
+ SchedulerConfig, SpeculativeConfig, VllmConfig)
10
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
11
from vllm.sampling_params import SamplingParams
12
from vllm.v1.core.sched.output import SchedulerOutput
@@ -31,6 +31,7 @@ def create_scheduler(
31
num_blocks: int = 10000,
32
block_size: int = 16,
33
max_model_len: Optional[int] = None,
34
+ num_speculative_tokens: Optional[int] = None,
35
) -> Scheduler:
36
'''Create scheduler under test.
37
@@ -81,11 +82,17 @@ def create_scheduler(
82
kv_connector_extra_config={"shared_storage_path": "local_storage"},
83
) if use_kv_connector else None
84
85
+ speculative_config: Optional[SpeculativeConfig] = None
86
+ if num_speculative_tokens is not None:
87
+ speculative_config = SpeculativeConfig(
88
+ model="ngram", num_speculative_tokens=num_speculative_tokens)
89
+
90
vllm_config = VllmConfig(
91
scheduler_config=scheduler_config,
92
model_config=model_config,
93
cache_config=cache_config,
94
kv_transfer_config=kv_transfer_config,
95
+ speculative_config=speculative_config,
96
)
97
kv_cache_config = KVCacheConfig(
98
num_blocks=num_blocks, # A large number of blocks to hold all requests
@@ -429,7 +436,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
436
437
def test_stop_via_update_from_output():
438
"""Test stopping behavior through update_from_output"""
432
- scheduler = create_scheduler()
439
+ scheduler = create_scheduler(num_speculative_tokens=1)
440
441
# Test case 1: Stop on EOS token
442
requests = create_requests(num_requests=2, max_tokens=10)
@@ -480,7 +487,7 @@ def test_stop_via_update_from_output():
487
assert list(requests[1].output_token_ids) == [10, 11]
488
489
# Test case 2: Stop on custom stop token
483
- scheduler = create_scheduler()
490
+ scheduler = create_scheduler(num_speculative_tokens=2)
491
requests = create_requests(num_requests=2,
492
max_tokens=10,
493
stop_token_ids=[42, 43])
@@ -531,7 +538,7 @@ def test_stop_via_update_from_output():
538
assert list(requests[1].output_token_ids) == [13, 14]
539
540
# Test case 3: Stop on max tokens
534
- scheduler = create_scheduler()
541
+ scheduler = create_scheduler(num_speculative_tokens=2)
542
requests = create_requests(num_requests=2, max_tokens=2)
543
for req in requests:
544
req.num_computed_tokens = req.num_tokens
@@ -580,7 +587,7 @@ def test_stop_via_update_from_output():
587
assert list(requests[1].output_token_ids) == [13]
588
589
# Test case 4: Ignore EOS flag
583
- scheduler = create_scheduler()
590
+ scheduler = create_scheduler(num_speculative_tokens=2)
591
requests = create_requests(num_requests=1, max_tokens=10)
592
requests[0].sampling_params.ignore_eos = True
593
requests[0].num_computed_tokens = requests[0].num_tokens
@@ -682,13 +689,14 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],
689
@pytest.mark.parametrize(
690
"spec_tokens,output_tokens,expected",
691
[
685
- ([[1, 2, 3]], [[1, 2, 3, 4]], (3, 3)), # perfect match
686
- ([[1, 2, 3]], [[1, 5]], (3, 1)), # early mismatch
687
- ([[1, 2], [3]], [[1, 2, 5], [3, 4]], (3, 3)), # multiple sequences
688
- ([[1]], [[1, 2]], (1, 1)), # single token sequence
689
- ([[]], [[5]], (0, 0)), # empty sequence
692
+ ([[1, 2, 3]], [[1, 2, 3, 4]], (1, 3, 3, [1, 1, 1])), # perfect match
693
+ ([[1, 2, 3]], [[1, 5]], (1, 3, 1, [1, 0, 0])), # early mismatch
694
+ ([[1, 2], [3]], [[1, 2, 5], [3, 4]],
695
+ (2, 3, 3, [2, 1])), # multiple sequences
696
+ ([[1]], [[1, 2]], (1, 1, 1, [1])), # single token sequence
697
+ ([[]], [[5]], (0, 0, 0, [0])), # empty sequence
698
([[1, 2, 3], [4, 5, 6]], [[1, 2, 7], [4, 8]],
691
- (6, 3)), # multiple mismatches
699
+ (2, 6, 3, [2, 1, 0])), # multiple mismatches
700
])
701
def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
702
"""Test scheduling behavior with speculative decoding.
@@ -697,7 +705,8 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
705
1. Speculated tokens get scheduled correctly
706
2. Spec decoding stats properly count number of draft and accepted tokens
707
"""
700
- scheduler = create_scheduler()
708
+ num_spec_tokens = max(1, max(len(t) for t in spec_tokens))
709
+ scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens)
710
requests = create_requests(num_requests=len(spec_tokens), num_tokens=1)
711
req_ids = []
712
req_to_index = {}
@@ -770,8 +779,10 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
779
else:
780
assert scheduler_stats.spec_decoding_stats is not None
781
stats = scheduler_stats.spec_decoding_stats
773
- assert stats.num_draft_tokens == expected[0]
774
- assert stats.num_accepted_tokens == expected[1]
782
+ assert stats.num_drafts == expected[0]
783
+ assert stats.num_draft_tokens == expected[1]
784
+ assert stats.num_accepted_tokens == expected[2]
785
+ assert stats.num_accepted_tokens_per_pos == expected[3]
786
787
788
def _assert_right_scheduler_output(
vllm/v1/core/sched/scheduler.py
+9
-7
@@ -122,11 +122,12 @@ class Scheduler(SchedulerInterface):
122
self.encoder_cache_manager = EncoderCacheManager(
123
cache_size=encoder_cache_size)
124
125
- self.num_lookahead_tokens = 0
125
speculative_config = vllm_config.speculative_config
127
- if speculative_config and speculative_config.method == "eagle":
128
- self.num_lookahead_tokens = \
129
- speculative_config.num_speculative_tokens
126
+ self.num_spec_tokens = self.num_lookahead_tokens = 0
127
+ if speculative_config:
128
+ self.num_spec_tokens = speculative_config.num_speculative_tokens
129
+ if speculative_config.method == "eagle":
130
+ self.num_lookahead_tokens = self.num_spec_tokens
131
132
def schedule(self) -> SchedulerOutput:
133
# NOTE(woosuk) on the scheduling algorithm:
@@ -824,7 +825,8 @@ class Scheduler(SchedulerInterface):
825
if not self.log_stats:
826
return None
827
if spec_decoding_stats is None:
827
- spec_decoding_stats = SpecDecodingStats()
828
- spec_decoding_stats.observe(num_draft_tokens=num_draft_tokens,
829
- num_accepted_tokens=num_accepted_tokens)
828
+ spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens)
829
+ spec_decoding_stats.observe_draft(
830
+ num_draft_tokens=num_draft_tokens,
831
+ num_accepted_tokens=num_accepted_tokens)
832
return spec_decoding_stats
vllm/v1/metrics/loggers.py
+9
-26
@@ -12,7 +12,7 @@ from vllm.logger import init_logger
12
from vllm.v1.core.kv_cache_utils import PrefixCachingMetrics
13
from vllm.v1.engine import FinishReason
14
from vllm.v1.metrics.stats import IterationStats, SchedulerStats
15
-from vllm.v1.spec_decode.metrics import SpecDecodingMetrics
15
+from vllm.v1.spec_decode.metrics import SpecDecodingLogging, SpecDecodingProm
16
17
logger = init_logger(__name__)
18
@@ -39,7 +39,7 @@ class LoggingStatLogger(StatLoggerBase):
39
# Prefix cache metrics. This cannot be reset.
40
# TODO: Make the interval configurable.
41
self.prefix_caching_metrics = PrefixCachingMetrics()
42
- self.spec_decoding_metrics = SpecDecodingMetrics()
42
+ self.spec_decoding_logging = SpecDecodingLogging()
43
self.last_prompt_throughput: float = 0.0
44
self.last_generation_throughput: float = 0.0
45
@@ -70,7 +70,7 @@ class LoggingStatLogger(StatLoggerBase):
70
self.prefix_caching_metrics.observe(scheduler_stats.prefix_cache_stats)
71
72
if scheduler_stats.spec_decoding_stats is not None:
73
- self.spec_decoding_metrics.observe(
73
+ self.spec_decoding_logging.observe(
74
scheduler_stats.spec_decoding_stats)
75
76
self.last_scheduler_stats = scheduler_stats
@@ -112,7 +112,7 @@ class LoggingStatLogger(StatLoggerBase):
112
)
113
114
if scheduler_stats.spec_decoding_stats is not None:
115
- self.spec_decoding_metrics.log(log_fn=log_fn)
115
+ self.spec_decoding_logging.log(log_fn=log_fn)
116
117
118
class PrometheusStatLogger(StatLoggerBase):
@@ -133,6 +133,9 @@ class PrometheusStatLogger(StatLoggerBase):
133
134
max_model_len = vllm_config.model_config.max_model_len
135
136
+ self.spec_decoding_prom = SpecDecodingProm(
137
+ vllm_config.speculative_config, labelnames, labelvalues)
138
+
139
#
140
# Scheduler state
141
#
@@ -323,24 +326,6 @@ class PrometheusStatLogger(StatLoggerBase):
326
self.labelname_running_lora_adapters,
327
])
328
326
- #
327
- # Speculative Decoding metrics
328
- # The acceptance rate can be calculated using a PromQL query:
329
- #
330
- # rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) /
331
- # rate(vllm:spec_decode_num_draft_tokens_total[$interval])
332
- #
333
- self.counter_spec_decode_num_draft_tokens = \
334
- prometheus_client.Counter(
335
- name="vllm:spec_decode_num_draft_tokens_total",
336
- documentation="Number of draft tokens.",
337
- labelnames=labelnames).labels(*labelvalues)
338
- self.counter_spec_decode_num_accepted_tokens = \
339
- prometheus_client.Counter(
340
- name="vllm:spec_decode_num_accepted_tokens_total",
341
- documentation="Number of accepted tokens.",
342
- labelnames=labelnames).labels(*labelvalues)
343
-
329
#
330
# Cache config info metric
331
#
@@ -378,10 +363,8 @@ class PrometheusStatLogger(StatLoggerBase):
363
scheduler_stats.prefix_cache_stats.hits)
364
365
if scheduler_stats.spec_decoding_stats is not None:
381
- self.counter_spec_decode_num_draft_tokens.inc(
382
- scheduler_stats.spec_decoding_stats.num_draft_tokens)
383
- self.counter_spec_decode_num_accepted_tokens.inc(
384
- scheduler_stats.spec_decoding_stats.num_accepted_tokens)
366
+ self.spec_decoding_prom.observe(
367
+ scheduler_stats.spec_decoding_stats)
368
369
if iteration_stats is None:
370
return
vllm/v1/spec_decode/metrics.py
+115
-13
@@ -1,9 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from dataclasses import dataclass
3
+from dataclasses import dataclass, field
4
+from typing import Optional
5
6
import numpy as np
7
+import prometheus_client
8
9
+from vllm.config import SpeculativeConfig
10
from vllm.logger import init_logger
11
12
logger = init_logger(__name__)
@@ -11,52 +14,151 @@ logger = init_logger(__name__)
14
15
@dataclass
16
class SpecDecodingStats:
17
+ """Per-step iteration decoding stats from scheduler.
18
+
19
+ Each scheduler step, statistics on spec decoding performance are
20
+ aggregated across requests by the scheduler and returned to the
21
+ frontend in EngineCoreOutputs->SchedulerStats.
22
+ """
23
+
24
+ num_spec_tokens: int
25
+ num_drafts: int = 0
26
num_draft_tokens: int = 0
27
num_accepted_tokens: int = 0
28
+ num_accepted_tokens_per_pos: list[int] = field(default_factory=list)
29
17
- def take(self):
18
- copied = SpecDecodingStats(self.num_draft_tokens,
19
- self.num_accepted_tokens)
20
- self.reset()
21
- return copied
22
-
23
- def reset(self):
24
- self.num_draft_tokens = 0
25
- self.num_accepted_tokens = 0
30
+ @classmethod
31
+ def new(cls, num_spec_tokens: int) -> "SpecDecodingStats":
32
+ return cls(num_spec_tokens=num_spec_tokens,
33
+ num_accepted_tokens_per_pos=[0] * num_spec_tokens)
34
27
- def observe(self, num_draft_tokens: int, num_accepted_tokens: int):
35
+ def observe_draft(self, num_draft_tokens: int, num_accepted_tokens: int):
36
+ self.num_drafts += 1
37
self.num_draft_tokens += num_draft_tokens
38
self.num_accepted_tokens += num_accepted_tokens
39
+ assert num_accepted_tokens <= self.num_spec_tokens
40
+ for i in range(num_accepted_tokens):
41
+ self.num_accepted_tokens_per_pos[i] += 1
42
+
43
44
+class SpecDecodingLogging:
45
+ """Aggregate and log spec decoding metrics.
46
32
-class SpecDecodingMetrics:
47
+ LoggingStatLogger aggregates per-iteration metrics over a set
48
+ time interval using observe() and then logs them using log()
49
+ before resetting to zero.
50
+ """
51
52
def __init__(self):
53
self.reset()
54
55
def reset(self):
56
+ self.num_drafts: list[int] = []
57
self.num_draft_tokens: list[int] = []
58
self.num_accepted_tokens: list[int] = []
59
+ self.accepted_tokens_per_pos_lists: list[list[int]] = []
60
61
def observe(self, spec_decoding_stats: SpecDecodingStats):
62
+ self.num_drafts.append(spec_decoding_stats.num_drafts)
63
self.num_draft_tokens.append(spec_decoding_stats.num_draft_tokens)
64
self.num_accepted_tokens.append(
65
spec_decoding_stats.num_accepted_tokens)
66
+ self.accepted_tokens_per_pos_lists.append(
67
+ spec_decoding_stats.num_accepted_tokens_per_pos)
68
69
def log(self, log_fn=logger.info):
70
+ num_drafts = np.sum(self.num_drafts)
71
num_draft_tokens = np.sum(self.num_draft_tokens)
72
num_accepted_tokens = np.sum(self.num_accepted_tokens)
73
74
draft_acceptance_rate = (num_accepted_tokens / num_draft_tokens *
75
100 if num_draft_tokens > 0 else float("nan"))
76
+ mean_acceptance_length = (num_accepted_tokens / num_drafts)
77
+
78
+ pos_matrix = np.array(self.accepted_tokens_per_pos_lists)
79
+ acceptance_rates = np.sum(pos_matrix, axis=0) / num_drafts
80
+ rates_str = ", ".join(f"{p:.3f}" for p in acceptance_rates)
81
82
log_fn(
83
"SpecDecoding metrics: "
84
"Draft acceptance rate: %.1f%%, "
85
+ "Mean acceptance length: %.2f, "
86
"Accepted: %d tokens, "
57
- "Drafted: %d tokens",
87
+ "Drafted: %d tokens, "
88
+ "Per-position acceptance rate: %s",
89
draft_acceptance_rate,
90
+ mean_acceptance_length,
91
num_accepted_tokens,
92
num_draft_tokens,
93
+ rates_str,
94
)
95
self.reset()
96
+
97
+
98
+class SpecDecodingProm:
99
+ """Record spec decoding metrics in Prometheus.
100
+
101
+ The acceptance rate can be calculated using a PromQL query:
102
+
103
+ rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) /
104
+ rate(vllm:spec_decode_num_draft_tokens_total[$interval])
105
+
106
+ The mean acceptance length can be calculated using:
107
+
108
+ rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) /
109
+ rate(vllm:spec_decode_num_drafts[$interval])
110
+
111
+ A per-position acceptance rate vector can be computed using
112
+
113
+ vllm:spec_decode_num_accepted_tokens_per_pos[$interval] /
114
+ vllm:spec_decode_num_drafts[$interval]
115
+ """
116
+
117
+ def __init__(self, speculative_config: Optional[SpeculativeConfig],
118
+ labelnames: list[str], labelvalues: list[str]):
119
+ self.spec_decoding_enabled = speculative_config is not None
120
+ if not self.spec_decoding_enabled:
121
+ return
122
+
123
+ self.counter_spec_decode_num_drafts = \
124
+ prometheus_client.Counter(
125
+ name="vllm:spec_decode_num_drafts_total",
126
+ documentation="Number of spec decoding drafts.",
127
+ labelnames=labelnames).labels(*labelvalues)
128
+ self.counter_spec_decode_num_draft_tokens = \
129
+ prometheus_client.Counter(
130
+ name="vllm:spec_decode_num_draft_tokens_total",
131
+ documentation="Number of draft tokens.",
132
+ labelnames=labelnames).labels(*labelvalues)
133
+ self.counter_spec_decode_num_accepted_tokens = \
134
+ prometheus_client.Counter(
135
+ name="vllm:spec_decode_num_accepted_tokens_total",
136
+ documentation="Number of accepted tokens.",
137
+ labelnames=labelnames).labels(*labelvalues)
138
+
139
+ assert speculative_config is not None
140
+ num_spec_tokens = (speculative_config.num_speculative_tokens
141
+ if self.spec_decoding_enabled else 0)
142
+ pos_labelnames = labelnames + ["position"]
143
+ base_counter = prometheus_client.Counter(
144
+ name="vllm:spec_decode_num_accepted_tokens_per_pos",
145
+ documentation="Accepted tokens per draft position.",
146
+ labelnames=pos_labelnames)
147
+ self.counter_spec_decode_num_accepted_tokens_per_pos: \
148
+ list[prometheus_client.Counter] = []
149
+ for pos in range(num_spec_tokens):
150
+ pos_labelvalues = labelvalues + [str(pos)]
151
+ self.counter_spec_decode_num_accepted_tokens_per_pos.append(
152
+ base_counter.labels(*pos_labelvalues))
153
+
154
+ def observe(self, spec_decoding_stats: SpecDecodingStats):
155
+ if not self.spec_decoding_enabled:
156
+ return
157
+ self.counter_spec_decode_num_drafts.inc(spec_decoding_stats.num_drafts)
158
+ self.counter_spec_decode_num_draft_tokens.inc(
159
+ spec_decoding_stats.num_draft_tokens)
160
+ self.counter_spec_decode_num_accepted_tokens.inc(
161
+ spec_decoding_stats.num_accepted_tokens)
162
+ for pos, counter in enumerate(
163
+ self.counter_spec_decode_num_accepted_tokens_per_pos):
164
+ counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos])