diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 4c7db7921..3c26570ed 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1573,9 +1573,7 @@ class SchedulerDisaggregationDecodeMixin: while True: # Receive requests - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) self.process_decode_queue() if self._engine_paused: @@ -1603,9 +1601,7 @@ class SchedulerDisaggregationDecodeMixin: while True: # Receive requests - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) self.process_decode_queue() if self._engine_paused: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 0d8e080ca..5cb6f9efc 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -395,9 +395,7 @@ class SchedulerDisaggregationPrefillMixin: while True: # Receive requests - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() @@ -430,9 +428,7 @@ class SchedulerDisaggregationPrefillMixin: while True: # Receive requests - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index 7cbe67965..6e3ac9de7 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -168,9 +168,7 @@ class SchedulerMlxOverlapMixin: ) while True: - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: continue diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 359b3883d..58f1e5787 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -148,10 +148,6 @@ from sglang.srt.managers.io_struct import ( UpdateWeightsFromIPCReqInput, UpdateWeightsFromTensorReqInput, ) -from sglang.srt.managers.mm_utils import ( - has_shm_features, - unwrap_shm_features, -) from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors from sglang.srt.managers.prefill_delayer import ( PrefillDelayer, @@ -212,7 +208,6 @@ from sglang.srt.session.session_controller import SessionController from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( DynamicGradMode, - broadcast_pyobj, configure_gc_logger, configure_logger, freeze_gc, @@ -221,7 +216,6 @@ from sglang.srt.utils import ( get_int_env_var, is_mps, kill_itself_when_parent_died, - point_to_point_pyobj, require_mlp_sync, set_gpu_proc_affinity, set_random_seed, @@ -1388,9 +1382,7 @@ class Scheduler( """A normal scheduler loop.""" while True: # Receive requests - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: continue @@ -1426,9 +1418,7 @@ class Scheduler( while True: # Receive requests - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: continue @@ -1502,201 +1492,6 @@ class Scheduler( return disable_overlap_for_batch or need_grammar_sync - @staticmethod - def recv_limit_reached( - self: "SchedulerRequestReceiver", num_recv_reqs: int - ) -> bool: - if self.max_recv_per_poll < 0: - return False - return num_recv_reqs >= self.max_recv_per_poll - - @staticmethod - def recv_requests( - self: "SchedulerRequestReceiver", - ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]: - """Receive results at tp_rank = 0 and broadcast it to all other TP ranks.""" - - if self.recv_skipper is not None: - if not self.recv_skipper.handle(self.get_last_forward_mode()): - return [] - - if self.ps.pp_rank == 0: - if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: - recv_reqs = [] - - while True: - try: - if Scheduler.recv_limit_reached(self, len(recv_reqs)): - break - recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) - except zmq.ZMQError: - break - recv_reqs.append(recv_req) - - while True: - try: - if Scheduler.recv_limit_reached(self, len(recv_reqs)): - break - recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK) - except zmq.ZMQError: - break - recv_reqs.append(recv_rpc) - else: - recv_reqs = None - else: - if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: - dp_offset = self.ps.attn_dp_rank * self.ps.attn_tp_size - recv_reqs = point_to_point_pyobj( - [], - self.ps.pp_rank * self.ps.tp_size + dp_offset, - self.world_group.cpu_group, - (self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset, - self.ps.pp_rank * self.ps.tp_size + dp_offset, - ) - else: - recv_reqs = None - - if self.input_blocker is not None: - recv_reqs = self.input_blocker.handle(recv_reqs) - - if self.server_args.enable_dp_attention: - if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: - work_reqs, control_reqs = Scheduler._split_work_and_control_reqs( - self, recv_reqs - ) - else: - work_reqs = None - control_reqs = None - - if self.ps.attn_tp_size != 1: - work_reqs = broadcast_pyobj( - work_reqs, - self.attn_tp_group.rank, - self.attn_tp_cpu_group, - src=self.attn_tp_group.ranks[0], - ) - - if self.ps.attn_cp_size != 1: - work_reqs = broadcast_pyobj( - work_reqs, - self.attn_cp_group.rank, - self.attn_cp_cpu_group, - src=self.attn_cp_group.ranks[0], - ) - - # When dp_attention_local_control_broadcast is enabled, each DP - # group leader already receives control messages from the DP - # controller, so we broadcast within attn_tp_group + attn_cp_group - # instead of the full tp_group. This avoids an expensive - # all-ranks gloo sync. - _local_ctrl = self.server_args.enable_dp_attention_local_control_broadcast - if _local_ctrl: - if self.ps.attn_tp_size != 1: - control_reqs = broadcast_pyobj( - control_reqs, - self.attn_tp_group.rank, - self.attn_tp_cpu_group, - src=self.attn_tp_group.ranks[0], - ) - if self.ps.attn_cp_size != 1: - control_reqs = broadcast_pyobj( - control_reqs, - self.attn_cp_group.rank, - self.attn_cp_cpu_group, - src=self.attn_cp_group.ranks[0], - ) - elif self.ps.tp_size != 1: - control_reqs = broadcast_pyobj( - control_reqs, - self.tp_group.rank, - self.tp_cpu_group, - src=self.tp_group.ranks[0], - ) - recv_reqs = work_reqs + control_reqs - elif self.ps.tp_size != 1: - recv_reqs = broadcast_pyobj( - recv_reqs, - self.tp_group.rank, - self.tp_cpu_group, - src=self.tp_group.ranks[0], - ) - - # Process MM requests under EPD-disaggregation mode - if ( - self.ps.pp_rank == 0 - and self.server_args.language_only - and self.server_args.encoder_transfer_backend == "zmq_to_scheduler" - ): - recv_reqs, abort_reqs = self.mm_receiver.process_waiting_requests(recv_reqs) - for req, error_msg, error_code in abort_reqs: - status_code = ( - HTTPStatus.BAD_REQUEST - if error_code == 400 - else HTTPStatus.INTERNAL_SERVER_ERROR - ) - prepare_abort(req, error_msg, status_code=status_code) - self.stream_output([req], req.return_logprob) - - # Unwrap shared memory features AFTER all broadcasts complete, - # so that ShmPointerMMData metadata (not full tensor data) is what - # gets serialized during broadcast_pyobj. - if recv_reqs: - # Barrier for the non-DP-attention path only: there is a single - # broadcast_pyobj on tp_cpu_group where the source rank returns - # the original objects immediately while other ranks are still in - # pickle.loads (-> __setstate__ -> shm_open). Without a barrier - # the source can call materialize() / shm_unlink before others - # open the segment. recv_reqs is consistent across all ranks - # here (same broadcast), so the guard is deadlock-free. - # - # Under DP-attention no barrier is needed: the control_reqs - # broadcast on tp_cpu_group (step 3) is a collective that forces - # every rank to complete the earlier attn_tp / attn_cp work_reqs - # deserializations (steps 1-2, which call shm_open) before any - # rank returns from step 3. POSIX guarantees shm_unlink only - # removes the name; already-open handles stay valid. - if ( - not self.server_args.enable_dp_attention - and self.ps.tp_size > 1 - and self.model_config.is_multimodal - and has_shm_features(recv_reqs) - ): - barrier(group=self.tp_cpu_group) - for req in recv_reqs: - unwrap_shm_features(req) - - return recv_reqs - - @staticmethod - def _split_work_and_control_reqs(self: "SchedulerRequestReceiver", recv_reqs: List): - work_reqs = [ - req - for req in recv_reqs - if isinstance( - req, - ( - TokenizedGenerateReqInput, - TokenizedEmbeddingReqInput, - BatchTokenizedGenerateReqInput, - BatchTokenizedEmbeddingReqInput, - ), - ) - ] - control_reqs = [ - req - for req in recv_reqs - if not isinstance( - req, - ( - TokenizedGenerateReqInput, - TokenizedEmbeddingReqInput, - BatchTokenizedGenerateReqInput, - BatchTokenizedEmbeddingReqInput, - ), - ) - ] - return work_reqs, control_reqs - def process_input_requests(self, recv_reqs: List): now = time.monotonic() self.session_controller.maybe_reap(now) diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index aec7b4326..fe29bd690 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -1,9 +1,34 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Optional +from http import HTTPStatus +from typing import ( + TYPE_CHECKING, + Any, + Callable, + List, + Optional, + Union, +) import zmq +from torch.distributed import barrier + +from sglang.srt.disaggregation.utils import prepare_abort +from sglang.srt.managers.io_struct import ( + BatchTokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, +) +from sglang.srt.managers.mm_utils import ( + has_shm_features, + unwrap_shm_features, +) +from sglang.srt.utils import ( + broadcast_pyobj, + point_to_point_pyobj, +) if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig @@ -31,3 +56,191 @@ class SchedulerRequestReceiver: max_recv_per_poll: int stream_output: Callable[..., None] get_last_forward_mode: Callable[[], Any] + + def recv_limit_reached(self, num_recv_reqs: int) -> bool: + if self.max_recv_per_poll < 0: + return False + return num_recv_reqs >= self.max_recv_per_poll + + def recv_requests( + self, + ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]: + """Receive results at tp_rank = 0 and broadcast it to all other TP ranks.""" + + if self.recv_skipper is not None: + if not self.recv_skipper.handle(self.get_last_forward_mode()): + return [] + + if self.ps.pp_rank == 0: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + recv_reqs = [] + + while True: + try: + if self.recv_limit_reached(len(recv_reqs)): + break + recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) + except zmq.ZMQError: + break + recv_reqs.append(recv_req) + + while True: + try: + if self.recv_limit_reached(len(recv_reqs)): + break + recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK) + except zmq.ZMQError: + break + recv_reqs.append(recv_rpc) + else: + recv_reqs = None + else: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + dp_offset = self.ps.attn_dp_rank * self.ps.attn_tp_size + recv_reqs = point_to_point_pyobj( + [], + self.ps.pp_rank * self.ps.tp_size + dp_offset, + self.world_group.cpu_group, + (self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset, + self.ps.pp_rank * self.ps.tp_size + dp_offset, + ) + else: + recv_reqs = None + + if self.input_blocker is not None: + recv_reqs = self.input_blocker.handle(recv_reqs) + + if self.server_args.enable_dp_attention: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs) + else: + work_reqs = None + control_reqs = None + + if self.ps.attn_tp_size != 1: + work_reqs = broadcast_pyobj( + work_reqs, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + + if self.ps.attn_cp_size != 1: + work_reqs = broadcast_pyobj( + work_reqs, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + + # When dp_attention_local_control_broadcast is enabled, each DP + # group leader already receives control messages from the DP + # controller, so we broadcast within attn_tp_group + attn_cp_group + # instead of the full tp_group. This avoids an expensive + # all-ranks gloo sync. + _local_ctrl = self.server_args.enable_dp_attention_local_control_broadcast + if _local_ctrl: + if self.ps.attn_tp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + if self.ps.attn_cp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + elif self.ps.tp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.tp_group.rank, + self.tp_cpu_group, + src=self.tp_group.ranks[0], + ) + recv_reqs = work_reqs + control_reqs + elif self.ps.tp_size != 1: + recv_reqs = broadcast_pyobj( + recv_reqs, + self.tp_group.rank, + self.tp_cpu_group, + src=self.tp_group.ranks[0], + ) + + # Process MM requests under EPD-disaggregation mode + if ( + self.ps.pp_rank == 0 + and self.server_args.language_only + and self.server_args.encoder_transfer_backend == "zmq_to_scheduler" + ): + recv_reqs, abort_reqs = self.mm_receiver.process_waiting_requests(recv_reqs) + for req, error_msg, error_code in abort_reqs: + status_code = ( + HTTPStatus.BAD_REQUEST + if error_code == 400 + else HTTPStatus.INTERNAL_SERVER_ERROR + ) + prepare_abort(req, error_msg, status_code=status_code) + self.stream_output([req], req.return_logprob) + + # Unwrap shared memory features AFTER all broadcasts complete, + # so that ShmPointerMMData metadata (not full tensor data) is what + # gets serialized during broadcast_pyobj. + if recv_reqs: + # Barrier for the non-DP-attention path only: there is a single + # broadcast_pyobj on tp_cpu_group where the source rank returns + # the original objects immediately while other ranks are still in + # pickle.loads (-> __setstate__ -> shm_open). Without a barrier + # the source can call materialize() / shm_unlink before others + # open the segment. recv_reqs is consistent across all ranks + # here (same broadcast), so the guard is deadlock-free. + # + # Under DP-attention no barrier is needed: the control_reqs + # broadcast on tp_cpu_group (step 3) is a collective that forces + # every rank to complete the earlier attn_tp / attn_cp work_reqs + # deserializations (steps 1-2, which call shm_open) before any + # rank returns from step 3. POSIX guarantees shm_unlink only + # removes the name; already-open handles stay valid. + if ( + not self.server_args.enable_dp_attention + and self.ps.tp_size > 1 + and self.model_config.is_multimodal + and has_shm_features(recv_reqs) + ): + barrier(group=self.tp_cpu_group) + for req in recv_reqs: + unwrap_shm_features(req) + + return recv_reqs + + def _split_work_and_control_reqs(self, recv_reqs: List): + work_reqs = [ + req + for req in recv_reqs + if isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + control_reqs = [ + req + for req in recv_reqs + if not isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + return work_reqs, control_reqs diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index e6d378f04..354f5ef52 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -80,9 +80,7 @@ class SchedulerPPMixin: next_first_rank_mb_id = (mb_id + self.ps.pp_size) % self.pp_loop_size next_mb_id = (mb_id + 1) % self.pp_loop_size with torch.profiler.record_function("recv_requests"): - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if not self.pp_group.is_last_rank: self._pp_commit_comm_work(self.send_req_work) @@ -216,9 +214,7 @@ class SchedulerPPMixin: d2h_event = None next_batch_result = None - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if not self.pp_group.is_last_rank: @@ -364,9 +360,7 @@ class SchedulerPPMixin: d2h_event = None next_batch_result = None - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) if not self.pp_group.is_last_rank: diff --git a/python/sglang/srt/multiplex/multiplexing_mixin.py b/python/sglang/srt/multiplex/multiplexing_mixin.py index 7806c376b..d5fd32a33 100644 --- a/python/sglang/srt/multiplex/multiplexing_mixin.py +++ b/python/sglang/srt/multiplex/multiplexing_mixin.py @@ -110,9 +110,7 @@ class SchedulerMultiplexMixin: while True: with torch.cuda.stream(decode_stream): set_pdmux_status(False) - recv_reqs = self.recv_requests( - self.request_receiver, - ) + recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) with torch.cuda.stream(prefill_stream):