Stand up SchedulerOutputStreamer; migrate output-streaming state to it (#25634)

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(

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,

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)

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 (

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

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(
+        self._stream_output_generation(
-            batch.reqs, batch.return_logprob, is_idle_batch=True
+            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(
+            SchedulerOutputProcessorMixin._trigger_crash_for_tests(
-                envs.SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS.get()
+                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(
+    @staticmethod
-        self: Scheduler,
+    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(
+        load = self.load_inquirer_get_loads(GetLoadsReqInput(include=["core"]))
-            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)