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, 00:30 UTC
[Misc] Replace CUDA_VISIBLE_DEVICES in DP with torch.cuda.set_device for device selection on cuda-like devices (#27564)
Signed-off-by: ilmarkov <markovilya197@gmail.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Ilya Markov committed
Oct 30, 2025 at 16:41 UTC
60f76baa6688ce265a4205f183bd42a62d8f7179
4 files changed
+43
-7
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+8
-4
@@ -1008,11 +1008,14 @@ class NixlConnectorWorker:
1008
# Enable different block lengths for different layers when MLA is used.
1009
self.block_len_per_layer = list[int]()
1010
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
1011
+ self.device_id = self.tp_rank
1012
for layer_name, cache_or_caches in xfer_buffers.items():
1013
cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1014
1015
for cache in cache_list:
1016
base_addr = cache.data_ptr()
1017
+ if not self.use_host_buffer and current_platform.is_cuda_alike():
1018
+ self.device_id = cache.device.index
1019
if base_addr in seen_base_addresses:
1020
continue
1021
@@ -1040,7 +1043,7 @@ class NixlConnectorWorker:
1043
"All kv cache tensors must have the same size"
1044
)
1045
caches_data.append(
1043
- (base_addr, curr_tensor_size_bytes, self.tp_rank, "")
1046
+ (base_addr, curr_tensor_size_bytes, self.device_id, "")
1047
)
1048
1049
logger.debug(
@@ -1087,7 +1090,7 @@ class NixlConnectorWorker:
1090
block_offset = block_id * self.block_len_per_layer[i]
1091
addr = base_addr + block_offset
1092
# (addr, len, device id)
1090
- blocks_data.append((addr, kv_block_len, self.tp_rank))
1093
+ blocks_data.append((addr, kv_block_len, self.device_id))
1094
1095
if self._use_flashinfer:
1096
# Separate and interleave K/V regions to maintain the same
@@ -1098,12 +1101,13 @@ class NixlConnectorWorker:
1101
addr = base_addr + block_offset
1102
# Register addresses for V cache (K registered first).
1103
v_addr = addr + kv_block_len
1101
- blocks_data.append((v_addr, kv_block_len, self.tp_rank))
1104
+ blocks_data.append((v_addr, kv_block_len, self.device_id))
1105
logger.debug(
1103
- "Created %s blocks for src engine %s and rank %s",
1106
+ "Created %s blocks for src engine %s and rank %s on device id %s",
1107
len(blocks_data),
1108
self.engine_id,
1109
self.tp_rank,
1110
+ self.device_id,
1111
)
1112
1113
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
vllm/v1/engine/utils.py
+10
-1
@@ -134,9 +134,18 @@ class CoreEngineProcManager:
134
data_parallel = vllm_config.parallel_config.data_parallel_size > 1
135
try:
136
for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
137
+ # Adjust device control in DP for non-CUDA platforms
138
+ # as well as external and ray launchers
139
+ # For CUDA platforms, we use torch.cuda.set_device()
140
with (
141
set_device_control_env_var(vllm_config, local_dp_rank)
139
- if (data_parallel)
142
+ if (
143
+ data_parallel
144
+ and (
145
+ not current_platform.is_cuda_alike()
146
+ or vllm_config.parallel_config.use_ray
147
+ )
148
+ )
149
else contextlib.nullcontext()
150
):
151
proc.start()
vllm/v1/worker/dp_utils.py
+2
-2
@@ -8,7 +8,6 @@ import torch.distributed as dist
8
from vllm.config import ParallelConfig
9
from vllm.distributed.parallel_state import get_dp_group
10
from vllm.logger import init_logger
11
-from vllm.platforms import current_platform
11
from vllm.v1.worker.ubatch_utils import (
12
UBatchSlices,
13
check_ubatch_thresholds,
@@ -20,7 +19,8 @@ logger = init_logger(__name__)
19
20
21
def _get_device_and_group(parallel_config: ParallelConfig):
23
- device = current_platform.device_type
22
+ # Use the actual device assigned to the DP group, not just the device type
23
+ device = get_dp_group().device
24
group = get_dp_group().device_group
25
26
# Transfering this tensor from GPU to CPU will introduce a GPU sync
vllm/v1/worker/gpu_worker.py
+23
@@ -172,6 +172,29 @@ class Worker(WorkerBase):
172
if self.device_config.device.type == "cuda":
173
# This env var set by Ray causes exceptions with graph building.
174
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
175
+ if (
176
+ self.parallel_config.data_parallel_size > 1
177
+ and self.parallel_config.data_parallel_size_local > 0
178
+ and self.parallel_config.distributed_executor_backend
179
+ not in ["ray", "external_launcher"]
180
+ and self.vllm_config.parallel_config.data_parallel_backend != "ray"
181
+ ):
182
+ # Use local DP rank if available, otherwise use global DP rank.
183
+ dp_local_rank = self.parallel_config.data_parallel_rank_local
184
+ if dp_local_rank is None:
185
+ dp_local_rank = self.parallel_config.data_parallel_rank
186
+
187
+ tp_pp_world_size = (
188
+ self.parallel_config.pipeline_parallel_size
189
+ * self.parallel_config.tensor_parallel_size
190
+ )
191
+
192
+ # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
193
+ self.local_rank += dp_local_rank * tp_pp_world_size
194
+ assert self.local_rank < torch.cuda.device_count(), (
195
+ f"DP adjusted local rank {self.local_rank} is out of bounds. "
196
+ )
197
+
198
self.device = torch.device(f"cuda:{self.local_rank}")
199
current_platform.set_device(self.device)
200