[V1][Spec Decoding] Strip metrics back to acceptance rate

Now just num_accepted_tokens, num_draft_tokens, and acceptance rate.

Signed-off-by: Mark McLoughlin <[email protected]>

Author: markmc

vllm/v1/core/sched/scheduler.py

@@ -602,13 +602,9 @@ def update_from_output(
                 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))
+                        num_accepted_tokens=len(generated_token_ids) - 1)
 
             cached_encoder_input_ids = (
                 self.encoder_cache_manager.get_cached_input_ids(request))

vllm/v1/engine/async_llm.py

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

vllm/v1/metrics/loggers.py

@@ -32,15 +32,14 @@ def log(self):  # noqa
 
 class LoggingStatLogger(StatLoggerBase):
 
-    def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
+    def __init__(self, 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)
+        self.spec_decoding_metrics = SpecDecodingMetrics()
 
     def _reset(self, now):
         self.last_log_time = now
@@ -329,11 +328,6 @@ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
                 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
@@ -376,8 +370,6 @@ def record(self, scheduler_stats: SchedulerStats,
                 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/spec_decode/metrics.py

@@ -4,7 +4,6 @@
 
 import numpy as np
 
-from vllm.config import SpeculativeConfig
 from vllm.logger import init_logger
 
 logger = init_logger(__name__)
@@ -14,59 +13,47 @@
 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.num_accepted_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):
+    def observe(self, num_draft_tokens: int, num_accepted_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)
+    def __init__(self):
         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)
+            "Number of draft tokens: %d, ", draft_acceptance_rate,
+            num_accepted_tokens, num_draft_tokens)
         self.reset()

vllm/v1/worker/gpu_model_runner.py

@@ -1154,20 +1154,20 @@ def generate_draft_token_ids(
         self,
         sampled_token_ids: list[list[int]],
         sampling_metadata: SamplingMetadata,
-    ) -> list[Optional[list[int]]]:
+    ) -> list[list[int]]:
         # TODO(woosuk): Optimize.
-        draft_token_ids: list[Optional[list[int]]] = []
+        draft_token_ids: list[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(None)
+                draft_token_ids.append([])
                 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(None)
+                draft_token_ids.append([])
                 continue
 
             # Add sampled_token_ids to token_ids_cpu.