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 6, 12:00 UTC
[kv_offload+HMA][2/N]: Support multiple KV groups in GPULoadStoreSpec (#36642)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Mar 18, 2026 at 19:26 UTC
5dd8df070172ac20e99a7dbd3d96cb6b054f0f57
3 files changed
+43
-10
tests/v1/kv_offload/test_cpu_gpu.py
+4
-8
@@ -135,19 +135,19 @@ def test_transfer(
135
# set transfer direction
136
if gpu_to_cpu:
137
handler = handlers.gpu_to_cpu_handler
138
- src_spec_class = GPULoadStoreSpec
139
- dst_spec_class = CPULoadStoreSpec
138
src_blocks = gpu_blocks
139
dst_blocks = cpu_blocks
140
+ src_spec = GPULoadStoreSpec(src_blocks, group_sizes=(len(src_blocks),))
141
+ dst_spec = CPULoadStoreSpec(dst_blocks)
142
src_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
143
dst_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
144
dst_size_in_kernel_blocks = num_cpu_blocks * kernel_blocks_per_cpu_block
145
else:
146
handler = handlers.cpu_to_gpu_handler
147
- src_spec_class = CPULoadStoreSpec
148
- dst_spec_class = GPULoadStoreSpec
147
src_blocks = cpu_blocks
148
dst_blocks = gpu_blocks
149
+ src_spec = CPULoadStoreSpec(src_blocks)
150
+ dst_spec = GPULoadStoreSpec(dst_blocks, group_sizes=(len(dst_blocks),))
151
src_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
152
dst_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
153
dst_size_in_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
@@ -159,10 +159,6 @@ def test_transfer(
159
):
160
dst_to_src[dst_block] = src_block
161
162
- # build transfer specs
163
- src_spec = src_spec_class(src_blocks)
164
- dst_spec = dst_spec_class(dst_blocks)
165
-
162
# clone src and dst tensors before transfer
163
orig_src_caches = [x.clone() for x in handler.src_tensors]
164
orig_dst_caches = [x.clone() for x in handler.dst_tensors]
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+8
-2
@@ -173,7 +173,11 @@ class OffloadingConnectorScheduler:
173
)
174
175
src_spec = self.manager.prepare_load(block_hashes)
176
- dst_spec = GPULoadStoreSpec(block_ids[num_computed_gpu_blocks:])
176
+ dst_spec = GPULoadStoreSpec(
177
+ block_ids[num_computed_gpu_blocks:],
178
+ group_sizes=(num_pending_gpu_blocks,),
179
+ block_indices=(num_computed_gpu_blocks,),
180
+ )
181
182
block_hashes = self._get_block_hashes(
183
request, start_idx=start_block_idx, end_idx=num_blocks
@@ -246,7 +250,9 @@ class OffloadingConnectorScheduler:
250
gpu_block_idx = offloaded_block_idx * self.block_size_factor
251
for i in range(self.block_size_factor):
252
src_block_ids.append(block_ids[gpu_block_idx + i])
249
- src_spec = GPULoadStoreSpec(src_block_ids)
253
+ src_spec = GPULoadStoreSpec(
254
+ src_block_ids, group_sizes=(len(src_block_ids),)
255
+ )
256
257
reqs_to_store[req_id] = (src_spec, dst_spec)
258
self._reqs_being_stored[req_id] |= block_hashes_to_store
vllm/v1/kv_offload/mediums.py
+31
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from abc import ABC
4
+from collections.abc import Sequence
5
6
import numpy as np
7
@@ -22,8 +23,38 @@ class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
23
class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
24
"""
25
Spec for loading/storing a KV block to GPU memory.
26
+
27
+ If there are multiple KV groups, the blocks are expected to be
28
+ ordered by the group index.
29
+ In that case, group_sizes[i] determines the number of blocks
30
+ per the i-th KV group, and thus sum(group_sizes) == len(block_ids).
31
+ group_sizes=None indicates a single KV group.
32
+
33
+ If block_indices is given, each group (determined by group_sizes) of block IDs
34
+ will correspond to logically contiguous blocks, e.g. blocks 5-10 of a some request.
35
+ block_indices[i] will represent the block index of the first block in group #i.
36
+ Thus, len(block_indices) == len(group_sizes) = number of KV cache groups.
37
+ This information is required in order to support loading from offloaded blocks
38
+ which are larger than GPU blocks.
39
+ In such cases, the first GPU block per each group may be unaligned to the offloaded
40
+ block size, and so knowing block_indices[i] allows the worker to correctly
41
+ skip part of the first matching offloaded block.
42
+ Offloading from GPU is always aligned to offloaded block size, and so
43
+ block_indices will only be set by the offloading connector when loading into GPU.
44
"""
45
46
+ def __init__(
47
+ self,
48
+ block_ids: list[int],
49
+ group_sizes: Sequence[int],
50
+ block_indices: Sequence[int] | None = None,
51
+ ):
52
+ super().__init__(block_ids)
53
+ assert sum(group_sizes) == len(block_ids)
54
+ assert block_indices is None or len(block_indices) == len(group_sizes)
55
+ self.group_sizes: Sequence[int] = group_sizes
56
+ self.block_indices: Sequence[int] | None = block_indices
57
+
58
@staticmethod
59
def medium() -> str:
60
return "GPU"