From 0229025127f3928353c7af3bd1583e87bacaa397 Mon Sep 17 00:00:00 2001
From: YAMY <74099316+YAMY1234@users.noreply.github.com>
Date: Mon, 21 Sep 2026 17:47:03 -0500
Subject: [PATCH] [Spec][PP] Launch extend microbatches before the spec output
 exchange (#40499)

---
 python/sglang/srt/managers/scheduler.py       |  9 ++++--
 .../sglang/srt/managers/scheduler_pp_mixin.py | 30 +++++++++++++------
 .../unit/managers/test_pp_cp_rank_offsets.py  | 21 ++++++++++++-
 3 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index e368797dd..78bac308f 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/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:
-                        batch.seq_lens = batch_result.new_seq_lens
+                    new_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)
diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py
index 2757adb81..71bca7dab 100644
--- a/python/sglang/srt/managers/scheduler_pp_mixin.py
+++ b/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 the previous outputs before launching the next batch.
-                # Tree planning synchronizes CUDA on the host; sending alone
-                # leaves the peer's return send unmatched and can block that
-                # synchronization while the peer waits for our next proxy.
-                # 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)
+                exchange_outputs_before_forward = _pp_exchange_outputs_before_forward(
+                    cur_batch=cur_batch,
+                    spec_relay=self._pp_spec_relay,
+                    is_last_rank=self.pp_group.is_last_rank,
+                    async_batch_depth=get_parallel().pp_async_batch_depth,
                 )
                 if exchange_outputs_before_forward:
                     next_pp_outputs, next_batch_result, d2h_event = (
diff --git a/test/registered/unit/managers/test_pp_cp_rank_offsets.py b/test/registered/unit/managers/test_pp_cp_rank_offsets.py
index 6cbf68877..5db7cd72b 100644
--- a/test/registered/unit/managers/test_pp_cp_rank_offsets.py
+++ b/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()