[WIP][V1][Metrics] Speculative decoding metrics

Fixes #13990, part of #10582

Omitting system efficiency for now.

Signed-off-by: Mark McLoughlin <markmc@redhat.com>

Author: markmc

vllm/v1/core/sched/scheduler.py

@@ -23,6 +23,7 @@
 from vllm.v1.metrics.stats import SchedulerStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
+from vllm.v1.spec_decode.metrics import SpecDecodingStats
 from vllm.v1.structured_output import StructuredOutputManager
 
 logger = init_logger(__name__)
@@ -565,6 +566,7 @@ def update_from_output(
         spec_token_ids = model_runner_output.spec_token_ids
         logprobs = model_runner_output.logprobs
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
+        spec_decoding_stats = SpecDecodingStats() if self.log_stats else None
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
 
         new_running: list[Request] = []
@@ -597,6 +599,15 @@ def update_from_output(
                                        len(generated_token_ids))
                 request.num_computed_tokens -= num_tokens_rejected
 
+                if spec_decoding_stats is not None:
+                    # FIXME: If a drafter proposes zero tokens, we should
+                    # treat this as if num_spec_tokens were proposed and
+                    # all rejected to allow fair comparisons between drafters
+                    spec_decoding_stats.observe(
+                        num_draft_tokens=len(scheduled_spec_token_ids),
+                        num_accepted_tokens=len(generated_token_ids) - 1,
+                        num_emitted_tokens=len(generated_token_ids))
+
             cached_encoder_input_ids = (
                 self.encoder_cache_manager.get_cached_input_ids(request))
             # OPTIMIZATION: Avoid list(set) if the set is empty.
@@ -672,7 +683,7 @@ def update_from_output(
         self.running = new_running
         engine_core_outputs = EngineCoreOutputs(
             outputs=outputs,
-            scheduler_stats=self.make_stats(),
+            scheduler_stats=self.make_stats(spec_decoding_stats),
         )
         if self.include_finished_set:
             #TODO currently sending duplicates here, improve this
@@ -739,12 +750,16 @@ def get_num_unscheduled_requests(self) -> int:
     def reset_prefix_cache(self) -> bool:
         return self.kv_cache_manager.reset_prefix_cache()
 
-    def make_stats(self) -> Optional[SchedulerStats]:
+    def make_stats(
+        self,
+        spec_decoding_stats: Optional[SpecDecodingStats] = None,
+    ) -> Optional[SchedulerStats]:
         if not self.log_stats:
             return None
         return SchedulerStats(
             num_running_reqs=len(self.running),
             num_waiting_reqs=len(self.waiting),
             gpu_cache_usage=self.kv_cache_manager.usage,
             prefix_cache_stats=self.kv_cache_manager.make_prefix_cache_stats(),
+            spec_decoding_stats=spec_decoding_stats,
         )

vllm/v1/engine/async_llm.py

@@ -74,7 +74,8 @@ def __init__(
             for i in range(vllm_config.parallel_config.data_parallel_size):
                 loggers: list[StatLoggerBase] = []
                 if logger.isEnabledFor(logging.INFO):
-                    loggers.append(LoggingStatLogger(engine_index=i))
+                    loggers.append(
+                        LoggingStatLogger(vllm_config, engine_index=i))
                 loggers.append(
                     PrometheusStatLogger(vllm_config, engine_index=i))
                 self.stat_loggers.append(loggers)

vllm/v1/metrics/loggers.py

@@ -12,6 +12,7 @@
 from vllm.v1.core.kv_cache_utils import PrefixCachingMetrics
 from vllm.v1.engine import FinishReason
 from vllm.v1.metrics.stats import IterationStats, SchedulerStats
+from vllm.v1.spec_decode.metrics import SpecDecodingMetrics
 
 logger = init_logger(__name__)
 
@@ -31,13 +32,15 @@ def log(self):  # noqa
 
 class LoggingStatLogger(StatLoggerBase):
 
-    def __init__(self, engine_index: int = 0):
+    def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
         self.engine_index = engine_index
         self._reset(time.monotonic())
         self.last_scheduler_stats = SchedulerStats()
         # Prefix cache metrics. This cannot be reset.
         # TODO: Make the interval configurable.
         self.prefix_caching_metrics = PrefixCachingMetrics()
+        self.spec_decoding_metrics = SpecDecodingMetrics(
+            vllm_config.speculative_config)
 
     def _reset(self, now):
         self.last_log_time = now
@@ -65,6 +68,10 @@ def record(self, scheduler_stats: SchedulerStats,
 
         self.prefix_caching_metrics.observe(scheduler_stats.prefix_cache_stats)
 
+        if scheduler_stats.spec_decoding_stats is not None:
+            self.spec_decoding_metrics.observe(
+                scheduler_stats.spec_decoding_stats)
+
         self.last_scheduler_stats = scheduler_stats
 
     def log(self):
@@ -94,6 +101,9 @@ def log(self):
             self.prefix_caching_metrics.hit_rate * 100,
         )
 
+        if scheduler_stats.spec_decoding_stats is not None:
+            self.spec_decoding_metrics.log()
+
 
 class PrometheusStatLogger(StatLoggerBase):
 
@@ -302,6 +312,29 @@ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
                         self.labelname_running_lora_adapters,
                     ])
 
+        #
+        # Speculative Decoding metrics
+        # The acceptance rate can be calculated using a PromQL query:
+        #
+        #   rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) /
+        #   rate(vllm:spec_decode_num_draft_tokens_total[$interval])
+        #
+        self.counter_spec_decode_num_draft_tokens = \
+            prometheus_client.Counter(
+                name="vllm:spec_decode_num_draft_tokens_total",
+                documentation="Number of draft tokens.",
+                labelnames=labelnames).labels(*labelvalues)
+        self.counter_spec_decode_num_accepted_tokens = \
+            prometheus_client.Counter(
+                name="vllm:spec_decode_num_accepted_tokens_total",
+                documentation="Number of accepted tokens.",
+                labelnames=labelnames).labels(*labelvalues)
+        self.counter_spec_decode_num_emitted_tokens = \
+            prometheus_client.Counter(
+                name="vllm:spec_decode_num_emitted_tokens_total",
+                documentation="Number of emitted tokens.",
+                labelnames=labelnames).labels(*labelvalues)
+
         #
         # Cache config info metric
         #
@@ -338,6 +371,14 @@ def record(self, scheduler_stats: SchedulerStats,
         self.counter_gpu_prefix_cache_hits.inc(
             scheduler_stats.prefix_cache_stats.hits)
 
+        if scheduler_stats.spec_decoding_stats is not None:
+            self.counter_spec_decode_num_draft_tokens.inc(
+                scheduler_stats.spec_decoding_stats.num_draft_tokens)
+            self.counter_spec_decode_num_accepted_tokens.inc(
+                scheduler_stats.spec_decoding_stats.num_accepted_tokens)
+            self.counter_spec_decode_num_emitted_tokens.inc(
+                scheduler_stats.spec_decoding_stats.num_emitted_tokens)
+
         if iteration_stats is None:
             return
 

vllm/v1/metrics/stats.py

@@ -4,6 +4,8 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Optional
 
+from vllm.v1.spec_decode.metrics import SpecDecodingStats
+
 if TYPE_CHECKING:
     from vllm.v1.engine import EngineCoreEvent, EngineCoreOutput, FinishReason
     from vllm.v1.engine.output_processor import RequestState
@@ -35,6 +37,8 @@ class SchedulerStats:
     prefix_cache_stats: PrefixCacheStats = field(
         default_factory=PrefixCacheStats)
 
+    spec_decoding_stats: Optional[SpecDecodingStats] = None
+
 
 @dataclass
 class LoRAStats:

vllm/v1/spec_decode/metrics.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from vllm.config import SpeculativeConfig
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class SpecDecodingStats:
+    num_draft_tokens: int = 0
+    num_accepted_tokens: int = 0
+    num_emitted_tokens: int = 0
+
+    def take(self):
+        copied = SpecDecodingStats(self.num_draft_tokens,
+                                   self.num_accepted_tokens,
+                                   self.num_emitted_tokens)
+        self.reset()
+        return copied
+
+    def reset(self):
+        self.num_draft_tokens = 0
+        self.num_accepted_tokens = 0
+        self.num_emitted_tokens = 0
+
+    def observe(self, num_draft_tokens: int, num_accepted_tokens: int,
+                num_emitted_tokens: int):
+        self.num_draft_tokens += num_draft_tokens
+        self.num_accepted_tokens += num_accepted_tokens
+        self.num_emitted_tokens += num_emitted_tokens
+
+
+class SpecDecodingMetrics:
+
+    def __init__(self, speculative_config: SpeculativeConfig):
+        self.num_spec_tokens = (speculative_config.num_speculative_tokens
+                                if speculative_config is not None else 0)
+        self.reset()
+
+    def reset(self):
+        self.num_draft_tokens: list[int] = []
+        self.num_accepted_tokens: list[int] = []
+        self.num_emitted_tokens: list[int] = []
+
+    def observe(self, spec_decoding_stats: SpecDecodingStats):
+        self.num_draft_tokens.append(spec_decoding_stats.num_draft_tokens)
+        self.num_accepted_tokens.append(
+            spec_decoding_stats.num_accepted_tokens)
+        self.num_emitted_tokens.append(spec_decoding_stats.num_emitted_tokens)
+
+    def log(self):
+        num_draft_tokens = np.sum(self.num_draft_tokens)
+        num_accepted_tokens = np.sum(self.num_accepted_tokens)
+        num_emitted_tokens = np.sum(self.num_emitted_tokens)
+
+        draft_acceptance_rate = (num_accepted_tokens / num_draft_tokens
+                                 if num_draft_tokens > 0 else float("nan"))
+
+        logger.info(
+            "Speculative metrics: "
+            "Draft acceptance rate: %.3f, "
+            "Number of speculative tokens: %d, "
+            "Number of accepted tokens: %d, "
+            "Number of draft tokens: %d, "
+            "Number of emitted tokens: %d.", draft_acceptance_rate,
+            num_accepted_tokens, num_draft_tokens, num_emitted_tokens)
+        self.reset()

vllm/v1/worker/gpu_model_runner.py

@@ -1151,20 +1151,20 @@ def generate_draft_token_ids(
         self,
         sampled_token_ids: list[list[int]],
         sampling_metadata: SamplingMetadata,
-    ) -> list[list[int]]:
+    ) -> list[Optional[list[int]]]:
         # TODO(woosuk): Optimize.
-        draft_token_ids: list[list[int]] = []
+        draft_token_ids: list[Optional[list[int]]] = []
         for i, sampled_ids in enumerate(sampled_token_ids):
             num_sampled_ids = len(sampled_ids)
             if not num_sampled_ids:
                 # Skip speculative decoding.
-                draft_token_ids.append([])
+                draft_token_ids.append(None)
                 continue
 
             # Skip requests that require top-p, top-k, etc.
             req_id = self.input_batch.req_ids[i]
             if not is_spec_decode_supported(req_id, self.input_batch):
-                draft_token_ids.append([])
+                draft_token_ids.append(None)
                 continue
 
             # Add sampled_token_ids to token_ids_cpu.