[Spec][PP] Launch extend microbatches before the spec output exchange (#40499)

python/sglang/srt/managers/scheduler.py

@@ -4525,10 +4525,13 @@ class Scheduler(
                         if is_verify_round
                         else batch_result.next_draft_input
                     )
-                    if batch_result.new_seq_lens is not None:
+                    new_seq_lens = batch_result.new_seq_lens
-                        batch.seq_lens = batch_result.new_seq_lens
+                    # Extend rounds return batch.seq_lens itself; copying it back
+                    # would block the scheduler until the whole forward has run.
+                    if new_seq_lens is not None and new_seq_lens is not batch.seq_lens:
+                        batch.seq_lens = new_seq_lens
                         if batch.seq_lens_cpu is not None:
-                            batch.seq_lens_cpu = batch_result.new_seq_lens.to("cpu")
+                            batch.seq_lens_cpu = new_seq_lens.to("cpu")
                             batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
                     batch.input_ids = None  # rebuilt next iter from draft_token
                     self.update_cache_from_scheduler(batch, batch_result)

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -55,6 +55,22 @@ def _pp_can_skip_output_comm(batch: ScheduleBatch) -> bool:
     )
 
 
+def _pp_exchange_outputs_before_forward(
+    cur_batch: Optional[ScheduleBatch],
+    spec_relay: bool,
+    is_last_rank: bool,
+    async_batch_depth: int,
+) -> bool:
+    """Extend microbatches launch first: they need nothing from the relay, and
+    exchanging first caps every stage at (pp_size - 1) / pp_size. A verify round
+    must exchange first or the ring deadlocks on its tree rebuild."""
+    if async_batch_depth > 0:
+        return True
+    if not spec_relay or is_last_rank or cur_batch is None:
+        return False
+    return not (cur_batch.forward_mode.is_extend() or cur_batch.is_extend_in_batch)
+
+
 @dataclass
 class PPBatchMetadata:
     can_run_cuda_graph: bool
@@ -131,15 +147,11 @@ class SchedulerPPMixin:
                 next_pp_outputs = None
                 next_batch_result = None
                 d2h_event = None
-                # With zero async depth, non-last speculative ranks must
+                exchange_outputs_before_forward = _pp_exchange_outputs_before_forward(
-                # exchange the previous outputs before launching the next batch.
+                    cur_batch=cur_batch,
-                # Tree planning synchronizes CUDA on the host; sending alone
+                    spec_relay=self._pp_spec_relay,
-                # leaves the peer's return send unmatched and can block that
+                    is_last_rank=self.pp_group.is_last_rank,
-                # synchronization while the peer waits for our next proxy.
+                    async_batch_depth=get_parallel().pp_async_batch_depth,
-                # The last rank must launch first to produce its output.
-                exchange_outputs_before_forward = (
-                    get_parallel().pp_async_batch_depth > 0
-                    or (self._pp_spec_relay and not self.pp_group.is_last_rank)
                 )
                 if exchange_outputs_before_forward:
                     next_pp_outputs, next_batch_result, d2h_event = (

test/registered/unit/managers/test_pp_cp_rank_offsets.py

@@ -17,7 +17,11 @@ maybe_stub_sgl_kernel()
 from sglang.srt.managers.scheduler_components.request_receiver import (  # noqa: E402
     SchedulerRequestReceiver,
 )
-from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin  # noqa: E402
+from sglang.srt.managers.scheduler_pp_mixin import (  # noqa: E402
+    SchedulerPPMixin,
+    _pp_exchange_outputs_before_forward,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardMode  # noqa: E402
 
 register_cpu_ci(est_time=11, suite="base-a-test-cpu")
 
@@ -265,5 +269,20 @@ class TestDSparkPPOutput(CustomTestCase):
         self.assertEqual(payloads[0].hidden_states.numel(), 0)
 
 
+class TestPPSpecExchangeOrder(unittest.TestCase):
+    def test_extend_launches_before_the_relay_exchange(self):
+        kwargs = dict(spec_relay=True, is_last_rank=False, async_batch_depth=0)
+        extend = SimpleNamespace(
+            forward_mode=ForwardMode.EXTEND, is_extend_in_batch=False
+        )
+        decode = SimpleNamespace(
+            forward_mode=ForwardMode.DECODE, is_extend_in_batch=False
+        )
+        self.assertFalse(
+            _pp_exchange_outputs_before_forward(cur_batch=extend, **kwargs)
+        )
+        self.assertTrue(_pp_exchange_outputs_before_forward(cur_batch=decode, **kwargs))
+
+
 if __name__ == "__main__":
     unittest.main()