fix(dp-attn): consistent overlap disable decision across DP ranks (#20853)

This commit is contained in:
Liangsheng Yin
2026-03-18 15:16:39 -07:00
committed by GitHub
parent 4e8829e4cd
commit 8b9482e665
2 changed files with 18 additions and 4 deletions

View File

@@ -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,