[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>
This commit is contained in:
liuzhenwei
2026-03-11 13:19:28 +08:00
committed by GitHub
parent 76c6e6da08
commit f22d6e0267
2 changed files with 12 additions and 4 deletions

View File

@@ -24,9 +24,9 @@ class KVTransferConfig:
engine_id: str | None = None
"""The engine id for KV transfers."""
kv_buffer_device: str = "cuda"
"""The device used by kv connector to buffer the KV cache. Choices are
'cuda' and 'cpu'."""
kv_buffer_device: str | None = None
"""The device used by kv connector to buffer the KV cache. Choices are
'cuda','cpu' and 'xpu'."""
kv_buffer_size: float = 1e9
"""The buffer size for TorchDistributedConnector. Measured in number of
@@ -100,6 +100,11 @@ class KVTransferConfig:
f"is set, supported roles are {get_args(KVRole)}"
)
if self.kv_buffer_device is None:
from vllm.platforms import current_platform
self.kv_buffer_device = current_platform.device_type
@property
def is_kv_transfer_instance(self) -> bool:
return self.kv_connector is not None and self.kv_role in get_args(KVRole)

View File

@@ -998,7 +998,10 @@ class NixlConnectorWorker:
# KV Caches and nixl tracking data.
self.device_type = current_platform.device_type
self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device
kv_buffer_device = vllm_config.kv_transfer_config.kv_buffer_device
if kv_buffer_device is None:
raise ValueError("kv_buffer_device must be set for NixlConnector")
self.kv_buffer_device: str = kv_buffer_device
if self.device_type not in _NIXL_SUPPORTED_DEVICE:
raise RuntimeError(f"{self.device_type} is not supported.")
elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]: