From 55c616427d12f9c7151b08a68e37c0352c03c69b Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Tue, 20 Jan 2026 23:06:55 +0800 Subject: [PATCH] Add flag that enables NCCL mlp sync batch for overlap scheduler (#17288) --- docs/references/environment_variables.md | 1 + python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/scheduler_dp_attn_mixin.py | 5 ++++- 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 0dce9ff34..ac8a41dbe 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -40,6 +40,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_NONE` | Weight increment when forward mode is None in scheduler recv skipper. Works with `--scheduler-recv-interval` to control polling frequency when no specific forward mode is active. | `1` | | `SGLANG_MM_BUFFER_SIZE_MB` | Size of preallocated GPU buffer (in MB) for multi-modal feature hashing optimization. When set to a positive value, temporarily moves features to GPU for faster hash computation, then moves them back to CPU to save GPU memory. Larger features benefit more from GPU hashing. Set to `0` to disable. | `0` | | `SGLANG_MM_PRECOMPUTE_HASH` | Enable precomputing of hash values for MultimodalDataItem | `false` | +| `SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH` | Enable NCCL for gathering when preparing mlp sync batch under overlap scheduler (without this flag gloo is used for gathering) | `false` | ## DeepGEMM Configuration (Advanced Optimization) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5adec318c..5b66db887 100644 --- 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") diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index 65ac39da7..3a772d035 100644 --- 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: