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 14, 05:30 UTC
[Spec Decoding]Support Spec Decoding Metrics in DP Mode (#24049)
Signed-off-by: wuhang <wuhang6@huawei.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
wuhang committed
Sep 15, 2025 at 05:11 UTC
90f3f7d73e09aff12dc96767223fa54f48f117a8
2 files changed
+54
-38
vllm/v1/metrics/loggers.py
+6
-10
@@ -169,15 +169,11 @@ class PrometheusStatLogger(StatLoggerBase):
169
model_name = vllm_config.model_config.served_model_name
170
max_model_len = vllm_config.model_config.max_model_len
171
172
- if (len(self.engine_indexes) > 1
173
- and vllm_config.speculative_config is not None):
174
- raise NotImplementedError("Prometheus metrics with Spec Decoding "
175
- "with >1 EngineCore per AsyncLLM is not "
176
- "supported yet.")
177
- spec_decode_labelvalues = [
178
- vllm_config.model_config.served_model_name,
179
- str(self.engine_indexes[0])
180
- ]
172
+ spec_decode_labelvalues: dict[int, list[str]] = {
173
+ idx: [model_name, str(idx)]
174
+ for idx in engine_indexes
175
+ }
176
+
177
self.spec_decoding_prom = self._spec_decoding_cls(
178
vllm_config.speculative_config, labelnames,
179
spec_decode_labelvalues)
@@ -530,7 +526,7 @@ class PrometheusStatLogger(StatLoggerBase):
526
527
if scheduler_stats.spec_decoding_stats is not None:
528
self.spec_decoding_prom.observe(
533
- scheduler_stats.spec_decoding_stats)
529
+ scheduler_stats.spec_decoding_stats, engine_idx)
530
531
if iteration_stats is None:
532
return
vllm/v1/spec_decode/metrics.py
+48
-28
@@ -140,27 +140,32 @@ class SpecDecodingProm:
140
self,
141
speculative_config: Optional[SpeculativeConfig],
142
labelnames: list[str],
143
- labelvalues: list[str],
143
+ per_engine_labelvalues: dict[int, list[str]],
144
):
145
self.spec_decoding_enabled = speculative_config is not None
146
if not self.spec_decoding_enabled:
147
return
148
149
- self.counter_spec_decode_num_drafts = \
150
- self._counter_cls(
151
- name="vllm:spec_decode_num_drafts",
152
- documentation="Number of spec decoding drafts.",
153
- labelnames=labelnames).labels(*labelvalues)
154
- self.counter_spec_decode_num_draft_tokens = \
155
- self._counter_cls(
156
- name="vllm:spec_decode_num_draft_tokens",
157
- documentation="Number of draft tokens.",
158
- labelnames=labelnames,).labels(*labelvalues)
159
- self.counter_spec_decode_num_accepted_tokens = \
160
- self._counter_cls(
161
- name="vllm:spec_decode_num_accepted_tokens",
162
- documentation="Number of accepted tokens.",
163
- labelnames=labelnames).labels(*labelvalues)
149
+ counter_drafts = self._counter_cls(
150
+ name="vllm:spec_decode_num_drafts",
151
+ documentation="Number of spec decoding drafts.",
152
+ labelnames=labelnames)
153
+ self.counter_spec_decode_num_drafts = make_per_engine(
154
+ counter_drafts, per_engine_labelvalues)
155
+
156
+ counter_draft_tokens = self._counter_cls(
157
+ name="vllm:spec_decode_num_draft_tokens",
158
+ documentation="Number of draft tokens.",
159
+ labelnames=labelnames)
160
+ self.counter_spec_decode_num_draft_tokens = make_per_engine(
161
+ counter_draft_tokens, per_engine_labelvalues)
162
+
163
+ counter_accepted_tokens = self._counter_cls(
164
+ name="vllm:spec_decode_num_accepted_tokens",
165
+ documentation="Number of accepted tokens.",
166
+ labelnames=labelnames)
167
+ self.counter_spec_decode_num_accepted_tokens = make_per_engine(
168
+ counter_accepted_tokens, per_engine_labelvalues)
169
170
assert speculative_config is not None
171
num_spec_tokens = (speculative_config.num_speculative_tokens
@@ -171,21 +176,36 @@ class SpecDecodingProm:
176
documentation="Accepted tokens per draft position.",
177
labelnames=pos_labelnames,
178
)
174
- self.counter_spec_decode_num_accepted_tokens_per_pos: list[
175
- prometheus_client.Counter] = []
176
- for pos in range(num_spec_tokens):
177
- pos_labelvalues = labelvalues + [str(pos)]
178
- self.counter_spec_decode_num_accepted_tokens_per_pos.append(
179
- base_counter.labels(*pos_labelvalues))
180
-
181
- def observe(self, spec_decoding_stats: SpecDecodingStats):
179
+ self.counter_spec_decode_num_accepted_tokens_per_pos: dict[
180
+ int, list[prometheus_client.Counter]] = {
181
+ idx: [
182
+ base_counter.labels(*lv, str(pos))
183
+ for pos in range(num_spec_tokens)
184
+ ]
185
+ for idx, lv in per_engine_labelvalues.items()
186
+ }
187
+
188
+ def observe(self,
189
+ spec_decoding_stats: SpecDecodingStats,
190
+ engine_idx: int = 0):
191
if not self.spec_decoding_enabled:
192
return
184
- self.counter_spec_decode_num_drafts.inc(spec_decoding_stats.num_drafts)
185
- self.counter_spec_decode_num_draft_tokens.inc(
193
+ self.counter_spec_decode_num_drafts[engine_idx].inc(
194
+ spec_decoding_stats.num_drafts)
195
+ self.counter_spec_decode_num_draft_tokens[engine_idx].inc(
196
spec_decoding_stats.num_draft_tokens)
187
- self.counter_spec_decode_num_accepted_tokens.inc(
197
+ self.counter_spec_decode_num_accepted_tokens[engine_idx].inc(
198
spec_decoding_stats.num_accepted_tokens)
199
for pos, counter in enumerate(
190
- self.counter_spec_decode_num_accepted_tokens_per_pos):
200
+ self.
201
+ counter_spec_decode_num_accepted_tokens_per_pos[engine_idx]):
202
counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos])
203
+
204
+
205
+def make_per_engine(counter: prometheus_client.Counter,
206
+ per_engine_labelvalues: dict[int, list[str]]):
207
+ """Create a counter for each label value."""
208
+ return {
209
+ idx: counter.labels(*labelvalues)
210
+ for idx, labelvalues in per_engine_labelvalues.items()
211
+ }