From dc88b4eeb4646b6216cd5459387a64ab8a727a17 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 18 May 2026 18:43:48 +0800 Subject: [PATCH] Stand up SchedulerOutputStreamer; migrate output-streaming state to it (#25634) --- python/sglang/srt/disaggregation/decode.py | 20 ++++-- python/sglang/srt/disaggregation/prefill.py | 9 ++- python/sglang/srt/dllm/mixin/scheduler.py | 2 +- python/sglang/srt/managers/scheduler.py | 21 +++++- .../scheduler_components/output_streamer.py | 33 +++++++++ .../scheduler_output_processor_mixin.py | 68 ++++++++++++------- 6 files changed, 119 insertions(+), 34 deletions(-) create mode 100644 python/sglang/srt/managers/scheduler_components/output_streamer.py diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 1f0101540..2f9218388 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -532,7 +532,9 @@ class DecodePreallocQueue: message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_total_num_tokens}" logger.error(message) prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) - self.scheduler.stream_output([req], req.return_logprob) + self.scheduler.stream_output( + self.scheduler.output_streamer, [req], req.return_logprob + ) return True if self._uses_swa_tail_prealloc(): _, swa_required = self._prealloc_required_tokens(req) @@ -544,7 +546,9 @@ class DecodePreallocQueue: ) logger.error(message) prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) - self.scheduler.stream_output([req], req.return_logprob) + self.scheduler.stream_output( + self.scheduler.output_streamer, [req], req.return_logprob + ) return True return False @@ -779,7 +783,9 @@ class DecodePreallocQueue: continue if isinstance(decode_req.req.finished_reason, FINISH_ABORT): self.scheduler.stream_output( - [decode_req.req], decode_req.req.return_logprob + self.scheduler.output_streamer, + [decode_req.req], + decode_req.req.return_logprob, ) failed_reqs.append(decode_req) indices_to_remove.add(i) @@ -1509,7 +1515,9 @@ class DecodeTransferQueue: status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) self.scheduler.stream_output( - [decode_req.req], decode_req.req.return_logprob + self.scheduler.output_streamer, + [decode_req.req], + decode_req.req.return_logprob, ) if self.scheduler.enable_hisparse: self.scheduler.hisparse_coordinator.request_finished(decode_req.req) @@ -1526,7 +1534,9 @@ class DecodeTransferQueue: # Check if request was aborted due to corruption if isinstance(decode_req.req.finished_reason, FINISH_ABORT): self.scheduler.stream_output( - [decode_req.req], decode_req.req.return_logprob + self.scheduler.output_streamer, + [decode_req.req], + decode_req.req.return_logprob, ) if self.scheduler.enable_hisparse: self.scheduler.hisparse_coordinator.request_finished( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 9c09e01c8..d98a44134 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -252,7 +252,9 @@ class PrefillBootstrapQueue: logger.error(message) req.time_stats.trace_ctx.abort(abort_info={"reason": message}) prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) - self.scheduler.stream_output([req], req.return_logprob) + self.scheduler.stream_output( + self.scheduler.output_streamer, [req], req.return_logprob + ) return True return False @@ -309,7 +311,9 @@ class PrefillBootstrapQueue: prepare_abort( req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR ) - self.scheduler.stream_output([req], req.return_logprob) + self.scheduler.stream_output( + self.scheduler.output_streamer, [req], req.return_logprob + ) indices_to_remove.add(i) failed_reqs.append(req) if self.scheduler.metrics_reporter.enable_metrics: @@ -691,6 +695,7 @@ class SchedulerDisaggregationPrefillMixin: # Stream requests which have finished transfer self.stream_output( + self.output_streamer, done_reqs, any(req.return_logprob for req in done_reqs), None, diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 39ea06a99..703bf4988 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -89,7 +89,7 @@ class SchedulerDllmMixin: release_kv_cache(req, self.tree_cache) req.time_stats.set_completion_time() - self.stream_output(batch.reqs, batch.return_logprob) + self.stream_output(self.output_streamer, batch.reqs, batch.return_logprob) self.token_to_kv_pool_allocator.free_group_end() can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 531e76dd1..6e36b3398 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -184,6 +184,9 @@ from sglang.srt.managers.scheduler_components.metrics_reporter import ( PrefillStats, SchedulerMetricsReporter, ) +from sglang.srt.managers.scheduler_components.output_streamer import ( + SchedulerOutputStreamer, +) from sglang.srt.managers.scheduler_components.pool_stats_observer import ( SchedulerPoolStatsObserver, ) @@ -648,7 +651,9 @@ class Scheduler( server_args=self.server_args, model_config=self.model_config, max_recv_per_poll=self.max_recv_per_poll, - stream_output=self.stream_output, + stream_output=lambda *a, **kw: self.stream_output( + self.output_streamer, *a, **kw + ), get_last_forward_mode=lambda: ( self.last_batch.forward_mode if self.last_batch is not None else None ), @@ -742,6 +747,18 @@ class Scheduler( model_config=self.model_config, ) + self.output_streamer = SchedulerOutputStreamer( + send_to_detokenizer=self.send_to_detokenizer, + tree_cache=self.tree_cache, + ps=self.ps, + server_args=self.server_args, + is_generation=self.is_generation, + spec_algorithm=self.spec_algorithm, + disaggregation_mode=self.disaggregation_mode, + enable_hicache_storage=lambda: self.enable_hicache_storage, + load_inquirer_get_loads=lambda req: self.load_inquirer.get_loads(req), + ) + self.is_initializing = False def init_zbal_on_npu(self): @@ -1913,7 +1930,7 @@ class Scheduler( abort_info={"reason": error_msg} ) prepare_abort(req, error_msg, status_code=HTTPStatus.BAD_REQUEST) - self.stream_output([req], req.return_logprob) + self.stream_output(self.output_streamer, [req], req.return_logprob) return elif ( diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py new file mode 100644 index 000000000..eba315795 --- /dev/null +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any, Callable + +import zmq + +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.distributed.parallel_state_wrapper import ParallelState +from sglang.srt.environ import envs +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + +logger = logging.getLogger(__name__) + + +DEFAULT_FORCE_STREAM_INTERVAL = envs.SGLANG_FORCE_STREAM_INTERVAL.get() + + +@dataclass(kw_only=True, slots=True) +class SchedulerOutputStreamer: + send_to_detokenizer: zmq.Socket + tree_cache: BasePrefixCache + ps: ParallelState + server_args: ServerArgs + is_generation: bool + spec_algorithm: SpeculativeAlgorithm + disaggregation_mode: DisaggregationMode + enable_hicache_storage: Callable[[], bool] + load_inquirer_get_loads: Callable[..., Any] + _test_stream_output_count: int = 0 diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 2eddfc61d..848de4c42 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -33,6 +33,9 @@ if TYPE_CHECKING: ScheduleBatch, Scheduler, ) + from sglang.srt.managers.scheduler_components.output_streamer import ( + SchedulerOutputStreamer, + ) logger = logging.getLogger(__name__) @@ -47,7 +50,8 @@ class SchedulerOutputProcessorMixin: We put them into a separate file to make the `scheduler.py` shorter. """ - def _get_storage_backend_type(self) -> str: + @staticmethod + def _get_storage_backend_type(self: "SchedulerOutputStreamer") -> str: """Get storage backend type from tree_cache.""" storage_backend_type = "none" cache_controller = getattr(self.tree_cache, "cache_controller", None) @@ -57,7 +61,10 @@ class SchedulerOutputProcessorMixin: storage_backend_type = type(storage_backend).__name__ return storage_backend_type - def _get_cached_tokens_details(self: Scheduler, req: Req) -> Optional[dict]: + @staticmethod + def get_cached_tokens_details( + self: "SchedulerOutputStreamer", req: Req + ) -> Optional[dict]: """Get detailed cache breakdown for a request, if available. Returns: @@ -75,9 +82,11 @@ class SchedulerOutputProcessorMixin: "host": req.cached_tokens_host, } # Only include storage fields if L3 storage is enabled - if self.enable_hicache_storage: + if self.enable_hicache_storage(): details["storage"] = req.cached_tokens_storage - details["storage_backend"] = self._get_storage_backend_type() + details["storage_backend"] = ( + SchedulerOutputProcessorMixin._get_storage_backend_type(self) + ) return details if req.cached_tokens > 0: @@ -103,7 +112,7 @@ class SchedulerOutputProcessorMixin: release_kv_cache(req, self.tree_cache) # Note: Logprobs should be handled on the prefill engine. - self.stream_output(batch.reqs, batch.return_logprob) + self.stream_output(self.output_streamer, batch.reqs, batch.return_logprob) if use_free_group: self.token_to_kv_pool_allocator.free_group_end() @@ -402,7 +411,9 @@ class SchedulerOutputProcessorMixin: req.is_chunked -= 1 req.time_stats.set_last_chunked_prefill_finish_time() - self.stream_output(batch.reqs, batch.return_logprob, skip_stream_req) + self.stream_output( + self.output_streamer, batch.reqs, batch.return_logprob, skip_stream_req + ) can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) self.metrics_reporter.report_prefill_stats( @@ -467,8 +478,8 @@ class SchedulerOutputProcessorMixin: if result.copy_done is not None: result.copy_done.synchronize() - self.stream_output_generation( - batch.reqs, batch.return_logprob, is_idle_batch=True + self._stream_output_generation( + self.output_streamer, batch.reqs, batch.return_logprob, is_idle_batch=True ) def process_batch_result_decode( @@ -629,7 +640,7 @@ class SchedulerOutputProcessorMixin: self.abort_request(AbortReq(rid=req.rid)) req.grammar.finished = req.finished() - self.stream_output(batch.reqs, batch.return_logprob) + self.stream_output(self.output_streamer, batch.reqs, batch.return_logprob) self.token_to_kv_pool_allocator.free_group_end() self.metrics_reporter.forward_ct_decode = ( @@ -715,24 +726,28 @@ class SchedulerOutputProcessorMixin: actual_seq_len // mamba_track_interval * mamba_track_interval ) + @staticmethod def stream_output( - self: Scheduler, + self: "SchedulerOutputStreamer", reqs: List[Req], return_logprob: bool, skip_req: Optional[Req] = None, ): """Stream the output to detokenizer.""" if self.is_generation: - self.stream_output_generation(reqs, return_logprob, skip_req) + SchedulerOutputProcessorMixin._stream_output_generation( + self, reqs, return_logprob, skip_req + ) else: # embedding or reward model - self.stream_output_embedding(reqs) + SchedulerOutputProcessorMixin._stream_output_embedding(self, reqs) if envs.SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS.get() > 0: - self._trigger_crash_for_tests( - envs.SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS.get() + SchedulerOutputProcessorMixin._trigger_crash_for_tests( + self, envs.SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS.get() ) - def _trigger_crash_for_tests(self: Scheduler, crash_threshold: int): + @staticmethod + def _trigger_crash_for_tests(self: "SchedulerOutputStreamer", crash_threshold: int): # Crash trigger: crash after stream_output is called N times # This is used for testing purposes. if not hasattr(self, "_test_stream_output_count"): @@ -743,8 +758,9 @@ class SchedulerOutputProcessorMixin: f"Test crash after stream_output called {self._test_stream_output_count} times" ) - def stream_output_generation( - self: Scheduler, + @staticmethod + def _stream_output_generation( + self: "SchedulerOutputStreamer", reqs: List[Req], return_logprob: bool, skip_req: Optional[Req] = None, @@ -772,9 +788,7 @@ class SchedulerOutputProcessorMixin: spec_correct_drafts_histogram = [] retraction_counts = [] output_hidden_states = None - load = self.load_inquirer.get_loads( - GetLoadsReqInput(include=["core"]), - ) + load = self.load_inquirer_get_loads(GetLoadsReqInput(include=["core"])) routed_experts = None indexer_topk = None customized_info = {} @@ -821,7 +835,8 @@ class SchedulerOutputProcessorMixin: else: if req.stream: stream_interval = ( - req.sampling_params.stream_interval or self.stream_interval + req.sampling_params.stream_interval + or self.server_args.stream_interval ) # origin stream_interval logic @@ -872,7 +887,9 @@ class SchedulerOutputProcessorMixin: cached_tokens.append(req.cached_tokens) # Collect detailed cache breakdown if available - cached_tokens_details.append(self._get_cached_tokens_details(req)) + cached_tokens_details.append( + SchedulerOutputProcessorMixin.get_cached_tokens_details(self, req) + ) retraction_counts.append(req.retraction_count) @@ -1032,7 +1049,8 @@ class SchedulerOutputProcessorMixin: ) ) - def stream_output_embedding(self: Scheduler, reqs: List[Req]): + @staticmethod + def _stream_output_embedding(self: "SchedulerOutputStreamer", reqs: List[Req]): rids = [] http_worker_ipcs = [] finished_reasons: List[BaseFinishReason] = [] @@ -1055,7 +1073,9 @@ class SchedulerOutputProcessorMixin: cached_tokens.append(req.cached_tokens) # Collect detailed cache breakdown if available - cached_tokens_details.append(self._get_cached_tokens_details(req)) + cached_tokens_details.append( + SchedulerOutputProcessorMixin.get_cached_tokens_details(self, req) + ) time_stats.append(req.time_stats) retraction_counts.append(req.retraction_count)