[KV-Offloading] Fix tensors_per_block stride (#46888)
Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
Varun Sundar Rabindranath committed
Jun 27, 2026 at 21:01 UTC
798185d438c030f9b4fd62687440889e7b195251
2 files changed
+26
-2
tests/v1/kv_connector/unit/offloading_connector/utils.py
+11
-1
@@ -41,6 +41,7 @@ from vllm.v1.kv_cache_interface import (
41
FullAttentionSpec,
42
KVCacheConfig,
43
KVCacheGroupSpec,
44
+ KVCacheTensor,
45
)
46
from vllm.v1.kv_offload.base import (
47
CanonicalKVCaches,
@@ -241,9 +242,18 @@ class RequestRunner:
242
)
243
]
244
245
+ kv_cache_tensors = [
246
+ KVCacheTensor(
247
+ size=group.kv_cache_spec.page_size_bytes * num_gpu_blocks,
248
+ shared_by=[layer_name],
249
+ )
250
+ for group in kv_cache_groups
251
+ for layer_name in group.layer_names
252
+ ]
253
+
254
kv_cache_config = KVCacheConfig(
255
num_blocks=num_gpu_blocks,
246
- kv_cache_tensors=[],
256
+ kv_cache_tensors=kv_cache_tensors,
257
kv_cache_groups=kv_cache_groups,
258
)
259
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+15
-1
@@ -53,6 +53,16 @@ class OffloadingConnectorWorker:
53
kv_cache_config = self.spec.kv_cache_config
54
num_blocks = kv_cache_config.num_blocks
55
56
+ # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
57
+ # stride(0) as the manager-block stride (equals total_num_bytes_per_block).
58
+ # General (non-packed) layouts size the tensor at page_size_bytes per
59
+ # manager block, so page_size_bytes is the correct offloading stride.
60
+ layer_is_packed: dict[str, bool] = {
61
+ ln: bool(kv_tensor.block_stride)
62
+ for kv_tensor in kv_cache_config.kv_cache_tensors
63
+ for ln in kv_tensor.shared_by
64
+ }
65
+
66
# layer_name -> (num_blocks, page_size_bytes) tensor
67
tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
68
# layer_name -> size of (un-padded) page in bytes
@@ -77,7 +87,11 @@ class OffloadingConnectorWorker:
87
page = layer_kv_cache_spec.page_size_bytes
88
elem_size = layer_kv_cache.element_size()
89
byte_offset = layer_kv_cache.storage_offset() * elem_size
80
- block_stride_bytes = layer_kv_cache.stride(0) * elem_size
90
+ block_stride_bytes = (
91
+ layer_kv_cache.stride(0) * elem_size
92
+ if layer_is_packed[layer_name]
93
+ else page
94
+ )
95
tensors_per_block[layer_name] = (
96
torch.tensor(
97
[],