Add spec_accept_histogram request statistic (#18332)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user