| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import time |
| 5 | from dataclasses import dataclass, field |
| 6 | |
| 7 | import numpy as np |
| 8 | import prometheus_client |
| 9 | |
| 10 | from vllm.config import SpeculativeConfig |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.v1.metrics.utils import create_metric_per_engine |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | |
| 17 | @dataclass |
| 18 | class SpecDecodingStats: |
| 19 | """Per-step iteration decoding stats from scheduler. |
| 20 | |
| 21 | Each scheduler step, statistics on spec decoding performance are |
| 22 | aggregated across requests by the scheduler and returned to the |
| 23 | frontend in EngineCoreOutputs->SchedulerStats. |
| 24 | """ |
| 25 | |
| 26 | num_spec_tokens: int |
| 27 | num_drafts: int = 0 |
| 28 | num_draft_tokens: int = 0 |
| 29 | num_accepted_tokens: int = 0 |
| 30 | num_accepted_tokens_per_pos: list[int] = field(default_factory=list) |
| 31 | num_draft_tokens_per_pos: list[int] = field(default_factory=list) |
| 32 | |
| 33 | @classmethod |
| 34 | def new(cls, num_spec_tokens: int) -> "SpecDecodingStats": |
| 35 | return cls( |
| 36 | num_spec_tokens=num_spec_tokens, |
| 37 | num_accepted_tokens_per_pos=[0] * num_spec_tokens, |
| 38 | num_draft_tokens_per_pos=[0] * num_spec_tokens, |
| 39 | ) |
| 40 | |
| 41 | def observe_draft(self, num_draft_tokens: int, num_accepted_tokens: int): |
| 42 | self.num_drafts += 1 |
| 43 | self.num_draft_tokens += num_draft_tokens |
| 44 | self.num_accepted_tokens += num_accepted_tokens |
| 45 | assert num_accepted_tokens <= self.num_spec_tokens |
| 46 | for i in range(num_accepted_tokens): |
| 47 | self.num_accepted_tokens_per_pos[i] += 1 |
| 48 | for i in range(num_draft_tokens): |
| 49 | self.num_draft_tokens_per_pos[i] += 1 |
| 50 | |
| 51 | |
| 52 | class SpecDecodingLogging: |
| 53 | """Aggregate and log spec decoding metrics. |
| 54 | |
| 55 | LoggingStatLogger aggregates per-iteration metrics over a set |
| 56 | time interval using observe() and then logs them using log() |
| 57 | before resetting to zero. |
| 58 | """ |
| 59 | |
| 60 | def __init__(self, is_diffusion: bool = False): |
| 61 | # Diffusion (dLLM) models reuse the spec-decode data path with |
| 62 | # overloaded semantics, so the raw spec-decode framing (drafts, bonus |
| 63 | # token, per-position vector) is logged with diffusion-native terms. |
| 64 | self.is_diffusion = is_diffusion |
| 65 | self.reset() |
| 66 | |
| 67 | def reset(self): |
| 68 | self.num_drafts: list[int] = [] |
| 69 | self.num_draft_tokens: list[int] = [] |
| 70 | self.num_accepted_tokens: list[int] = [] |
| 71 | self.accepted_tokens_per_pos_lists: list[list[int]] = [] |
| 72 | self.last_log_time = time.monotonic() |
| 73 | |
| 74 | def observe(self, spec_decoding_stats: SpecDecodingStats): |
| 75 | self.num_drafts.append(spec_decoding_stats.num_drafts) |
| 76 | self.num_draft_tokens.append(spec_decoding_stats.num_draft_tokens) |
| 77 | self.num_accepted_tokens.append(spec_decoding_stats.num_accepted_tokens) |
| 78 | self.accepted_tokens_per_pos_lists.append( |
| 79 | spec_decoding_stats.num_accepted_tokens_per_pos |
| 80 | ) |
| 81 | |
| 82 | def log(self, log_fn=logger.info): |
| 83 | if not self.num_drafts: |
| 84 | return |
| 85 | num_drafts = np.sum(self.num_drafts) |
| 86 | num_draft_tokens = np.sum(self.num_draft_tokens) |
| 87 | num_accepted_tokens = np.sum(self.num_accepted_tokens) |
| 88 | draft_throughput = 0 |
| 89 | accepted_throughput = 0 |
| 90 | |
| 91 | elapsed_time = time.monotonic() - self.last_log_time |
| 92 | if elapsed_time > 0: |
| 93 | draft_throughput = num_draft_tokens / elapsed_time |
| 94 | accepted_throughput = num_accepted_tokens / elapsed_time |
| 95 | |
| 96 | if self.is_diffusion: |
| 97 | self._log_diffusion( |
| 98 | log_fn, |
| 99 | num_denoising_steps=num_drafts, |
| 100 | num_canvas_tokens=num_draft_tokens, |
| 101 | num_committed_tokens=num_accepted_tokens, |
| 102 | committed_throughput=accepted_throughput, |
| 103 | ) |
| 104 | self.reset() |
| 105 | return |
| 106 | |
| 107 | draft_acceptance_rate = ( |
| 108 | num_accepted_tokens / num_draft_tokens * 100 |
| 109 | if num_draft_tokens > 0 |
| 110 | else float("nan") |
| 111 | ) |
| 112 | |
| 113 | # Conventionally, mean acceptance length includes the bonus token |
| 114 | mean_acceptance_length = 1 + (num_accepted_tokens / num_drafts) |
| 115 | |
| 116 | pos_matrix = np.array(self.accepted_tokens_per_pos_lists) |
| 117 | acceptance_rates = np.sum(pos_matrix, axis=0) / num_drafts |
| 118 | rates_str = ", ".join(f"{p:.3f}" for p in acceptance_rates) |
| 119 | |
| 120 | log_fn( |
| 121 | "SpecDecoding metrics: " |
| 122 | "Mean acceptance length: %.2f, " |
| 123 | "Accepted throughput: %.2f tokens/s, " |
| 124 | "Drafted throughput: %.2f tokens/s, " |
| 125 | "Accepted: %d tokens, " |
| 126 | "Drafted: %d tokens, " |
| 127 | "Per-position acceptance rate: %s, " |
| 128 | "Avg Draft acceptance rate: %.1f%%", |
| 129 | mean_acceptance_length, |
| 130 | accepted_throughput, |
| 131 | draft_throughput, |
| 132 | num_accepted_tokens, |
| 133 | num_draft_tokens, |
| 134 | rates_str, |
| 135 | draft_acceptance_rate, |
| 136 | ) |
| 137 | self.reset() |
| 138 | |
| 139 | def _log_diffusion( |
| 140 | self, |
| 141 | log_fn, |
| 142 | num_denoising_steps: int, |
| 143 | num_canvas_tokens: int, |
| 144 | num_committed_tokens: int, |
| 145 | committed_throughput: float, |
| 146 | ): |
| 147 | # Each "draft" is one denoising step that re-evaluates the canvas block |
| 148 | # and finalizes some of its positions. |
| 149 | mean_committed_per_step = ( |
| 150 | num_committed_tokens / num_denoising_steps |
| 151 | if num_denoising_steps > 0 |
| 152 | else float("nan") |
| 153 | ) |
| 154 | mean_steps_per_canvas = ( |
| 155 | num_canvas_tokens / num_committed_tokens |
| 156 | if num_committed_tokens > 0 |
| 157 | else float("nan") |
| 158 | ) |
| 159 | |
| 160 | log_fn( |
| 161 | "DiffusionDecoding metrics: " |
| 162 | "Committed token throughput: %.2f tokens/s, " |
| 163 | "Mean denoising steps per canvas: %.2f, " |
| 164 | "Mean tokens committed per denoising step: %.2f, " |
| 165 | "Committed: %d tokens, " |
| 166 | "Denoising steps: %d, " |
| 167 | "Canvas positions evaluated: %d", |
| 168 | committed_throughput, |
| 169 | mean_steps_per_canvas, |
| 170 | mean_committed_per_step, |
| 171 | num_committed_tokens, |
| 172 | num_denoising_steps, |
| 173 | num_canvas_tokens, |
| 174 | ) |
| 175 | |
| 176 | |
| 177 | class SpecDecodingProm: |
| 178 | """Record spec decoding metrics in Prometheus. |
| 179 | |
| 180 | The acceptance rate can be calculated using a PromQL query: |
| 181 | |
| 182 | rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) / |
| 183 | rate(vllm:spec_decode_num_draft_tokens_total[$interval]) |
| 184 | |
| 185 | The mean acceptance length (conventionally including bonus tokens) |
| 186 | can be calculated using: |
| 187 | |
| 188 | 1 + ( |
| 189 | rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) / |
| 190 | rate(vllm:spec_decode_num_drafts_total[$interval])) |
| 191 | |
| 192 | A per-position acceptance rate vector can be computed using |
| 193 | |
| 194 | vllm:spec_decode_num_accepted_tokens_per_pos_total[$interval] / |
| 195 | vllm:spec_decode_num_drafts_total[$interval] |
| 196 | """ |
| 197 | |
| 198 | _counter_cls = prometheus_client.Counter |
| 199 | |
| 200 | def __init__( |
| 201 | self, |
| 202 | speculative_config: SpeculativeConfig | None, |
| 203 | labelnames: list[str], |
| 204 | per_engine_labelvalues: dict[int, list[object]], |
| 205 | is_diffusion: bool = False, |
| 206 | ): |
| 207 | # Diffusion (dLLM) models reuse the spec-decode counters but expose them |
| 208 | # under diffusion-native names; the per-position acceptance vector does |
| 209 | # not apply, so it is omitted. |
| 210 | self.is_diffusion = is_diffusion |
| 211 | self.spec_decoding_enabled = speculative_config is not None or is_diffusion |
| 212 | if not self.spec_decoding_enabled: |
| 213 | return |
| 214 | |
| 215 | if is_diffusion: |
| 216 | counter_specs = [ |
| 217 | ("vllm:diffusion_num_denoising_steps", "Number of denoising steps."), |
| 218 | ( |
| 219 | "vllm:diffusion_num_canvas_positions", |
| 220 | "Number of canvas positions evaluated.", |
| 221 | ), |
| 222 | ( |
| 223 | "vllm:diffusion_num_committed_tokens", |
| 224 | "Number of committed (finalized) tokens.", |
| 225 | ), |
| 226 | ] |
| 227 | else: |
| 228 | counter_specs = [ |
| 229 | ("vllm:spec_decode_num_drafts", "Number of spec decoding drafts."), |
| 230 | ("vllm:spec_decode_num_draft_tokens", "Number of draft tokens."), |
| 231 | ("vllm:spec_decode_num_accepted_tokens", "Number of accepted tokens."), |
| 232 | ] |
| 233 | |
| 234 | counters = [ |
| 235 | create_metric_per_engine( |
| 236 | self._counter_cls(name=name, documentation=doc, labelnames=labelnames), |
| 237 | per_engine_labelvalues, |
| 238 | ) |
| 239 | for name, doc in counter_specs |
| 240 | ] |
| 241 | # num_drafts/num_draft_tokens/num_accepted_tokens map onto denoising |
| 242 | # steps/canvas positions/committed tokens in the diffusion path. |
| 243 | self.counter_spec_decode_num_drafts = counters[0] |
| 244 | self.counter_spec_decode_num_draft_tokens = counters[1] |
| 245 | self.counter_spec_decode_num_accepted_tokens = counters[2] |
| 246 | |
| 247 | self.counter_spec_decode_num_accepted_tokens_per_pos: dict[ |
| 248 | int, list[prometheus_client.Counter] |
| 249 | ] = {} |
| 250 | if not is_diffusion: |
| 251 | assert speculative_config is not None |
| 252 | num_spec_tokens = speculative_config.num_speculative_tokens |
| 253 | pos_labelnames = labelnames + ["position"] |
| 254 | base_counter = self._counter_cls( |
| 255 | name="vllm:spec_decode_num_accepted_tokens_per_pos", |
| 256 | documentation="Accepted tokens per draft position.", |
| 257 | labelnames=pos_labelnames, |
| 258 | ) |
| 259 | self.counter_spec_decode_num_accepted_tokens_per_pos = { |
| 260 | idx: [ |
| 261 | base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens) |
| 262 | ] |
| 263 | for idx, lv in per_engine_labelvalues.items() |
| 264 | } |
| 265 | |
| 266 | def observe(self, spec_decoding_stats: SpecDecodingStats, engine_idx: int = 0): |
| 267 | if not self.spec_decoding_enabled: |
| 268 | return |
| 269 | self.counter_spec_decode_num_drafts[engine_idx].inc( |
| 270 | spec_decoding_stats.num_drafts |
| 271 | ) |
| 272 | self.counter_spec_decode_num_draft_tokens[engine_idx].inc( |
| 273 | spec_decoding_stats.num_draft_tokens |
| 274 | ) |
| 275 | self.counter_spec_decode_num_accepted_tokens[engine_idx].inc( |
| 276 | spec_decoding_stats.num_accepted_tokens |
| 277 | ) |
| 278 | for pos, counter in enumerate( |
| 279 | self.counter_spec_decode_num_accepted_tokens_per_pos.get(engine_idx, []) |
| 280 | ): |
| 281 | counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos]) |