Add spec_accept_histogram request statistic (#18332)

This commit is contained in:
Scott Lee
2026-02-12 05:09:21 -08:00
committed by GitHub
parent 0abe4a22c6
commit c59b9223e6
8 changed files with 54 additions and 8 deletions

View File

@@ -378,6 +378,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
cached_tokens_details=recv_obj.cached_tokens_details,
spec_verify_ct=recv_obj.spec_verify_ct,
spec_accepted_tokens=recv_obj.spec_accepted_tokens,
spec_acceptance_histogram=recv_obj.spec_acceptance_histogram,
input_token_logprobs_val=recv_obj.input_token_logprobs_val,
input_token_logprobs_idx=recv_obj.input_token_logprobs_idx,
output_token_logprobs_val=recv_obj.output_token_logprobs_val,

View File

@@ -117,6 +117,12 @@ class SpeculativeDecodingMetricsMixin:
# Accepted tokens: Number of accepted tokens during speculative decoding
spec_accepted_tokens: List[int]
# Acceptance histogram: List of lists, where each inner list represents histogram counts.
# List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens.
# Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens.
# Empty list [] when speculative decoding is disabled.
spec_acceptance_histogram: List[List[int]]
@dataclass
class APIServingTimingMixin:

View File

@@ -128,6 +128,9 @@ def _handle_output_by_index(output, i):
spec_accepted_tokens=_extract_field_by_index(
output, "spec_accepted_tokens", i
),
spec_acceptance_histogram=_extract_field_by_index(
output, "spec_acceptance_histogram", i
),
queue_time=_extract_field_by_index(output, "queue_time", i),
forward_entry_time=_extract_field_by_index(output, "forward_entry_time", i),
prefill_launch_delay=_extract_field_by_index(
@@ -222,6 +225,9 @@ def _handle_output_by_index(output, i):
spec_accepted_tokens=_extract_field_by_index(
output, "spec_accepted_tokens", i
),
spec_acceptance_histogram=_extract_field_by_index(
output, "spec_acceptance_histogram", i
),
queue_time=_extract_field_by_index(output, "queue_time", i),
forward_entry_time=_extract_field_by_index(output, "forward_entry_time", i),
prefill_launch_delay=_extract_field_by_index(

View File

@@ -46,7 +46,7 @@ from enum import Enum, auto
from functools import lru_cache
from http import HTTPStatus
from itertools import chain
from typing import TYPE_CHECKING, Any, List, Optional, Set, Tuple, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
import numpy as np
import torch
@@ -770,6 +770,11 @@ class Req(ReqDllmMixin):
# This is used to compute the acceptance rate and average acceptance length per request.
self.spec_accepted_tokens = 0
# Acceptance histogram for speculative decoding.
# List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens.
# Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens.
self.spec_acceptance_histogram: List[int] = []
# The number of times this request has been retracted / preempted.
self.retraction_count = 0
self.retraction_mb_id = None
@@ -858,6 +863,18 @@ class Req(ReqDllmMixin):
)
self.last_tic = now
def update_spec_acceptance_histogram(self, accepted_draft_tokens: int):
"""Update the speculative decoding acceptance histogram.
Args:
accepted_draft_tokens: Number of draft tokens accepted in this step.
"""
if len(self.spec_acceptance_histogram) <= accepted_draft_tokens:
self.spec_acceptance_histogram.extend(
[0] * (accepted_draft_tokens - len(self.spec_acceptance_histogram) + 1)
)
self.spec_acceptance_histogram[accepted_draft_tokens] += 1
def extend_image_inputs(self, image_inputs):
if self.multimodal_inputs is None:
self.multimodal_inputs = image_inputs

View File

@@ -368,7 +368,10 @@ class SchedulerOutputProcessorMixin:
next_token_ids[i * stride : i * stride + accept_lens[i]]
)
req.spec_verify_ct += 1
req.spec_accepted_tokens += accept_lens[i] - 1
accepted_draft_tokens = result.accept_length_per_req_cpu[i]
req.spec_accepted_tokens += accepted_draft_tokens
req.update_spec_acceptance_histogram(accepted_draft_tokens)
return predict_tokens
@@ -929,6 +932,7 @@ class SchedulerOutputProcessorMixin:
cached_tokens_details = [] # Detailed breakdown by cache source
spec_verify_ct = []
spec_accepted_tokens = []
spec_acceptance_histogram = []
retraction_counts = []
output_hidden_states = None
load = self.get_load()
@@ -1059,6 +1063,7 @@ class SchedulerOutputProcessorMixin:
if not self.spec_algorithm.is_none():
spec_verify_ct.append(req.spec_verify_ct)
spec_accepted_tokens.append(req.spec_accepted_tokens)
spec_acceptance_histogram.append(req.spec_acceptance_histogram)
if return_logprob:
if (
@@ -1160,6 +1165,7 @@ class SchedulerOutputProcessorMixin:
http_worker_ipcs=http_worker_ipcs,
spec_verify_ct=spec_verify_ct,
spec_accepted_tokens=spec_accepted_tokens,
spec_acceptance_histogram=spec_acceptance_histogram,
queue_time=queue_times,
forward_entry_time=forward_entry_times,
prefill_launch_delay=prefill_launch_delays,

View File

@@ -1863,6 +1863,16 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
meta_info["spec_draft_token_num"] = total_draft_tokens
meta_info["spec_verify_ct"] = recv_obj.spec_verify_ct[i]
# Acceptance histogram: tracks how many decoding steps accepted a certain number of draft tokens.
if (
recv_obj.spec_acceptance_histogram
and len(recv_obj.spec_acceptance_histogram) > i
and recv_obj.spec_acceptance_histogram[i]
):
meta_info["spec_accept_histogram"] = recv_obj.spec_acceptance_histogram[
i
]
def _calculate_timing_metrics(
self,
meta_info: Dict[str, Any],

View File

@@ -426,9 +426,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
else:
unfinished_accept_index.append(accept_index[i])
req.spec_verify_ct += 1
req.spec_accepted_tokens += (
sum(1 for idx in accept_index_row if idx != -1) - 1
)
accepted_draft_tokens = sum(1 for idx in accept_index_row if idx != -1) - 1
req.spec_accepted_tokens += accepted_draft_tokens
req.update_spec_acceptance_histogram(accepted_draft_tokens)
if has_finished:
accept_length = (accept_index != -1).sum(dim=1) - 1

View File

@@ -188,9 +188,9 @@ class NgramVerifyInput(SpecInput):
)
raise e
req.spec_verify_ct += 1
req.spec_accepted_tokens += (
sum(1 for idx in accept_index_row if idx != -1) - 1
)
accepted_draft_tokens = sum(1 for idx in accept_index_row if idx != -1) - 1
req.spec_accepted_tokens += accepted_draft_tokens
req.update_spec_acceptance_histogram(accepted_draft_tokens)
if has_finished:
self.accept_length = (self.accepted_indices != -1).sum(dim=1) - 1