fix(metrics): Accurately track emitted tokens for speculative decoding

tests/entrypoints/openai/completion/test_completion_error.py

@@ -196,8 +196,8 @@ def test_completion_per_request_metrics_suppressed_for_multiple_prompts():
 def _spec_decode_metrics() -> RequestSpecDecodeMetrics:
     # Two verify steps: accept 3 drafts, then 1 -> histogram [0, 1, 0, 1].
     m = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
-    m.observe(num_draft_tokens=3, num_accepted=3)
-    m.observe(num_draft_tokens=3, num_accepted=1)
+    m.observe(num_draft_tokens=3, num_accepted=3, num_emitted=4)
+    m.observe(num_draft_tokens=3, num_accepted=1, num_emitted=2)
     return m
 
 

tests/v1/spec_decode/test_request_acceptance.py

@@ -18,7 +18,9 @@
 def _metrics(pairs, num_spec_tokens=3, detailed=False):
     s = RequestSpecDecodeMetrics.new(num_spec_tokens)
     for k, j in pairs:
-        s.observe(num_draft_tokens=k, num_accepted=j, detailed=detailed)
+        s.observe(
+            num_draft_tokens=k, num_accepted=j, num_emitted=j + 1, detailed=detailed
+        )
     return s
 
 
@@ -60,6 +62,7 @@ def test_to_dict_summary_omits_per_step_arrays():
         "num_accepted_draft_tokens": 9,
         "num_draft_tokens": 15,
         "num_spec_tokens": 3,
+        "num_emitted_tokens": 14,
     }
 
 
@@ -91,6 +94,7 @@ def test_empty_metrics_do_not_divide_by_zero():
     d = RequestSpecDecodeMetrics.new(3).to_dict()
     assert d["num_spec_steps"] == 0
     assert d["num_draft_tokens"] == 0
+    assert d["num_emitted_tokens"] == 0
     assert d["draft_acceptance_rate"] == 0.0
     assert d["mean_acceptance_length"] == 1.0
     assert "per_step_accepted" not in d
@@ -102,11 +106,12 @@ def test_observe_records_proposed_and_accepted_independently():
     # grammar-invalidated-draft subtraction happens in the scheduler before
     # observe() -- see test_per_request_spec_decode_subtracts_invalid_drafts.)
     s = RequestSpecDecodeMetrics.new(num_spec_tokens=3)
-    s.observe(num_draft_tokens=2, num_accepted=1)
-    s.observe(num_draft_tokens=3, num_accepted=1)
+    s.observe(num_draft_tokens=2, num_accepted=1, num_emitted=2)
+    s.observe(num_draft_tokens=3, num_accepted=1, num_emitted=2)
     d = s.to_dict()
     assert d["acceptance_histogram"] == [0, 2, 0, 0]  # both steps accepted 1
     assert d["num_draft_tokens"] == 5  # proposed summed independently: 2 + 3
+    assert d["num_emitted_tokens"] == 4
 
 
 def test_engine_core_output_round_trips_spec_decode_metrics():

vllm/v1/core/sched/scheduler.py

@@ -1963,6 +1963,7 @@ def update_from_output(
                 sampled_token_ids[req_index] if sampled_token_ids else []
             )
 
+            spec_observe_args = None
             scheduled_spec_token_ids = (
                 scheduler_output.scheduled_spec_decode_tokens.get(req_id)
             )
