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 10, 07:45 UTC
OffloadingConnector: Add cpu_bytes_to_use configuration (#24498)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Jan 12, 2026 at 17:00 UTC
9cddbdba6d714bb0b0303b0648c8f8bb14d0b8f2
9 files changed
+49
-24
docs/features/disagg_prefill.md
+2
-2
@@ -37,10 +37,10 @@ For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
37
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
38
```
39
40
-- **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and number of blocks to allocate (per worker):
40
+- **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and total CPU memory bytes to allocate:
41
42
```bash
43
- --kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "num_cpu_blocks": 1000}}'
43
+ --kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
44
```
45
46
## Benchmarks
tests/v1/kv_connector/unit/test_config.py
+2
-3
@@ -15,7 +15,7 @@ pytestmark = pytest.mark.cpu_test
15
[
16
("native", 4.0, 1, 1, "OffloadingConnector", 4.0 * (1 << 30)),
17
# bytes per rank: 8.0 GiB / (2 * 2) = 2.0 GiB
18
- ("native", 8.0, 2, 2, "OffloadingConnector", 8.0 * (1 << 30) / 4),
18
+ ("native", 8.0, 2, 2, "OffloadingConnector", 8.0 * (1 << 30)),
19
("lmcache", 4.0, 1, 1, "LMCacheConnectorV1", 4.0),
20
# size per rank: 8.0 GiB / (2 * 2) = 2.0 GiB
21
("lmcache", 8.0, 2, 2, "LMCacheConnectorV1", 2.0),
@@ -54,8 +54,7 @@ def test_kv_connector(
54
assert kv_transfer_config.kv_role == "kv_both"
55
56
if kv_offloading_backend == "native":
57
- assert kv_connector_extra_config["kv_bytes_per_rank"] == expected_bytes
58
- assert kv_connector_extra_config["num_cpu_blocks"] == 0
57
+ assert kv_connector_extra_config["cpu_bytes_to_use"] == expected_bytes
58
# Existing config should be preserved
59
assert kv_connector_extra_config["existing_key"] == "existing_value"
60
elif kv_offloading_backend == "lmcache":
tests/v1/kv_connector/unit/test_offloading_connector.py
+3
-2
@@ -26,6 +26,7 @@ from vllm.v1.core.kv_cache_utils import (
26
init_none_hash,
27
)
28
from vllm.v1.core.sched.scheduler import Scheduler
29
+from vllm.v1.kv_cache_interface import KVCacheConfig
30
from vllm.v1.kv_offload.abstract import (
31
LoadStoreSpec,
32
OffloadingEvent,
@@ -93,8 +94,8 @@ class MockOffloadingHandler(OffloadingHandler):
94
95
96
class MockOffloadingSpec(OffloadingSpec):
96
- def __init__(self, vllm_config: VllmConfig):
97
- super().__init__(vllm_config)
97
+ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
98
+ super().__init__(vllm_config, kv_cache_config)
99
100
self.manager = MagicMock(spec=OffloadingManager)
101
self.manager.lookup.return_value = 0
tests/v1/kv_offload/test_cpu_offloading.py
+1
-1
@@ -161,7 +161,7 @@ def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
161
kv_connector="OffloadingConnector",
162
kv_role="kv_both",
163
kv_connector_extra_config={
164
- "num_cpu_blocks": 1000,
164
+ "cpu_bytes_to_use": 500 << 20,
165
"block_size": cpu_block_size,
166
},
167
)
vllm/config/vllm.py
+1
-5
@@ -516,12 +516,8 @@ class VllmConfig:
516
517
if kv_offloading_backend == "native":
518
self.kv_transfer_config.kv_connector = "OffloadingConnector"
519
- kv_bytes_per_rank = kv_offloading_size * (1 << 30) / num_kv_ranks
520
-
521
- # NOTE(ApostaC): the actual calculation for num_cpu_blocks should be
522
- # done after the model's KV cache is initialized
519
self.kv_transfer_config.kv_connector_extra_config.update(
524
- {"kv_bytes_per_rank": kv_bytes_per_rank, "num_cpu_blocks": 0}
520
+ {"cpu_bytes_to_use": kv_offloading_size * (1 << 30)}
521
)
522
elif kv_offloading_backend == "lmcache":
523
self.kv_transfer_config.kv_connector = "LMCacheConnectorV1"
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+1
-1
@@ -56,7 +56,7 @@ class OffloadingConnector(KVConnectorBase_V1):
56
):
57
super().__init__(vllm_config, role, kv_cache_config)
58
59
- spec = OffloadingSpecFactory.create_spec(vllm_config)
59
+ spec = OffloadingSpecFactory.create_spec(vllm_config, kv_cache_config)
60
61
self.connector_scheduler: OffloadingConnectorScheduler | None = None
62
self.connector_worker: OffloadingConnectorWorker | None = None
vllm/v1/kv_offload/cpu.py
+31
-8
@@ -7,6 +7,7 @@ import torch
7
from vllm.config import VllmConfig
8
from vllm.platforms import current_platform
9
from vllm.v1.attention.backend import AttentionBackend
10
+from vllm.v1.kv_cache_interface import KVCacheConfig
11
from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
12
from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
13
from vllm.v1.kv_offload.backends.cpu import CPUBackend
@@ -18,15 +19,37 @@ from vllm.v1.kv_offload.worker.worker import OffloadingHandler
19
20
21
class CPUOffloadingSpec(OffloadingSpec):
21
- def __init__(self, vllm_config: VllmConfig):
22
- super().__init__(vllm_config)
22
+ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
23
+ super().__init__(vllm_config, kv_cache_config)
24
24
- num_cpu_blocks = self.extra_config.get("num_cpu_blocks")
25
- if not num_cpu_blocks:
25
+ cpu_bytes_to_use = self.extra_config.get("cpu_bytes_to_use")
26
+ if not cpu_bytes_to_use:
27
raise Exception(
27
- "num_cpu_blocks must be specified in kv_connector_extra_config"
28
+ "cpu_bytes_to_use must be specified in kv_connector_extra_config"
29
)
29
- self.num_cpu_blocks: int = num_cpu_blocks
30
+
31
+ # calculate kv_bytes_per_offloaded_block
32
+ assert kv_cache_config is not None
33
+ page_sizes = {
34
+ kv_cache_group.kv_cache_spec.page_size_bytes
35
+ for kv_cache_group in kv_cache_config.kv_cache_groups
36
+ }
37
+ assert len(page_sizes) == 1
38
+ page_size_bytes = page_sizes.pop()
39
+ kv_bytes_per_block = (
40
+ page_size_bytes
41
+ * len(kv_cache_config.kv_cache_tensors)
42
+ * vllm_config.parallel_config.world_size
43
+ )
44
+ kv_bytes_per_offloaded_block = kv_bytes_per_block * (
45
+ self.offloaded_block_size // self.gpu_block_size
46
+ )
47
+
48
+ self.num_blocks = (
49
+ int(cpu_bytes_to_use) // kv_bytes_per_offloaded_block
50
+ if kv_bytes_per_offloaded_block > 0
51
+ else 0
52
+ )
53
54
# scheduler-side
55
self._manager: OffloadingManager | None = None
@@ -44,7 +67,7 @@ class CPUOffloadingSpec(OffloadingSpec):
67
)
68
69
backend = CPUBackend(
47
- block_size=self.offloaded_block_size, num_blocks=self.num_cpu_blocks
70
+ block_size=self.offloaded_block_size, num_blocks=self.num_blocks
71
)
72
73
if self.eviction_policy == "lru":
@@ -77,7 +100,7 @@ class CPUOffloadingSpec(OffloadingSpec):
100
attn_backends=attn_backends,
101
gpu_block_size=self.gpu_block_size,
102
cpu_block_size=self.offloaded_block_size,
80
- num_cpu_blocks=self.num_cpu_blocks,
103
+ num_cpu_blocks=self.num_blocks,
104
gpu_caches=kv_caches,
105
)
106
vllm/v1/kv_offload/factory.py
+3
-1
@@ -9,6 +9,7 @@ from vllm.v1.kv_offload.spec import OffloadingSpec
9
10
if TYPE_CHECKING:
11
from vllm.config import VllmConfig
12
+ from vllm.v1.kv_cache_interface import KVCacheConfig
13
14
logger = init_logger(__name__)
15
@@ -32,6 +33,7 @@ class OffloadingSpecFactory:
33
def create_spec(
34
cls,
35
config: "VllmConfig",
36
+ kv_cache_config: "KVCacheConfig | None",
37
) -> OffloadingSpec:
38
kv_transfer_config = config.kv_transfer_config
39
assert kv_transfer_config is not None
@@ -47,7 +49,7 @@ class OffloadingSpecFactory:
49
spec_cls = getattr(spec_module, spec_name)
50
assert issubclass(spec_cls, OffloadingSpec)
51
logger.info("Creating offloading spec with name: %s", spec_name)
50
- return spec_cls(config)
52
+ return spec_cls(config, kv_cache_config)
53
54
55
# Register various specs here.
vllm/v1/kv_offload/spec.py
+5
-1
@@ -13,6 +13,7 @@ from vllm.v1.kv_offload.worker.worker import OffloadingHandler
13
14
if TYPE_CHECKING:
15
from vllm.config import VllmConfig
16
+ from vllm.v1.kv_cache_interface import KVCacheConfig
17
18
logger = init_logger(__name__)
19
@@ -20,12 +21,15 @@ logger = init_logger(__name__)
21
class OffloadingSpec(ABC):
22
"""Spec for an offloading connector"""
23
23
- def __init__(self, vllm_config: "VllmConfig"):
24
+ def __init__(
25
+ self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig | None"
26
+ ):
27
logger.warning(
28
"Initializing OffloadingSpec. This API is experimental and "
29
"subject to change in the future as we iterate the design."
30
)
31
self.vllm_config = vllm_config
32
+ self.kv_cache_config = kv_cache_config
33
34
kv_transfer_config = vllm_config.kv_transfer_config
35
assert kv_transfer_config is not None