[Refactor / Style] Unify all event loops (except for PP) (#12959)

python/sglang/srt/disaggregation/base/conn.py

@@ -141,6 +141,18 @@ class BaseKVReceiver(ABC):
         """
         ...
 
+    def clear(self):
+        """
+        Clear any internal states.
+        """
+        pass
+
+    def abort(self):
+        """
+        Abort the current transfer.
+        """
+        pass
+
 
 class BaseKVBootstrapServer(ABC):
     @abstractmethod

python/sglang/srt/disaggregation/decode.py

@@ -693,6 +693,50 @@ class DecodeTransferQueue:
     def extend(self, decode_reqs: List[DecodeRequest]) -> None:
         self.queue.extend(decode_reqs)
 
+    def _commit_transfer_to_req(self, decode_req: DecodeRequest) -> None:
+        idx = decode_req.metadata_buffer_index
+        (
+            output_id,
+            cached_tokens,
+            output_token_logprobs_val,
+            output_token_logprobs_idx,
+            output_top_logprobs_val,
+            output_top_logprobs_idx,
+            output_topk_p,
+            output_topk_index,
+            output_hidden_states,
+        ) = self.metadata_buffers.get_buf(idx)
+
+        decode_req.req.output_ids.append(output_id[0].item())
+        decode_req.req.cached_tokens = cached_tokens[0].item()
+        if not self.spec_algorithm.is_none():
+            decode_req.req.output_topk_p = output_topk_p
+            decode_req.req.output_topk_index = output_topk_index
+            decode_req.req.hidden_states_tensor = output_hidden_states
+
+        if decode_req.req.return_logprob:
+            decode_req.req.output_token_logprobs_val.append(
+                output_token_logprobs_val[0].item()
+            )
+            decode_req.req.output_token_logprobs_idx.append(
+                output_token_logprobs_idx[0].item()
+            )
+            decode_req.req.output_top_logprobs_val.append(
+                output_top_logprobs_val[: decode_req.req.top_logprobs_num].tolist()
+            )
+            decode_req.req.output_top_logprobs_idx.append(
+                output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist()
+            )
+
+        decode_req.kv_receiver.clear()
+        decode_req.kv_receiver = None
+        trace_slice_end(
+            RequestStage.DECODE_TRANSFERRED,
+            decode_req.req.rid,
+            auto_next_anon=True,
+        )
+        decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
+
     def pop_transferred(self) -> List[Req]:
         if not self.queue:
             return []
@@ -725,58 +769,9 @@ class DecodeTransferQueue:
                     self.scheduler.metrics_collector.increment_transfer_failed_reqs()
                 continue
             elif poll == KVPoll.Success:
-
+                self._commit_transfer_to_req(decode_req)
-                idx = decode_req.metadata_buffer_index
-                (
-                    output_id,
-                    cached_tokens,
-                    output_token_logprobs_val,
-                    output_token_logprobs_idx,
-                    output_top_logprobs_val,
-                    output_top_logprobs_idx,
-                    output_topk_p,
-                    output_topk_index,
-                    output_hidden_states,
-                ) = self.metadata_buffers.get_buf(idx)
-
-                decode_req.req.output_ids.append(output_id[0].item())
-                decode_req.req.cached_tokens = cached_tokens[0].item()
-                if not self.spec_algorithm.is_none():
-                    decode_req.req.output_topk_p = output_topk_p
-                    decode_req.req.output_topk_index = output_topk_index
-                    decode_req.req.hidden_states_tensor = output_hidden_states
-
-                if decode_req.req.return_logprob:
-                    decode_req.req.output_token_logprobs_val.append(
-                        output_token_logprobs_val[0].item()
-                    )
-                    decode_req.req.output_token_logprobs_idx.append(
-                        output_token_logprobs_idx[0].item()
-                    )
-                    decode_req.req.output_top_logprobs_val.append(
-                        output_top_logprobs_val[
-                            : decode_req.req.top_logprobs_num
-                        ].tolist()
-                    )
-                    decode_req.req.output_top_logprobs_idx.append(
-                        output_top_logprobs_idx[
-                            : decode_req.req.top_logprobs_num
-                        ].tolist()
-                    )
-
-                if hasattr(decode_req.kv_receiver, "clear"):
-                    decode_req.kv_receiver.clear()
-                decode_req.kv_receiver = None
-
                 indices_to_remove.add(i)
-                decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
                 transferred_reqs.append(decode_req.req)
-                trace_slice_end(
-                    RequestStage.DECODE_TRANSFERRED,
-                    decode_req.req.rid,
-                    auto_next_anon=True,
-                )
-
             elif poll in [
                 KVPoll.Bootstrapping,
                 KVPoll.WaitingForInput,

python/sglang/srt/disaggregation/prefill.py

@@ -310,6 +310,37 @@ class SchedulerDisaggregationPrefillMixin:
     Mixin for Scheduler to handle disaggregation prefill
     """
 
+    def get_next_disagg_prefill_batch_to_run(
+        self: Scheduler,
+    ) -> Optional[ScheduleBatch]:
+        if self.last_batch and self.last_batch.forward_mode.is_extend():
+            if self.chunked_req:
+                # Move the chunked request out of the batch so that we can merge
+                # only finished requests to running_batch.
+                self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
+                self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
+                if self.enable_overlap:
+                    # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
+                    self.chunked_req.tmp_end_idx = min(
+                        len(self.chunked_req.fill_ids),
+                        len(self.chunked_req.origin_input_ids),
+                    )
+                else:
+                    self.send_kv_chunk(self.chunked_req)
+                # chunked request keeps its rid but will get a new req_pool_idx
+                self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
+                self.running_batch.batch_is_full = False
+
+        batch = self.get_new_batch_prefill()
+        if self.require_mlp_sync:
+            batch = self.prepare_mlp_sync_batch(batch)
+
+        if batch:
+            attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
+            trace_event_batch("schedule", batch.reqs, attrs=attrs)
+
+        return batch
+
     @torch.no_grad()
     def event_loop_normal_disagg_prefill(self: Scheduler) -> None:
         """A normal scheduler loop for prefill worker in disaggregation mode."""
@@ -320,26 +351,17 @@ class SchedulerDisaggregationPrefillMixin:
             self.waiting_queue.extend(
                 self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
             )
-            self.process_prefill_chunk()
+            batch = self.get_next_disagg_prefill_batch_to_run()
-            batch = self.get_new_batch_prefill()
-            if batch:
-                attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
-                trace_event_batch("schedule", batch.reqs, attrs=attrs)
-
-            if self.require_mlp_sync:
-                batch = self.prepare_mlp_sync_batch(batch)
             self.cur_batch = batch
 
             if batch:
                 result = self.run_batch(batch)
                 self.process_batch_result_disagg_prefill(batch, result)
-
+            else:
-            if len(self.disagg_prefill_inflight_queue) > 0:
-                self.process_disagg_prefill_inflight_queue()
-
-            if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
                 self.self_check_during_idle()
 
+            self.process_disagg_prefill_inflight_queue()
+
             self.last_batch = batch
             # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
             # Otherwise, it hangs under high concurrency
@@ -355,14 +377,7 @@ class SchedulerDisaggregationPrefillMixin:
             self.waiting_queue.extend(
                 self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
             )
-            self.process_prefill_chunk()
+            batch = self.get_next_disagg_prefill_batch_to_run()
-            batch = self.get_new_batch_prefill()
-            if batch:
-                attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
-                trace_event_batch("schedule", batch.reqs, attrs=attrs)
-
-            if self.require_mlp_sync:
-                batch = self.prepare_mlp_sync_batch(batch)
             self.cur_batch = batch
 
             batch_result = None
@@ -373,15 +388,13 @@ class SchedulerDisaggregationPrefillMixin:
             if self.last_batch:
                 tmp_batch, tmp_result = self.result_queue.popleft()
                 self.process_batch_result_disagg_prefill(tmp_batch, tmp_result)
+            elif batch is None:
+                self.self_check_during_idle()
 
-            if len(self.disagg_prefill_inflight_queue) > 0:
             self.process_disagg_prefill_inflight_queue()
 
             self.launch_batch_sample_if_needed(batch_result)
 
-            if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
-                self.self_check_during_idle()
-
             self.last_batch = batch
             # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
             # Otherwise, it hangs under high concurrency
@@ -599,25 +612,6 @@ class SchedulerDisaggregationPrefillMixin:
 
         return transferred_rids
 
-    def process_prefill_chunk(self: Scheduler) -> None:
-        if self.last_batch and self.last_batch.forward_mode.is_extend():
-            if self.chunked_req:
-                # Move the chunked request out of the batch so that we can merge
-                # only finished requests to running_batch.
-                self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
-                self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
-                if self.enable_overlap:
-                    # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
-                    self.chunked_req.tmp_end_idx = min(
-                        len(self.chunked_req.fill_ids),
-                        len(self.chunked_req.origin_input_ids),
-                    )
-                else:
-                    self.send_kv_chunk(self.chunked_req)
-                # chunked request keeps its rid but will get a new req_pool_idx
-                self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
-                self.running_batch.batch_is_full = False
-
     def send_kv_chunk(
         self: Scheduler,
         req: Req,

python/sglang/srt/managers/scheduler.py

@@ -2443,14 +2443,12 @@ class Scheduler(
             for decode_req in self.disagg_decode_prealloc_queue.queue:
                 if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
                     logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}")
-                    if hasattr(decode_req.kv_receiver, "abort"):
                     decode_req.kv_receiver.abort()
 
             # Abort requests waiting for kvcache to release tree cache
             for decode_req in self.disagg_decode_transfer_queue.queue:
                 if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
                     logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}")
-                    if hasattr(decode_req.kv_receiver, "abort"):
                     decode_req.kv_receiver.abort()
 
         # Delete requests in the running batch

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -217,7 +217,10 @@ class SchedulerRuntimeCheckerMixin:
             self.tree_cache.sanity_check()
 
     def self_check_during_idle(self: Scheduler):
-        if self.disaggregation_mode == DisaggregationMode.DECODE:
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            if len(self.disagg_prefill_inflight_queue) > 0:
+                return
+        elif self.disaggregation_mode == DisaggregationMode.DECODE:
             queue_size = (
                 len(self.waiting_queue)
                 + len(self.disagg_decode_transfer_queue.queue)