@@ -1982,27 +1983,13 @@ def update_from_output(
                         request.num_computed_tokens -= num_rejected
                     if request.num_output_placeholders > 0:
                         request.num_output_placeholders -= num_rejected
-                spec_decoding_stats = self.make_spec_decoding_stats(
-                    spec_decoding_stats,
-                    num_draft_tokens=num_draft_tokens,
-                    num_accepted_tokens=num_accepted,
-                    num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
-                    request_id=req_id,
-                )
-                if request.spec_decode_metrics is not None:
-                    # Exclude grammar-invalidated drafts from the proposed
-                    # count, mirroring make_spec_decoding_stats; the accepted
-                    # bucket (j) is unaffected.
-                    adj_draft_tokens = num_draft_tokens
-                    if scheduler_output.num_invalid_spec_tokens:
-                        adj_draft_tokens -= (
-                            scheduler_output.num_invalid_spec_tokens.get(req_id, 0)
-                        )
-                    request.spec_decode_metrics.observe(
-                        num_draft_tokens=adj_draft_tokens,
-                        num_accepted=num_accepted,
-                        detailed=self.spec_decode_metrics_level == "detailed",
+                # Exclude grammar-invalidated drafts from the proposed count.
+                adj_draft_tokens = num_draft_tokens
+                if scheduler_output.num_invalid_spec_tokens:
+                    adj_draft_tokens -= scheduler_output.num_invalid_spec_tokens.get(
+                        req_id, 0
                     )
+                spec_observe_args = (adj_draft_tokens, num_accepted)
 
             # Free encoder inputs only after the step has actually executed.
             if request.has_encoder_inputs:
@@ -2040,6 +2027,22 @@ def update_from_output(
                 request.status = RequestStatus.FINISHED_STOPPED
                 stopped = True
 
+            if spec_observe_args is not None:
+                num_emitted = len(new_token_ids) if new_token_ids else 0
+                if request.spec_decode_metrics is not None:
+                    request.spec_decode_metrics.observe(
+                        num_draft_tokens=spec_observe_args[0],
+                        num_accepted=spec_observe_args[1],
+                        num_emitted=num_emitted,
+                        detailed=self.spec_decode_metrics_level == "detailed",
+                    )
+                spec_decoding_stats = self.make_spec_decoding_stats(
+                    spec_decoding_stats,
+                    num_draft_tokens=spec_observe_args[0],
+                    num_accepted_tokens=spec_observe_args[1],
+                    num_emitted_tokens=num_emitted,
+                )
+
             if new_token_ids and self.structured_output_manager.should_advance(
                 request, new_token_ids=new_token_ids
             ):
@@ -2802,17 +2805,16 @@ def make_spec_decoding_stats(
         spec_decoding_stats: SpecDecodingStats | None,
         num_draft_tokens: int,
         num_accepted_tokens: int,
-        num_invalid_spec_tokens: dict[str, int] | None,
-        request_id: str,
+        num_emitted_tokens: int,
     ) -> SpecDecodingStats | None:
         if not self.log_stats or not num_draft_tokens:
             return None
         if spec_decoding_stats is None:
             spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens)
-        if num_invalid_spec_tokens:
-            num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0)
         spec_decoding_stats.observe_draft(
-            num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens
+            num_draft_tokens=num_draft_tokens,
+            num_accepted_tokens=num_accepted_tokens,
+            num_emitted_tokens=num_emitted_tokens,
         )
         return spec_decoding_stats
 

vllm/v1/metrics/stats.py

@@ -328,6 +328,7 @@ class RequestSpecDecodeMetrics:
     num_spec_tokens: int
     histogram: list[int] = field(default_factory=list)
     num_draft_tokens: int = 0
+    num_emitted_tokens: int = 0
     per_step_accepted: list[int] = field(default_factory=list)
     per_step_drafted: list[int] = field(default_factory=list)
 
@@ -339,10 +340,15 @@ def new(cls, num_spec_tokens: int) -> "RequestSpecDecodeMetrics":
         )
 
     def observe(
-        self, num_draft_tokens: int, num_accepted: int, detailed: bool = False
+        self,
+        num_draft_tokens: int,
+        num_accepted: int,
+        num_emitted: int,
+        detailed: bool = False,
     ) -> None:
         self.histogram[num_accepted] += 1
         self.num_draft_tokens += num_draft_tokens
+        self.num_emitted_tokens += num_emitted
         if detailed:
             self.per_step_accepted.append(num_accepted)
             self.per_step_drafted.append(num_draft_tokens)
@@ -352,13 +358,13 @@ def to_dict(self) -> dict[str, Any]:
 
         ``acceptance_histogram`` is a dense list indexed by accepted draft count
         ``j`` (length ``num_spec_tokens + 1``). ``mean_acceptance_length``
