From afd7b1dce94fed484351fafd5bf5ea6601ac621e Mon Sep 17 00:00:00 2001 From: Revital Sur Date: Tue, 19 May 2026 06:12:04 +0300 Subject: [PATCH] [Bugfix] Use platform-agnostic device in example_connector load (#42926) Signed-off-by: Revital Sur Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../kv_transfer/kv_connector/v1/example_connector.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py index a0279120ef0..8ed18894e8a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py @@ -188,7 +188,9 @@ class ExampleConnector(KVConnectorBase_V1): filename = self._generate_filename_debug( layer_name, request.token_ids, request.mm_hashes ) - kv_cache = safetensors.torch.load_file(filename)["kv_cache"].cuda() + kv_cache = safetensors.torch.load_file( + filename, device=str(kv_cache_layer.device) + )["kv_cache"] if isinstance(attn_metadata, dict): inject_kv_into_layer( kv_cache_layer,