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

Fixes vllm-project#13990, part of vllm-project#10582

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

Author: markmc

vllm/v1/core/sched/scheduler.py

@@ -22,6 +22,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__)
@@ -535,6 +536,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()
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
 
         new_running: list[Request] = []
@@ -559,6 +561,7 @@ def update_from_output(
                 # Otherwise, we ignore the sampler output for the request.
                 request.num_computed_tokens += num_tokens_scheduled
                 assert request.num_computed_tokens <= request.num_tokens
+                spec_decoding_stats.num_emitted_tokens += num_tokens_scheduled
             else:
                 # num_computed_tokens_step represents the number of tokens
                 # processed in the current step, considering scheduled
@@ -576,6 +579,13 @@ def update_from_output(
                     len(generated_token_ids))
                 request.num_computed_tokens += num_computed_tokens_step
 
+                spec_decoding_stats.num_draft_tokens += len(
+                    scheduled_spec_token_ids)
+                spec_decoding_stats.num_accepted_tokens += len(
+                    generated_token_ids) - 1
+                spec_decoding_stats.num_emitted_tokens += \
+                    num_computed_tokens_step
+
             cached_encoder_input_ids = (
                 self.encoder_cache_manager.get_cached_input_ids(request))
             # OPTIMIZATION: Avoid list(set) if the set is empty.
@@ -647,7 +657,7 @@ def update_from_output(
         self.running = new_running
         return EngineCoreOutputs(
             outputs=outputs,
-            scheduler_stats=self.make_stats(),
+            scheduler_stats=self.make_stats(spec_decoding_stats),
         )
 
     def add_request(self, request: Request) -> None:
@@ -708,12 +718,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

@@ -68,7 +68,7 @@ def __init__(
         self.stat_loggers: list[StatLoggerBase] = []
         if self.log_stats:
             if logger.isEnabledFor(logging.INFO):
-                self.stat_loggers.append(LoggingStatLogger())
+                self.stat_loggers.append(LoggingStatLogger(vllm_config))
             self.stat_loggers.append(PrometheusStatLogger(vllm_config))
 
         # Tokenizer (+ ensure liveness if running in another process).

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,12 +32,14 @@ def log(self):  # noqa
 
 class LoggingStatLogger(StatLoggerBase):
 
-    def __init__(self):
+    def __init__(self, vllm_config: VllmConfig):
         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
@@ -64,6 +67,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):
@@ -91,6 +98,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):
 
@@ -296,6 +306,26 @@ def __init__(self, vllm_config: VllmConfig):
                         self.labelname_running_lora_adapters,
                     ])
 
+        #
+        # Speculative Decoding metrics
+        # FIXME: add note on acceptance rate and system efficiency
+        #
+        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
         #
@@ -332,6 +362,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,96 @@
+# 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
+
+
+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)
+        # FIXME: relies on num_draft_tokens % k == 0 assumption
+        #max_num_emitted_tokens = get_max_num_emitted_tokens(
+        #    draft_tokens=num_draft_tokens, k=self.num_spec_tokens)
+        draft_acceptance_rate = (num_accepted_tokens / num_draft_tokens
+                                 if num_draft_tokens > 0 else float("nan"))
+        #system_efficiency = (num_emitted_tokens / max_num_emitted_tokens
+        #                     if max_num_emitted_tokens > 0 else float("nan"))
+        system_efficiency = float("nan")
+        logger.info(
+            "Speculative metrics: "
+            "Draft acceptance rate: %.3f, "
+            "System efficiency: %.3f, "
+            "Number of speculative tokens: %d, "
+            "Number of accepted tokens: %d, "
+            "Number of draft tokens: %d, "
+            "Number of emitted tokens: %d.", draft_acceptance_rate,
+            system_efficiency, self.num_spec_tokens, num_accepted_tokens,
+            num_draft_tokens, num_emitted_tokens)
+        self.reset()
+
+
+def get_max_num_emitted_tokens(draft_tokens: int, k: int) -> int:
+    """Calculate the number of emitted tokens, assuming all tokens accepted.
+
+    This is equal to the number of sequences that have been speculated on,
+    times (speculation len + 1). The +1 comes from the bonus token.
+    """
+    # Determine the number of sequences that have been speculated on. Since
+    # the batch size can be variable, we divide by k.
+    print(f"DRAFT TOKENS {draft_tokens} K {k}")
+    # Cannot assume this - ngram proposer says "If there are less than k
+    # tokens follow the match, we will return the maximum amount of tokens
+    # until the end."
+    assert draft_tokens % k == 0
+    total_num_spec_seqs = draft_tokens // k
+
+    # A single sequence may emit k accepted tokens and one bonus token in
+    # the best case.
+    num_emitted_per_seq_if_all_accepted = k + 1
+
+    # The max num of emitted tokens is the number of speculated sequences
+    # times the max emitted per seq.
+    return total_num_spec_seqs * num_emitted_per_seq_if_all_accepted