-        includes the bonus token (``j + 1``); ``draft_acceptance_rate`` is
-        draft-only, full precision. Per-step arrays are included only when
-        populated (``detailed`` level).
+        is the average number of emitted tokens per step (including the bonus token);
+        ``draft_acceptance_rate`` is draft-only, full precision. Per-step arrays
+        are included only when populated (``detailed`` level).
         """
         num_spec_steps = sum(self.histogram)
         num_accepted = sum(j * count for j, count in enumerate(self.histogram))
-        mean_al = 1.0 + num_accepted / num_spec_steps if num_spec_steps else 1.0
+        mean_al = self.num_emitted_tokens / num_spec_steps if num_spec_steps else 1.0
         rate = num_accepted / self.num_draft_tokens if self.num_draft_tokens else 0.0
         result: dict[str, Any] = {
             "mean_acceptance_length": mean_al,
@@ -367,6 +373,7 @@ def to_dict(self) -> dict[str, Any]:
             "num_spec_steps": num_spec_steps,
             "num_accepted_draft_tokens": num_accepted,
             "num_draft_tokens": self.num_draft_tokens,
+            "num_emitted_tokens": self.num_emitted_tokens,
             "num_spec_tokens": self.num_spec_tokens,
         }
         if self.per_step_accepted:

vllm/v1/spec_decode/metrics.py

@@ -27,6 +27,7 @@ class SpecDecodingStats:
     num_drafts: int = 0
     num_draft_tokens: int = 0
     num_accepted_tokens: int = 0
+    num_emitted_tokens: int = 0
     num_accepted_tokens_per_pos: list[int] = field(default_factory=list)
     num_draft_tokens_per_pos: list[int] = field(default_factory=list)
 
@@ -38,10 +39,13 @@ def new(cls, num_spec_tokens: int) -> "SpecDecodingStats":
             num_draft_tokens_per_pos=[0] * num_spec_tokens,
         )
 
-    def observe_draft(self, num_draft_tokens: int, num_accepted_tokens: int):
+    def observe_draft(
+        self, num_draft_tokens: int, num_accepted_tokens: int, num_emitted_tokens: int
+    ):
         self.num_drafts += 1
         self.num_draft_tokens += num_draft_tokens
         self.num_accepted_tokens += num_accepted_tokens
+        self.num_emitted_tokens += num_emitted_tokens
         assert num_accepted_tokens <= self.num_spec_tokens
         for i in range(num_accepted_tokens):
             self.num_accepted_tokens_per_pos[i] += 1
@@ -68,13 +72,15 @@ def reset(self):
         self.num_drafts: list[int] = []
         self.num_draft_tokens: list[int] = []
         self.num_accepted_tokens: list[int] = []
+        self.num_emitted_tokens: list[int] = []
         self.accepted_tokens_per_pos_lists: list[list[int]] = []
         self.last_log_time = time.monotonic()
 
     def observe(self, spec_decoding_stats: SpecDecodingStats):
         self.num_drafts.append(spec_decoding_stats.num_drafts)
         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)
         self.accepted_tokens_per_pos_lists.append(
             spec_decoding_stats.num_accepted_tokens_per_pos
         )
@@ -85,6 +91,7 @@ def log(self, log_fn=logger.info):
         num_drafts = np.sum(self.num_drafts)
         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_throughput = 0
         accepted_throughput = 0
 
@@ -110,8 +117,9 @@ def log(self, log_fn=logger.info):
             else float("nan")
         )
 
-        # Conventionally, mean acceptance length includes the bonus token
-        mean_acceptance_length = 1 + (num_accepted_tokens / num_drafts)
+        mean_acceptance_length = (
+            num_emitted_tokens / num_drafts if num_drafts > 0 else 1.0
+        )
 
         pos_matrix = np.array(self.accepted_tokens_per_pos_lists)
         acceptance_rates = np.sum(pos_matrix, axis=0) / num_drafts