[Bugfix] Fix offloading set_ overflow for packed non-uniform KV caches (#48530)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
Elvir Crnčević committed
Jul 16, 2026 at 09:04 UTC
12f2c515a7def384d2f775ff28ff2089947d06c2
1 file changed
+8
-7
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+8
-7
@@ -92,16 +92,17 @@ class OffloadingConnectorWorker:
92
if layer_is_packed[layer_name]
93
else page
94
)
95
+ raw = torch.empty(
96
+ 0,
97
+ dtype=torch.int8,
98
+ device=layer_kv_cache.device,
99
+ ).set_(layer_kv_cache.untyped_storage())
100
tensors_per_block[layer_name] = (
96
- torch.tensor(
97
- [],
98
- dtype=torch.int8,
99
- device=layer_kv_cache.device,
100
- ).set_(
101
- layer_kv_cache.untyped_storage(),
102
- byte_offset,
101
+ torch.as_strided(
102
+ raw,
103
(num_blocks, page),
104
(block_stride_bytes, 1),
105
+ byte_offset,
106
),
107
)
108
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes