biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Misc] Remove use of CUDA_VISIBLE_DEVICES for device selection (fix DP slow startup time &c) (#26709)

Browse Source 
Signed-off-by: ilmarkov <markovilya197@gmail.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
This commit is contained in:
Ilya Markov
2025-10-23 14:58:39 +02:00
committed by GitHub
parent faee3ccdc2
commit 237cf6d32a
4 changed files with 35 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
View File

@@ -991,11 +991,14 @@ class NixlConnectorWorker:
# Enable different block lengths for different layers when MLA is used.
self.block_len_per_layer = list[int]()
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
self.device_id = self.tp_rank
for layer_name, cache_or_caches in xfer_buffers.items():
cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
for cache in cache_list:
base_addr = cache.data_ptr()
if not self.use_host_buffer and current_platform.is_cuda_alike():
self.device_id = cache.device.index
if base_addr in seen_base_addresses:
continue
@@ -1023,7 +1026,7 @@ class NixlConnectorWorker:
"All kv cache tensors must have the same size"
)
caches_data.append(
(base_addr, curr_tensor_size_bytes, self.tp_rank, "")
(base_addr, curr_tensor_size_bytes, self.device_id, "")
)
logger.debug(
@@ -1070,7 +1073,7 @@ class NixlConnectorWorker:
block_offset = block_id * self.block_len_per_layer[i]
addr = base_addr + block_offset
# (addr, len, device id)
blocks_data.append((addr, kv_block_len, self.tp_rank))
blocks_data.append((addr, kv_block_len, self.device_id))
if self._use_flashinfer:
# Separate and interleave K/V regions to maintain the same
@@ -1081,12 +1084,13 @@ class NixlConnectorWorker:
addr = base_addr + block_offset
# Register addresses for V cache (K registered first).
v_addr = addr + kv_block_len
blocks_data.append((v_addr, kv_block_len, self.tp_rank))
blocks_data.append((v_addr, kv_block_len, self.device_id))
logger.debug(
"Created %s blocks for src engine %s and rank %s",
"Created %s blocks for src engine %s and rank %s on device id %s",
len(blocks_data),
self.engine_id,
self.tp_rank,
self.device_id,
)
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)