minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fzyzcjy
2026-05-18 18:43:48 +08:00
committed by GitHub
parent 2cbe01d044
commit dc88b4eeb4
6 changed files with 119 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+15 -5
View File
@@ -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
+7 -2
View File
@@ -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
+1 -1
View File
@@ -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
+19 -2
View File
@@ -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
+33
View File
@@ -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
+44 -24
View File
@@ -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)