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 13, 09:30 UTC
add cpu option for p/d in nixl_connector (#28356)
Signed-off-by: Hongming Zheng <hongming.zheng@intel.com>
ZhengHongming888 committed
Nov 12, 2025 at 03:53 UTC
c5f10cc139ec87e217f2bb56a677dd57394729f5
1 file changed
+16
-2
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+16
-2
@@ -91,6 +91,7 @@ _NIXL_SUPPORTED_DEVICE = {
91
),
92
"tpu": ("cpu",),
93
"xpu": ("cpu",),
94
+ "cpu": ("cpu",),
95
}
96
# support for oot platform by providing mapping in current_platform
97
_NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices())
@@ -348,7 +349,13 @@ class NixlConnectorScheduler:
349
+ vllm_config.parallel_config.data_parallel_rank
350
)
351
assert vllm_config.kv_transfer_config is not None
351
- self.use_host_buffer = vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
352
+ if current_platform.device_type == "cpu":
353
+ self.use_host_buffer = False
354
+ else:
355
+ self.use_host_buffer = (
356
+ vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
357
+ )
358
+
359
logger.info("Initializing NIXL Scheduler %s", engine_id)
360
361
# Background thread for handling new handshake requests.
@@ -820,7 +827,11 @@ class NixlConnectorWorker:
827
# cpu kv buffer for xfer
828
# used when device memory can not be registered under nixl
829
self.host_xfer_buffers: dict[str, torch.Tensor] = {}
823
- self.use_host_buffer = self.kv_buffer_device == "cpu"
830
+ if self.device_type == "cpu":
831
+ self.use_host_buffer = False
832
+ else:
833
+ self.use_host_buffer = self.kv_buffer_device == "cpu"
834
+
835
# support for oot platform which can't register nixl memory
836
# type based on kv_buffer_device
837
nixl_memory_type = current_platform.get_nixl_memory_type()
@@ -1021,6 +1032,9 @@ class NixlConnectorWorker:
1032
# Set a no-op if the host buffer is not cpu.
1033
if self.kv_buffer_device != "cpu":
1034
return
1035
+ # Set a no-op if self.device_type is 'cpu'.
1036
+ if self.device_type == "cpu":
1037
+ return
1038
assert self.use_host_buffer
1039
self.copy_blocks = copy_operation
1040