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, 19:15 UTC
[kv_offload+HMA][1/N]: Support multiple KV groups in OffloadingSpec (#36610)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Mar 13, 2026 at 10:04 UTC
cfaf4668f7100a279e6ac8c07921213169d5230c
5 files changed
+80
-22
tests/v1/kv_connector/unit/test_offloading_connector.py
+37
-5
@@ -26,8 +26,13 @@ from vllm.v1.core.kv_cache_utils import (
26
get_request_block_hasher,
27
init_none_hash,
28
)
29
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
30
from vllm.v1.core.sched.scheduler import Scheduler
30
-from vllm.v1.kv_cache_interface import KVCacheConfig
31
+from vllm.v1.kv_cache_interface import (
32
+ FullAttentionSpec,
33
+ KVCacheConfig,
34
+ KVCacheGroupSpec,
35
+)
36
from vllm.v1.kv_offload.abstract import (
37
LoadStoreSpec,
38
OffloadingEvent,
@@ -43,11 +48,11 @@ from vllm.v1.kv_offload.worker.worker import (
48
)
49
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
50
from vllm.v1.request import Request, RequestStatus
51
+from vllm.v1.structured_output import StructuredOutputManager
52
53
from .utils import (
54
EOS_TOKEN_ID,
55
create_model_runner_output,
50
- create_scheduler,
56
create_vllm_config,
57
)
58
@@ -175,10 +180,37 @@ class RequestRunner:
180
},
181
)
182
178
- self.scheduler: Scheduler = create_scheduler(
179
- vllm_config, num_blocks=num_gpu_blocks
183
+ block_size = vllm_config.cache_config.block_size
184
+ kv_cache_config = KVCacheConfig(
185
+ num_blocks=num_gpu_blocks,
186
+ kv_cache_tensors=[],
187
+ kv_cache_groups=[
188
+ KVCacheGroupSpec(
189
+ ["layer"],
190
+ FullAttentionSpec(
191
+ block_size=block_size,
192
+ num_kv_heads=1,
193
+ head_size=1,
194
+ dtype=torch.float32,
195
+ ),
196
+ )
197
+ ],
198
+ )
199
+ vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
200
+ self.num_kv_groups = len(kv_cache_config.kv_cache_groups)
201
+
202
+ scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
203
+ self.scheduler = scheduler_cls(
204
+ vllm_config=vllm_config,
205
+ kv_cache_config=kv_cache_config,
206
+ log_stats=True,
207
+ structured_output_manager=StructuredOutputManager(vllm_config),
208
+ block_size=block_size,
209
+ )
210
+
211
+ self.worker_connector = OffloadingConnector(
212
+ vllm_config, KVConnectorRole.WORKER, kv_cache_config
213
)
181
- self.worker_connector = OffloadingConnector(vllm_config, KVConnectorRole.WORKER)
214
215
# register worker kv_caches to enable OffloadingWorker creations
216
self.worker_connector.register_cross_layers_kv_cache(
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+5
-3
@@ -126,6 +126,7 @@ class OffloadingConnector(KVConnectorBase_V1):
126
):
127
super().__init__(vllm_config, role, kv_cache_config)
128
129
+ assert kv_cache_config is not None
130
spec = OffloadingSpecFactory.create_spec(vllm_config, kv_cache_config)
131
132
self.connector_scheduler: OffloadingConnectorScheduler | None = None
@@ -245,9 +246,10 @@ class OffloadingConnectorScheduler:
246
"""Implementation of Scheduler side methods"""
247
248
def __init__(self, spec: OffloadingSpec):
248
- self.gpu_block_size = spec.gpu_block_size
249
- self.offloaded_block_size = spec.offloaded_block_size
250
- self.block_size_factor = self.offloaded_block_size // self.gpu_block_size
249
+ assert len(spec.gpu_block_size) == 1
250
+ self.gpu_block_size = spec.gpu_block_size[0]
251
+ self.offloaded_block_size = self.gpu_block_size * spec.block_size_factor
252
+ self.block_size_factor = spec.block_size_factor
253
self.manager: OffloadingManager = spec.get_manager()
254
255
self._requests: dict[ReqId, Request] = {}
vllm/v1/kv_offload/cpu.py
+10
-6
@@ -42,10 +42,8 @@ class CPUOffloadingSpec(OffloadingSpec):
42
* len(kv_cache_config.kv_cache_tensors)
43
* vllm_config.parallel_config.world_size
44
)
45
- kv_bytes_per_offloaded_block = kv_bytes_per_block * (
46
- self.offloaded_block_size // self.gpu_block_size
47
- )
45
46
+ kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
47
self.num_blocks = (
48
int(cpu_bytes_to_use) // kv_bytes_per_offloaded_block
49
if kv_bytes_per_offloaded_block > 0
@@ -67,8 +65,11 @@ class CPUOffloadingSpec(OffloadingSpec):
65
kv_events_config is not None and kv_events_config.enable_kv_cache_events
66
)
67
68
+ assert len(self.gpu_block_size) == 1
69
+ gpu_block_size = self.gpu_block_size[0]
70
+ offloaded_block_size = gpu_block_size * self.block_size_factor
71
backend = CPUBackend(
71
- block_size=self.offloaded_block_size, num_blocks=self.num_blocks
72
+ block_size=offloaded_block_size, num_blocks=self.num_blocks
73
)
74
75
if self.eviction_policy == "lru":
@@ -111,10 +112,13 @@ class CPUOffloadingSpec(OffloadingSpec):
112
"CPU Offloading is currently only supported on CUDA-alike GPUs"
113
)
114
115
+ assert len(self.gpu_block_size) == 1
116
+ gpu_block_size = self.gpu_block_size[0]
117
+
118
self._handlers = CpuGpuOffloadingHandlers(
119
attn_backends=attn_backends,
116
- gpu_block_size=self.gpu_block_size,
117
- cpu_block_size=self.offloaded_block_size,
120
+ gpu_block_size=gpu_block_size,
121
+ cpu_block_size=gpu_block_size * self.block_size_factor,
122
num_cpu_blocks=self.num_blocks,
123
gpu_caches=kv_caches,
124
)
vllm/v1/kv_offload/factory.py
+1
-1
@@ -33,7 +33,7 @@ class OffloadingSpecFactory:
33
def create_spec(
34
cls,
35
config: "VllmConfig",
36
- kv_cache_config: "KVCacheConfig | None",
36
+ kv_cache_config: "KVCacheConfig",
37
) -> OffloadingSpec:
38
kv_transfer_config = config.kv_transfer_config
39
assert kv_transfer_config is not None
vllm/v1/kv_offload/spec.py
+27
-7
@@ -21,9 +21,7 @@ logger = init_logger(__name__)
21
class OffloadingSpec(ABC):
22
"""Spec for an offloading connector"""
23
24
- def __init__(
25
- self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig | None"
26
- ):
24
+ def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
25
logger.warning(
26
"Initializing OffloadingSpec. This API is experimental and "
27
"subject to change in the future as we iterate the design."
@@ -35,12 +33,34 @@ class OffloadingSpec(ABC):
33
assert kv_transfer_config is not None
34
self.extra_config = kv_transfer_config.kv_connector_extra_config
35
38
- self.gpu_block_size = vllm_config.cache_config.block_size
39
- self.offloaded_block_size = int(
40
- self.extra_config.get("block_size", self.gpu_block_size)
36
+ # block size used by vLLM for hashing request tokens for the sake
37
+ # of enabling prefix caching
38
+ self.hash_block_size = vllm_config.cache_config.block_size
39
+ # gpu block size per group
40
+ self.gpu_block_size: tuple[int, ...] = tuple(
41
+ kv_cache_group.kv_cache_spec.block_size
42
+ for kv_cache_group in kv_cache_config.kv_cache_groups
43
)
44
43
- assert self.offloaded_block_size % self.gpu_block_size == 0
45
+ for block_size in self.gpu_block_size:
46
+ assert block_size % self.hash_block_size == 0
47
+
48
+ # offloaded_block_size / gpu_block_size
49
+ self.block_size_factor: int = 1
50
+
51
+ offloaded_block_size = self.extra_config.get("block_size")
52
+ if offloaded_block_size is not None:
53
+ offloaded_block_size_int = int(offloaded_block_size)
54
+ gpu_block_sizes = set(self.gpu_block_size)
55
+ assert len(gpu_block_sizes) == 1, (
56
+ "If 'block_size' is specified in kv_connector_extra_config, "
57
+ "there must be at least one KV cache group, "
58
+ "and all groups must have the same block size."
59
+ )
60
+ gpu_block_size = gpu_block_sizes.pop()
61
+
62
+ assert offloaded_block_size_int % gpu_block_size == 0
63
+ self.block_size_factor = offloaded_block_size_int // gpu_block_size
64
65
@abstractmethod
66
def get_manager(self) -> OffloadingManager: