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 3, 09:15 UTC
[Bugfix] Use platform-agnostic device in example_connector load (#42926)
Signed-off-by: Revital Sur <eres@il.ibm.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Revital Sur committed
May 19, 2026 at 06:12 UTC
afd7b1dce94fed484351fafd5bf5ea6601ac621e
1 file changed
+3
-1
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+3
-1
@@ -188,7 +188,9 @@ class ExampleConnector(KVConnectorBase_V1):
188
filename = self._generate_filename_debug(
189
layer_name, request.token_ids, request.mm_hashes
190
)
191
- kv_cache = safetensors.torch.load_file(filename)["kv_cache"].cuda()
191
+ kv_cache = safetensors.torch.load_file(
192
+ filename, device=str(kv_cache_layer.device)
193
+ )["kv_cache"]
194
if isinstance(attn_metadata, dict):
195
inject_kv_into_layer(
196
kv_cache_layer,