[Bugfix][KV Offload] Handle padded GPU cache storage (#54021)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Jiangyun Zhu committed
Aug 27, 2026 at 18:52 UTC
478ec3ea41358d410a0997d8852ef5a344c9a677
2 files changed
+39
-3
tests/v1/simple_kv_offload/test_worker.py
+32
-2
@@ -209,7 +209,10 @@ def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout):
209
device="cuda",
210
)
211
caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
212
- cache_config = MagicMock(num_blocks=num_blocks)
212
+ cache_config = MagicMock(
213
+ num_blocks=num_blocks,
214
+ kv_cache_tensors=[MagicMock(size=raw.nbytes)],
215
+ )
216
worker = SimpleCPUOffloadWorker(
217
vllm_config=None,
218
kv_cache_config=cache_config,
@@ -238,6 +241,30 @@ def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout):
241
}
242
243
244
+def test_register_kv_cache_storage_with_trailing_padding(monkeypatch):
245
+ num_blocks = 4
246
+ block_bytes = 32
247
+ cache_bytes = num_blocks * block_bytes
248
+ raw = torch.zeros(4096, dtype=torch.int8, device="cuda")
249
+ cache = raw[:cache_bytes].view(num_blocks, block_bytes)
250
+ worker = SimpleCPUOffloadWorker(
251
+ vllm_config=None,
252
+ kv_cache_config=MagicMock(
253
+ num_blocks=num_blocks,
254
+ kv_cache_tensors=[MagicMock(size=cache_bytes)],
255
+ ),
256
+ cpu_capacity_bytes=cache_bytes,
257
+ )
258
+ worker._backend = MagicMock()
259
+ monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
260
+
261
+ worker.register_kv_caches({"layer.0": cache})
262
+
263
+ assert worker.gpu_kv_caches is not None
264
+ assert list(worker.gpu_kv_caches) == ["layer.0"]
265
+ assert worker.gpu_kv_caches["layer.0"].shape == (num_blocks, block_bytes)
266
+
267
+
268
def test_register_separate_kv_head_groups(monkeypatch):
269
# LHBNC hoists the K/V head groups outside the block dim, so each layer's
270
# blocks are registered as one region per group (K, V).
@@ -260,7 +287,10 @@ def test_register_separate_kv_head_groups(monkeypatch):
287
caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
288
worker = SimpleCPUOffloadWorker(
289
vllm_config=None,
263
- kv_cache_config=MagicMock(num_blocks=num_blocks),
290
+ kv_cache_config=MagicMock(
291
+ num_blocks=num_blocks,
292
+ kv_cache_tensors=[MagicMock(size=raw.nbytes)],
293
+ ),
294
cpu_capacity_bytes=raw.nbytes,
295
)
296
worker._backend = MagicMock()
vllm/v1/simple_kv_offload/worker.py
+7
-1
@@ -99,6 +99,8 @@ class SimpleCPUOffloadWorker:
99
100
assert self.kv_cache_config is not None
101
num_blocks = self.kv_cache_config.num_blocks
102
+ assert self.kv_cache_config.kv_cache_tensors
103
+ logical_storage_bytes = self.kv_cache_config.kv_cache_tensors[0].size
104
105
# The DMA backend copies whole blocks as base + block_id * stride(0),
106
# so view each unique allocation as [num_blocks, block_bytes].
@@ -118,7 +120,11 @@ class SimpleCPUOffloadWorker:
120
)
121
block_bytes = tensor.stride(0) * tensor.element_size() * physical_per_block
122
raw = torch.empty(0, dtype=torch.int8, device=tensor.device).set_(storage)
121
- regions = raw.view(-1, num_blocks, block_bytes)
123
+ assert raw.numel() >= logical_storage_bytes, (
124
+ f"KV cache {name!r} storage has {raw.numel()} bytes, smaller "
125
+ f"than the configured {logical_storage_bytes}-byte allocation"
126
+ )
127
+ regions = raw[:logical_storage_bytes].view(-1, num_blocks, block_bytes)
128
for idx, region in enumerate(regions):
129
key_name = name if len(regions) == 1 else f"{name}.{idx}"
130
unique_gpu_caches[key_name] = region