Add SchedulerRequestReceiver and route request-ingress state through it (#25609)

python/sglang/srt/disaggregation/decode.py

@@ -1573,7 +1573,9 @@ class SchedulerDisaggregationDecodeMixin:
 
         while True:
             # Receive requests
-            recv_reqs = self.recv_requests()
+            recv_reqs = self.recv_requests(
+                self.request_receiver,
+            )
             self.process_input_requests(recv_reqs)
             self.process_decode_queue()
             if self._engine_paused:
@@ -1601,7 +1603,9 @@ class SchedulerDisaggregationDecodeMixin:
 
         while True:
             # Receive requests
-            recv_reqs = self.recv_requests()
+            recv_reqs = self.recv_requests(
+                self.request_receiver,
+            )
             self.process_input_requests(recv_reqs)
             self.process_decode_queue()
             if self._engine_paused:

python/sglang/srt/disaggregation/prefill.py

@@ -395,7 +395,9 @@ class SchedulerDisaggregationPrefillMixin:
 
         while True:
             # Receive requests
-            recv_reqs = self.recv_requests()
+            recv_reqs = self.recv_requests(
+                self.request_receiver,
+            )
             self.process_input_requests(recv_reqs)
             self.waiting_queue.extend(
                 self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
@@ -428,7 +430,9 @@ class SchedulerDisaggregationPrefillMixin:
 
         while True:
             # Receive requests
-            recv_reqs = self.recv_requests()
+            recv_reqs = self.recv_requests(
+                self.request_receiver,
+            )
             self.process_input_requests(recv_reqs)
             self.waiting_queue.extend(
                 self.disagg_prefill_bootstrap_queue.pop_bootstrapped()

python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py

@@ -168,7 +168,9 @@ class SchedulerMlxOverlapMixin:
             )
 
         while True:
-            recv_reqs = self.recv_requests()
+            recv_reqs = self.recv_requests(
+                self.request_receiver,
+            )
             self.process_input_requests(recv_reqs)
             if self._engine_paused:
                 continue

python/sglang/srt/managers/scheduler.py

@@ -168,6 +168,9 @@ from sglang.srt.managers.schedule_policy import (
     PrefillAdder,
     SchedulePolicy,
 )
+from sglang.srt.managers.scheduler_components.request_receiver import (
+    SchedulerRequestReceiver,
+)
 from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin
 from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker
 from sglang.srt.managers.scheduler_output_processor_mixin import (
@@ -563,6 +566,29 @@ class Scheduler(
         # Init the grammar backend for constrained generation
         self.grammar_manager = GrammarManager(self)
 
+        self.request_receiver = SchedulerRequestReceiver(
+            recv_from_tokenizer=self.recv_from_tokenizer,
+            recv_from_rpc=self.recv_from_rpc,
+            recv_skipper=self.recv_skipper,
+            input_blocker=self.input_blocker,
+            mm_receiver=self.mm_receiver,
+            ps=self.ps,
+            tp_group=self.tp_group,
+            tp_cpu_group=self.tp_cpu_group,
+            attn_tp_group=self.attn_tp_group,
+            attn_tp_cpu_group=self.attn_tp_cpu_group,
+            attn_cp_group=self.attn_cp_group,
+            attn_cp_cpu_group=self.attn_cp_cpu_group,
+            world_group=self.world_group,
+            server_args=self.server_args,
+            model_config=self.model_config,
+            max_recv_per_poll=self.max_recv_per_poll,
+            stream_output=self.stream_output,
+            get_last_forward_mode=lambda: (
+                self.last_batch.forward_mode if self.last_batch is not None else None
+            ),
+        )
+
         self.is_initializing = False
 
     def init_zbal_on_npu(self):
@@ -1362,7 +1388,9 @@ class Scheduler(
         """A normal scheduler loop."""
         while True:
             # Receive requests
-            recv_reqs = self.recv_requests()
+            recv_reqs = self.recv_requests(
+                self.request_receiver,
+            )
             self.process_input_requests(recv_reqs)
             if self._engine_paused:
                 continue
@@ -1398,7 +1426,9 @@ class Scheduler(
 
         while True:
             # Receive requests
-            recv_reqs = self.recv_requests()
+            recv_reqs = self.recv_requests(
+                self.request_receiver,
+            )
             self.process_input_requests(recv_reqs)
             if self._engine_paused:
                 continue
@@ -1472,21 +1502,22 @@ class Scheduler(
 
         return disable_overlap_for_batch or need_grammar_sync
 
-    def recv_limit_reached(self, num_recv_reqs: int) -> bool:
+    @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,
+        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:
-            last_forward_mode = (
+            if not self.recv_skipper.handle(self.get_last_forward_mode()):
-                self.last_batch.forward_mode if self.last_batch is not None else None
-            )
-            if not self.recv_skipper.handle(last_forward_mode):
                 return []
 
         if self.ps.pp_rank == 0:
@@ -1495,7 +1526,7 @@ class Scheduler(
 
                 while True:
                     try:
-                        if self.recv_limit_reached(len(recv_reqs)):
+                        if Scheduler.recv_limit_reached(self, len(recv_reqs)):
                             break
                         recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
                     except zmq.ZMQError:
@@ -1504,7 +1535,7 @@ class Scheduler(
 
                 while True:
                     try:
-                        if self.recv_limit_reached(len(recv_reqs)):
+                        if Scheduler.recv_limit_reached(self, len(recv_reqs)):
                             break
                         recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK)
                     except zmq.ZMQError:
@@ -1530,7 +1561,9 @@ class Scheduler(
 
         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)
+                work_reqs, control_reqs = Scheduler._split_work_and_control_reqs(
+                    self, recv_reqs
+                )
             else:
                 work_reqs = None
                 control_reqs = None
@@ -1634,7 +1667,8 @@ class Scheduler(
 
         return recv_reqs
 
-    def _split_work_and_control_reqs(self, recv_reqs: List):
+    @staticmethod
+    def _split_work_and_control_reqs(self: "SchedulerRequestReceiver", recv_reqs: List):
         work_reqs = [
             req
             for req in recv_reqs

python/sglang/srt/managers/scheduler_components/request_receiver.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import zmq
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+    from sglang.srt.server_args import ServerArgs
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerRequestReceiver:
+    recv_from_tokenizer: zmq.Socket
+    recv_from_rpc: Optional[zmq.Socket]
+    recv_skipper: Any
+    input_blocker: Any
+    mm_receiver: Any
+    ps: "ParallelState"
+    tp_group: Any
+    tp_cpu_group: Any
+    attn_tp_group: Any
+    attn_tp_cpu_group: Any
+    attn_cp_group: Any
+    attn_cp_cpu_group: Any
+    world_group: Any
+    server_args: "ServerArgs"
+    model_config: "ModelConfig"
+    max_recv_per_poll: int
+    stream_output: Callable[..., None]
+    get_last_forward_mode: Callable[[], Any]

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -80,7 +80,9 @@ 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()
+                    recv_reqs = self.recv_requests(
+                        self.request_receiver,
+                    )
                     self.process_input_requests(recv_reqs)
                 if not self.pp_group.is_last_rank:
                     self._pp_commit_comm_work(self.send_req_work)
@@ -214,7 +216,9 @@ class SchedulerPPMixin:
                 d2h_event = None
                 next_batch_result = None
 
-                recv_reqs = self.recv_requests()
+                recv_reqs = self.recv_requests(
+                    self.request_receiver,
+                )
                 self.process_input_requests(recv_reqs)
 
                 if not self.pp_group.is_last_rank:
@@ -360,7 +364,9 @@ class SchedulerPPMixin:
                 d2h_event = None
                 next_batch_result = None
 
-                recv_reqs = self.recv_requests()
+                recv_reqs = self.recv_requests(
+                    self.request_receiver,
+                )
                 self.process_input_requests(recv_reqs)
 
                 if not self.pp_group.is_last_rank:

python/sglang/srt/multiplex/multiplexing_mixin.py

@@ -110,7 +110,9 @@ class SchedulerMultiplexMixin:
         while True:
             with torch.cuda.stream(decode_stream):
                 set_pdmux_status(False)
-                recv_reqs = self.recv_requests()
+                recv_reqs = self.recv_requests(
+                    self.request_receiver,
+                )
                 self.process_input_requests(recv_reqs)
 
             with torch.cuda.stream(prefill_stream):