From 12f2c515a7def384d2f775ff28ff2089947d06c2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Thu, 16 Jul 2026 09:04:39 +0200 Subject: [PATCH] [Bugfix] Fix offloading set_ overflow for packed non-uniform KV caches (#48530) Signed-off-by: Elvir Crncevic --- .../kv_connector/v1/offloading/worker.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 60e0beb6845..260746be051 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -92,16 +92,17 @@ class OffloadingConnectorWorker: if layer_is_packed[layer_name] else page ) + raw = torch.empty( + 0, + dtype=torch.int8, + device=layer_kv_cache.device, + ).set_(layer_kv_cache.untyped_storage()) tensors_per_block[layer_name] = ( - torch.tensor( - [], - dtype=torch.int8, - device=layer_kv_cache.device, - ).set_( - layer_kv_cache.untyped_storage(), - byte_offset, + torch.as_strided( + raw, (num_blocks, page), (block_stride_bytes, 1), + byte_offset, ), ) page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes