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, 19:15 UTC
[Misc] Set default `kv_buffer_device` in a better way (#36862)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 13, 2026 at 19:07 UTC
5a3f1eb62fb8a5d114001488832f8bd7f93df5b8
2 files changed
+9
-11
vllm/config/kv_transfer.py
+8
-7
@@ -13,6 +13,12 @@ KVConsumer = Literal["kv_consumer", "kv_both"]
13
KVRole = Literal[KVProducer, KVConsumer]
14
15
16
+def kv_buffer_device_default_factory() -> str:
17
+ from vllm.platforms import current_platform
18
+
19
+ return current_platform.device_type
20
+
21
+
22
@config
23
class KVTransferConfig:
24
"""Configuration for distributed KV cache transfer."""
@@ -24,9 +30,9 @@ class KVTransferConfig:
30
engine_id: str | None = None
31
"""The engine id for KV transfers."""
32
27
- kv_buffer_device: str | None = None
33
+ kv_buffer_device: str = field(default_factory=kv_buffer_device_default_factory)
34
"""The device used by kv connector to buffer the KV cache. Choices are
29
- 'cuda','cpu' and 'xpu'."""
35
+ 'cuda', 'cpu' and 'xpu'."""
36
37
kv_buffer_size: float = 1e9
38
"""The buffer size for TorchDistributedConnector. Measured in number of
@@ -100,11 +106,6 @@ class KVTransferConfig:
106
f"is set, supported roles are {get_args(KVRole)}"
107
)
108
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
-
109
@property
110
def is_kv_transfer_instance(self) -> bool:
111
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
+1
-4
@@ -998,10 +998,7 @@ class NixlConnectorWorker:
998
999
# KV Caches and nixl tracking data.
1000
self.device_type = current_platform.device_type
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
1001
+ self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device
1002
if self.device_type not in _NIXL_SUPPORTED_DEVICE:
1003
raise RuntimeError(f"{self.device_type} is not supported.")
1004
elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]: