Cleanup disagg decode prebuilt flow and add cross-stream sync in merge_batch (#19568)

Co-authored-by: vincent <vincent@vincentdeMacBook-Pro.local>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>

python/sglang/srt/disaggregation/decode.py

@@ -65,7 +65,6 @@ from sglang.srt.observability.req_time_stats import (
     set_schedule_time_batch,
     set_time_batch,
 )
-from sglang.srt.utils import get_int_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
 logger = logging.getLogger(__name__)
@@ -75,7 +74,7 @@ if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import Scheduler
     from sglang.srt.server_args import ServerArgs
 
-CLIP_MAX_NEW_TOKEN = get_int_env_var("SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION", 4096)
+CLIP_MAX_NEW_TOKEN = envs.SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION.get()
 
 
 def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool:
@@ -1036,39 +1035,27 @@ class SchedulerDisaggregationDecodeMixin:
     def get_next_disagg_decode_batch_to_run(
         self: Scheduler,
     ) -> Optional[ScheduleBatch]:
-        """Create fake completed prefill if possible and merge with running batch"""
-        # Merge the prefill batch into the running batch
-        last_batch = self.last_batch
-        if last_batch and last_batch.forward_mode.is_prebuilt():
-            # chunked prefill doesn't happen in decode instance.
-            assert self.chunked_req is None
-            # Filter finished batches.
-            last_batch.filter_batch()
-            if not last_batch.is_empty():
-                if self.running_batch.is_empty():
-                    self.running_batch = last_batch
-                else:
-                    # merge running_batch with prefill batch
-                    self.running_batch.merge_batch(last_batch)
-
+        """Process prebuilt batch and schedule the next decode batch."""
+        # Process pending prebuilt batch: output processing + filter + merge
         new_prebuilt_batch = self.get_new_prebuilt_batch()
-
-        ret: Optional[ScheduleBatch] = None
         if new_prebuilt_batch:
-            ret = new_prebuilt_batch
+            assert self.chunked_req is None
+            self.process_batch_result_prebuilt(new_prebuilt_batch)
+            new_prebuilt_batch.filter_batch()
+            if not new_prebuilt_batch.is_empty():
+                if self.running_batch.is_empty():
+                    self.running_batch = new_prebuilt_batch
+                else:
+                    self.running_batch.merge_batch(new_prebuilt_batch)
+
+        # Schedule decode batch
+        if self.running_batch.is_empty():
+            ret = None
         else:
-            if self.running_batch.is_empty():
-                ret = None
-            else:
-                self.running_batch = self.update_running_batch(self.running_batch)
-                ret = self.running_batch if not self.running_batch.is_empty() else None
+            self.running_batch = self.update_running_batch(self.running_batch)
+            ret = self.running_batch if not self.running_batch.is_empty() else None
 
-        # 1. decode + None -> decode + idle
-        # 2. decode + prebuilt -> decode + idle (idle forward, prebuilt returns)
-        # 3. prebuilt + None -> None (None forward, prebuilt returns) + None
-        # 4. prebuilt + decode + None -> idle (idle forward, prebuilt returns) + decode + idle
         ret = self.maybe_prepare_mlp_sync_batch(ret)
-
         if ret:
             set_schedule_time_batch(ret)
         return ret

python/sglang/srt/managers/schedule_batch.py

@@ -2088,9 +2088,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             )
 
     def merge_batch(self, other: "ScheduleBatch"):
-        # NOTE: in spec v2 mode, we do not need wait verify here because
-        # 1) current batch is always prefill, whose seq_lens is not a future
-        # 2) other batch is always decode, which is finished in previous step
+        # In the regular scheduler path:
+        # 1) self is always prefill, whose seq_lens is not a future
+        # 2) other is always decode, which is finished in previous step
+        # so verify_done is already synced and this is a no-op.
+        # In disagg decode + overlap, merge_batch can be called before
+        # filter_batch, so running_batch.seq_lens may still be a forward_stream
+        # future. Synchronize here to avoid a cross-stream data race.
+        self.maybe_wait_verify_done()
 
         # Penalizer orchestrator must be merged before Batch.reqs is merged. This is because
         # orchestrator.merge() depends on Batch.reqs during preparation of each penalizers, so it
@@ -2213,9 +2218,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         )
 
     def copy(self):
-        # Only contain fields that will be used by process_batch_result
+        # Only contain fields that will be used by process_batch_result.
+        # Shallow-copy the reqs list so that in-place mutations (filter_batch,
+        # merge_batch) on the original don't corrupt this snapshot.
         return ScheduleBatch(
-            reqs=self.reqs,
+            reqs=self.reqs[:],
             req_to_token_pool=self.req_to_token_pool,
             req_pool_indices=self.req_pool_indices,
             model_config=self.model_config,