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
Revert "[Misc] Remove use of CUDA_VISIBLE_DEVICES for device selectio… (#27502)
Zhuohan Li committed
Oct 24, 2025 at 22:31 UTC
56ed7609a9d61ac43e89d62e0e9494ceea9a450d
4 files changed
+7
-35
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+4
-8
@@ -991,14 +991,11 @@ class NixlConnectorWorker:
991
# Enable different block lengths for different layers when MLA is used.
992
self.block_len_per_layer = list[int]()
993
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
994
- self.device_id = self.tp_rank
994
for layer_name, cache_or_caches in xfer_buffers.items():
995
cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
996
997
for cache in cache_list:
998
base_addr = cache.data_ptr()
1000
- if not self.use_host_buffer and current_platform.is_cuda_alike():
1001
- self.device_id = cache.device.index
999
if base_addr in seen_base_addresses:
1000
continue
1001
@@ -1026,7 +1023,7 @@ class NixlConnectorWorker:
1023
"All kv cache tensors must have the same size"
1024
)
1025
caches_data.append(
1029
- (base_addr, curr_tensor_size_bytes, self.device_id, "")
1026
+ (base_addr, curr_tensor_size_bytes, self.tp_rank, "")
1027
)
1028
1029
logger.debug(
@@ -1073,7 +1070,7 @@ class NixlConnectorWorker:
1070
block_offset = block_id * self.block_len_per_layer[i]
1071
addr = base_addr + block_offset
1072
# (addr, len, device id)
1076
- blocks_data.append((addr, kv_block_len, self.device_id))
1073
+ blocks_data.append((addr, kv_block_len, self.tp_rank))
1074
1075
if self._use_flashinfer:
1076
# Separate and interleave K/V regions to maintain the same
@@ -1084,13 +1081,12 @@ class NixlConnectorWorker:
1081
addr = base_addr + block_offset
1082
# Register addresses for V cache (K registered first).
1083
v_addr = addr + kv_block_len
1087
- blocks_data.append((v_addr, kv_block_len, self.device_id))
1084
+ blocks_data.append((v_addr, kv_block_len, self.tp_rank))
1085
logger.debug(
1089
- "Created %s blocks for src engine %s and rank %s on device id %s",
1086
+ "Created %s blocks for src engine %s and rank %s",
1087
len(blocks_data),
1088
self.engine_id,
1089
self.tp_rank,
1093
- self.device_id,
1090
)
1091
1092
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
vllm/v1/engine/utils.py
+1
-4
@@ -134,12 +134,9 @@ 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
- # For CUDA platforms, setting same device id for different DP
139
- # processes affects NCCL init performance.
137
with (
138
set_device_control_env_var(vllm_config, local_dp_rank)
142
- if (data_parallel and not current_platform.is_cuda_alike())
139
+ if (data_parallel)
140
else contextlib.nullcontext()
141
):
142
proc.start()
vllm/v1/worker/dp_utils.py
+2
-2
@@ -8,6 +8,7 @@ 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
12
from vllm.v1.worker.ubatch_utils import (
13
UBatchSlices,
14
check_ubatch_thresholds,
@@ -19,8 +20,7 @@ logger = init_logger(__name__)
20
21
22
def _get_device_and_group(parallel_config: ParallelConfig):
22
- # Use the actual device assigned to the DP group, not just the device type
23
- device = get_dp_group().device
23
+ device = current_platform.device_type
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
-21
@@ -172,27 +172,6 @@ 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.data_parallel_backend != "ray"
179
- ):
180
- # Use local DP rank if available, otherwise use global DP rank.
181
- dp_local_rank = self.parallel_config.data_parallel_rank_local
182
- if dp_local_rank is None:
183
- dp_local_rank = self.parallel_config.data_parallel_rank
184
-
185
- tp_pp_world_size = (
186
- self.parallel_config.pipeline_parallel_size
187
- * self.parallel_config.tensor_parallel_size
188
- )
189
-
190
- # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
191
- self.local_rank += dp_local_rank * tp_pp_world_size
192
- assert self.local_rank <= torch.cuda.device_count(), (
193
- f"DP adjusted local rank {self.local_rank} is out of bounds. "
194
- )
195
-
175
self.device = torch.device(f"cuda:{self.local_rank}")
176
current_platform.set_device(self.device)
177