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 6, 21:45 UTC
[Hardware][NIXL] set default kv buffer type for different platform (#36438)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
liuzhenwei committed
Mar 11, 2026 at 13:19 UTC
f22d6e026798a74e6542a52ef776c054f2de572a
2 files changed
+12
-4
vllm/config/kv_transfer.py
+8
-3
@@ -24,9 +24,9 @@ class KVTransferConfig:
24
engine_id: str | None = None
25
"""The engine id for KV transfers."""
26
27
- kv_buffer_device: str = "cuda"
28
- """The device used by kv connector to buffer the KV cache. Choices are
29
- 'cuda' and 'cpu'."""
27
+ kv_buffer_device: str | None = None
28
+ """The device used by kv connector to buffer the KV cache. Choices are
29
+ 'cuda','cpu' and 'xpu'."""
30
31
kv_buffer_size: float = 1e9
32
"""The buffer size for TorchDistributedConnector. Measured in number of
@@ -100,6 +100,11 @@ class KVTransferConfig:
100
f"is set, supported roles are {get_args(KVRole)}"
101
)
102
103
+ if self.kv_buffer_device is None:
104
+ from vllm.platforms import current_platform
105
+
106
+ self.kv_buffer_device = current_platform.device_type
107
+
108
@property
109
def is_kv_transfer_instance(self) -> bool:
110
return self.kv_connector is not None and self.kv_role in get_args(KVRole)
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+4
-1
@@ -998,7 +998,10 @@ class NixlConnectorWorker:
998
999
# KV Caches and nixl tracking data.
1000
self.device_type = current_platform.device_type
1001
- self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device
1001
+ kv_buffer_device = vllm_config.kv_transfer_config.kv_buffer_device
1002
+ if kv_buffer_device is None:
1003
+ raise ValueError("kv_buffer_device must be set for NixlConnector")
1004
+ self.kv_buffer_device: str = kv_buffer_device
1005
if self.device_type not in _NIXL_SUPPORTED_DEVICE:
1006
raise RuntimeError(f"{self.device_type} is not supported.")
1007
elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]: