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 12, 21:00 UTC
[BugFix] kv_offloading: Fix bug in loading of partial cpu blocks (#28951)
Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Or Ozeri committed
Nov 20, 2025 at 12:55 UTC
c0c2dd1e0b75c70706f4d8dbcd1d75f1c1750e14
2 files changed
+11
-13
tests/v1/kv_offload/test_cpu_gpu.py
+2
-2
@@ -103,8 +103,8 @@ def test_transfer(
103
for i in range(gpu_blocks_per_cpu_block):
104
cpu_blocks_in_gpu_block_size.append(i + base_block_id)
105
106
- # maybe skip a GPU block to test writing to the middle of a CPU block
107
- if gpu_to_cpu:
106
+ # maybe skip a GPU block to test reading from the middle of a CPU block
107
+ if not gpu_to_cpu:
108
gpu_blocks = gpu_blocks[gpu_blocks_per_cpu_block - 1 :]
109
cpu_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size[
110
gpu_blocks_per_cpu_block - 1 :
vllm/v1/kv_offload/worker/cpu_gpu.py
+9
-11
@@ -135,22 +135,20 @@ class CpuGpuOffloadingHandler(OffloadingHandler):
135
assert src_blocks.ndim == 1
136
assert dst_blocks.ndim == 1
137
138
- dst_sub_blocks_to_skip = -src_blocks.size % dst_block_size_factor
138
src_sub_block_count = src_blocks.size * src_block_size_factor
139
+ dst_sub_block_count = dst_blocks.size * dst_block_size_factor
140
+ src_sub_blocks_to_skip = -dst_blocks.size % src_block_size_factor
141
141
- assert (
142
- src_sub_block_count
143
- == dst_blocks.size * dst_block_size_factor - dst_sub_blocks_to_skip
144
- )
142
+ assert dst_sub_block_count == src_sub_block_count - src_sub_blocks_to_skip
143
146
- src_to_dst = np.empty((src_sub_block_count, 2), dtype=np.int64)
147
- expand_block_ids(src_blocks, src_block_size_factor, src_to_dst[:, 0])
144
+ src_to_dst = np.empty((dst_sub_block_count, 2), dtype=np.int64)
145
expand_block_ids(
149
- dst_blocks,
150
- dst_block_size_factor,
151
- src_to_dst[:, 1],
152
- skip_count=dst_sub_blocks_to_skip,
146
+ src_blocks,
147
+ src_block_size_factor,
148
+ src_to_dst[:, 0],
149
+ skip_count=src_sub_blocks_to_skip,
150
)
151
+ expand_block_ids(dst_blocks, dst_block_size_factor, src_to_dst[:, 1])
152
src_to_dst_tensor = torch.from_numpy(src_to_dst)
153
154
event = self.events_pool.pop() if self.events_pool else torch.Event()