A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 01:00 UTC
[kv_offload]: Fix num CPU blocks for UniformTypeKVCacheSpecs (#39617)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Apr 17, 2026 at 15:13 UTC
79a5b6325396683d3063c41fd1a140c56c89e982
1 file changed
+7
-11
vllm/v1/kv_offload/cpu/spec.py
+7
-11
@@ -26,17 +26,13 @@ class CPUOffloadingSpec(OffloadingSpec):
26
27
# calculate kv_bytes_per_offloaded_block
28
assert kv_cache_config is not None
29
- page_sizes = {
30
- kv_cache_group.kv_cache_spec.page_size_bytes
31
- for kv_cache_group in kv_cache_config.kv_cache_groups
32
- }
33
- assert len(page_sizes) == 1
34
- page_size_bytes = page_sizes.pop()
35
- kv_bytes_per_block = (
36
- page_size_bytes
37
- * len(kv_cache_config.kv_cache_tensors)
38
- * vllm_config.parallel_config.world_size
39
- )
29
+ if kv_cache_config.num_blocks > 0:
30
+ total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
31
+ kv_bytes_per_block = (
32
+ total_gpu_kv_bytes // kv_cache_config.num_blocks
33
+ ) * vllm_config.parallel_config.world_size
34
+ else:
35
+ kv_bytes_per_block = 0
36
37
kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
38
self.num_blocks = (