A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Update num_tokens_across_dp to use nccl instead of gloo (#24105)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Sep 15, 2025 at 12:05 UTC
49bfc538e4047fce58d3a43b975aef0bab795b25
2 files changed
+25
-4
vllm/envs.py
+8
@@ -92,6 +92,7 @@ if TYPE_CHECKING:
92
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
93
VLLM_SKIP_P2P_CHECK: bool = False
94
VLLM_DISABLED_KERNELS: list[str] = []
95
+ VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION: bool = False
96
VLLM_USE_V1: bool = True
97
VLLM_ROCM_USE_AITER: bool = False
98
VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
@@ -745,6 +746,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
746
lambda: [] if "VLLM_DISABLED_KERNELS" not in os.environ else os.environ[
747
"VLLM_DISABLED_KERNELS"].split(","),
748
749
+ # Swaps the all reduce backend that we use to coordinate the DP padding
750
+ # information from NCCL to gloo.
751
+ "VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION":
752
+ lambda:
753
+ (os.getenv("VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION", "False").lower() in
754
+ ("true", "1")),
755
+
756
# If set, use the V1 code path.
757
"VLLM_USE_V1":
758
lambda: bool(int(os.getenv("VLLM_USE_V1", "1"))),
vllm/forward_context.py
+17
-4
@@ -13,6 +13,7 @@ import torch.distributed as dist
13
import vllm.envs as envs
14
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
15
from vllm.logger import init_logger
16
+from vllm.platforms import current_platform
17
18
if TYPE_CHECKING:
19
from vllm.attention.backends.abstract import AttentionMetadata
@@ -75,14 +76,26 @@ class DPMetadata:
76
Gather the num_tokens across all DP ranks and return results in a
77
CPU tensor of size dp_size.
78
"""
79
+ from vllm.distributed.parallel_state import get_dp_group
80
+ device = current_platform.device_type
81
+ group = get_dp_group().device_group
82
+
83
+ # Transfering this tensor from GPU to CPU will introduce a GPU sync
84
+ # point that could adversely affect performance of vllm with asynch
85
+ # scheduling. This environment variable exists to quickly disable
86
+ # this optimization if we run into this case.
87
+ if envs.VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION:
88
+ logger.info_once(
89
+ "Using CPU all reduce to syncronize DP padding between ranks.")
90
+ device = "cpu"
91
+ group = get_dp_group().cpu_group
92
num_tokens_across_dp = [0] * dp_size
93
num_tokens_across_dp[dp_rank] = num_tokens
94
num_tokens_tensor = torch.tensor(num_tokens_across_dp,
81
- device="cpu",
95
+ device=device,
96
dtype=torch.int32)
83
- from vllm.distributed.parallel_state import get_dp_group
84
- dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)
85
- return num_tokens_tensor
97
+ dist.all_reduce(num_tokens_tensor, group=group)
98
+ return num_tokens_tensor.cpu()
99
100
@staticmethod
101
def make(