diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a11df2205..fcacdf865 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1326,12 +1326,21 @@ class Scheduler( def is_disable_overlap_for_batch(self, batch: ScheduleBatch) -> bool: # For two consecutive prefill batches, we disable overlap to improve the TTFT of the first batch. # This might slightly hurt the throughput, so we use an environment variable to control it. + # In DP attention mode, use the globally synchronized is_extend_in_batch + # so all DP ranks make the same overlap decision (avoiding deadlock). + # In non-DP mode, use the local forward_mode directly. + if self.require_mlp_sync: + is_extend = lambda b: b and b.is_extend_in_batch + else: + is_extend = lambda b: b and b.forward_mode.is_extend() + + batch_is_extend = is_extend(batch) + last_batch_is_extend = is_extend(self.last_batch) + disable_overlap_for_batch = ( envs.SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get() - and batch - and batch.forward_mode.is_extend() - and self.last_batch - and self.last_batch.forward_mode.is_extend() + and batch_is_extend + and last_batch_is_extend ) # We do not support overlap + spec + grammar yet, diff --git a/test/registered/distributed/test_dp_attention.py b/test/registered/distributed/test_dp_attention.py index 1824f6f71..4d9d1fb1f 100644 --- a/test/registered/distributed/test_dp_attention.py +++ b/test/registered/distributed/test_dp_attention.py @@ -38,6 +38,10 @@ class TestDPAttentionDP2TP2( def setUpClass(cls): cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST cls.base_url = DEFAULT_URL_FOR_TEST + cls._env_override = envs.SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.override( + True + ) + cls._env_override.__enter__() cls.process = popen_launch_server( cls.model, cls.base_url, @@ -58,6 +62,7 @@ class TestDPAttentionDP2TP2( @classmethod def tearDownClass(cls): kill_process_tree(cls.process.pid) + cls._env_override.__exit__(None, None, None) def test_mgsm_en(self): args = SimpleNamespace(