xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-13 19:34:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

add cpu option for p/d in nixl_connector (#28356)

Browse Source 
Signed-off-by: Hongming Zheng <hongming.zheng@intel.com>
This commit is contained in:
ZhengHongming888 2025-11-12 03:53:08 -08:00 committed by GitHub
parent d143152308
commit c5f10cc139
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 16 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -91,6 +91,7 @@ _NIXL_SUPPORTED_DEVICE = {
), ),
"tpu": ("cpu",), "tpu": ("cpu",),
"xpu": ("cpu",), "xpu": ("cpu",),
"cpu": ("cpu",),
} }
# support for oot platform by providing mapping in current_platform # support for oot platform by providing mapping in current_platform
_NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices()) _NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices())
@ -348,7 +349,13 @@ class NixlConnectorScheduler:
+ vllm_config.parallel_config.data_parallel_rank + vllm_config.parallel_config.data_parallel_rank
) )
assert vllm_config.kv_transfer_config is not None assert vllm_config.kv_transfer_config is not None
self.use_host_buffer = vllm_config.kv_transfer_config.kv_buffer_device == "cpu" if current_platform.device_type == "cpu":
self.use_host_buffer = False
else:
self.use_host_buffer = (
vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
)
logger.info("Initializing NIXL Scheduler %s", engine_id) logger.info("Initializing NIXL Scheduler %s", engine_id)
# Background thread for handling new handshake requests. # Background thread for handling new handshake requests.
@ -820,7 +827,11 @@ class NixlConnectorWorker:
# cpu kv buffer for xfer # cpu kv buffer for xfer
# used when device memory can not be registered under nixl # used when device memory can not be registered under nixl
self.host_xfer_buffers: dict[str, torch.Tensor] = {} self.host_xfer_buffers: dict[str, torch.Tensor] = {}
if self.device_type == "cpu":
self.use_host_buffer = False
else:
self.use_host_buffer = self.kv_buffer_device == "cpu" self.use_host_buffer = self.kv_buffer_device == "cpu"
# support for oot platform which can't register nixl memory # support for oot platform which can't register nixl memory
# type based on kv_buffer_device # type based on kv_buffer_device
nixl_memory_type = current_platform.get_nixl_memory_type() nixl_memory_type = current_platform.get_nixl_memory_type()
@ -1021,6 +1032,9 @@ class NixlConnectorWorker:
# Set a no-op if the host buffer is not cpu. # Set a no-op if the host buffer is not cpu.
if self.kv_buffer_device != "cpu": if self.kv_buffer_device != "cpu":
return return
# Set a no-op if self.device_type is 'cpu'.
if self.device_type == "cpu":
return
assert self.use_host_buffer assert self.use_host_buffer
self.copy_blocks = copy_operation self.copy_blocks = copy_operation