Add flag that enables NCCL mlp sync batch for overlap scheduler (#17288)

2026-01-20 23:06:55 +08:00
parent 91a4cd8648
commit 55c616427d
3 changed files with 6 additions and 1 deletions
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -246,6 +246,7 @@ class Envs:
    SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK = EnvFloat(None)
    SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1)
    SGLANG_QUEUED_TIMEOUT_MS = EnvInt(-1)
+    SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH = EnvBool(False)

    # Test: pd-disaggregation
    SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake")
--- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py
+++ b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py
@@ -170,7 +170,10 @@ def prepare_mlp_sync_batch_raw(
        local_batch.is_extend_in_batch = is_extend_in_batch

    tbo_preparer = TboDPAttentionPreparer()
-    if len(offload_tags) == 0 and disable_overlap_schedule:
+    if len(offload_tags) == 0 and (
+        disable_overlap_schedule
+        or envs.SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH.get()
+    ):
        group = tp_group.device_group
        device = tp_group.device
    else: