diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 450e5d169..652227860 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -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, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 348cdba3c..c80b54be5 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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: diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 8b0bc6ffd..e1236aa0f 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -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( diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 519235db6..c07995798 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 5baaa5aed..f284bf961 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -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, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 2a21ef666..6c277d2fe 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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], diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 7718ba245..ac629c7ee 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -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 diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index 210ee505c..c18cf7965 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/python/sglang/srt/speculative/ngram_info.py @@ -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