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 14, 05:30 UTC
[KV offload][5/N] Add `CPUOffloadingSpec` (#24251)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Sep 22, 2025 at 22:30 UTC
8db29392899b12843787b0a03880f70a6e6ac88e
5 files changed
+146
docs/features/disagg_prefill.md
+6
@@ -31,6 +31,12 @@ Now supports 5 types of connectors:
31
--kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both"},{"kv_connector":"SharedStorageConnector","kv_role":"kv_both","kv_connector_extra_config":{"shared_storage_path":"local_storage"}}]}}'
32
```
33
34
+- **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and number of blocks to allocate (per worker):
35
+
36
+ ```bash
37
+ --kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "num_cpu_blocks": 1000}}'
38
+ ```
39
+
40
## Benchmarks
41
42
Please refer to <gh-file:benchmarks/disagg_benchmarks> for disaggregated prefilling benchmarks.
tests/v1/kv_offload/test_cpu_manager.py
renamed
tests/v1/kv_offload/test_cpu_offloading.py
new
+62
@@ -0,0 +1,62 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import time
4
+
5
+import pytest
6
+
7
+from vllm import LLM, SamplingParams
8
+from vllm.config import KVTransferConfig
9
+
10
+CPU_BLOCK_SIZES = [16, 48]
11
+
12
+
13
+@pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
14
+def test_cpu_offloading(cpu_block_size: int) -> None:
15
+ """
16
+ Tests OffloadingConnector with CPUOffloadingSpec.
17
+ """
18
+
19
+ # configure OffloadingConnector (spec_name=CPUOffloadingSpec by default)
20
+ kv_transfer_config = KVTransferConfig(
21
+ kv_connector="OffloadingConnector",
22
+ kv_role="kv_both",
23
+ kv_connector_extra_config={
24
+ "num_cpu_blocks": 100,
25
+ "block_size": cpu_block_size
26
+ },
27
+ )
28
+
29
+ llm = LLM(
30
+ model="meta-llama/Llama-3.2-1B-Instruct",
31
+ gpu_memory_utilization=0.5,
32
+ kv_transfer_config=kv_transfer_config,
33
+ )
34
+
35
+ prompts = ["Hi " * 100]
36
+ sampling_params = SamplingParams(temperature=0, max_tokens=20)
37
+
38
+ # run generation - this should trigger saving KV cache
39
+ start_time = time.time()
40
+ llm.generate(prompts, sampling_params, use_tqdm=False)
41
+ cold_time = time.time() - start_time
42
+
43
+ # run generation again - should hit the GPU prefix cache
44
+ start_time = time.time()
45
+ llm.generate(prompts, sampling_params, use_tqdm=False)
46
+ gpu_hit_time = time.time() - start_time
47
+
48
+ # reset prefix cache to avoid GPU hit.
49
+ llm.reset_prefix_cache()
50
+
51
+ # sleep for a sec to make sure CPU finished storing
52
+ time.sleep(1)
53
+
54
+ # run generation again - this should trigger loading from CPU
55
+ start_time = time.time()
56
+ llm.generate(prompts, sampling_params, use_tqdm=False)
57
+ cpu_hit_time = time.time() - start_time
58
+
59
+ print("Generation times:")
60
+ print(f" Cold: {cold_time * 1000:.2f}ms")
61
+ print(f" GPU hit: {gpu_hit_time * 1000:.2f}ms")
62
+ print(f" CPU hit: {cpu_hit_time * 1000:.2f}ms")
vllm/v1/kv_offload/cpu.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterator
4
+from typing import Optional
5
+
6
+import torch
7
+
8
+from vllm.config import VllmConfig, get_layers_from_vllm_config
9
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
+from vllm.platforms import current_platform
11
+from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
12
+from vllm.v1.kv_offload.backends.cpu import CPUBackend
13
+from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager
14
+from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
15
+from vllm.v1.kv_offload.spec import OffloadingSpec
16
+from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandler
17
+from vllm.v1.kv_offload.worker.worker import OffloadingHandler
18
+
19
+
20
+class CPUOffloadingSpec(OffloadingSpec):
21
+
22
+ def __init__(self, vllm_config: VllmConfig):
23
+ super().__init__(vllm_config)
24
+
25
+ num_cpu_blocks = self.extra_config.get("num_cpu_blocks")
26
+ if not num_cpu_blocks:
27
+ raise Exception("num_cpu_blocks must be specified "
28
+ "in kv_connector_extra_config")
29
+ self.num_cpu_blocks: int = num_cpu_blocks
30
+
31
+ # scheduler-side
32
+ self._manager: Optional[OffloadingManager] = None
33
+
34
+ # worker-side
35
+ self._handler: Optional[OffloadingHandler] = None
36
+
37
+ def get_manager(self) -> OffloadingManager:
38
+ if not self._manager:
39
+ kv_events_config = self.vllm_config.kv_events_config
40
+ enable_events = (kv_events_config is not None
41
+ and kv_events_config.enable_kv_cache_events)
42
+ self._manager = LRUOffloadingManager(CPUBackend(
43
+ block_size=self.offloaded_block_size,
44
+ num_blocks=self.num_cpu_blocks),
45
+ enable_events=enable_events)
46
+ return self._manager
47
+
48
+ def get_handlers(
49
+ self, kv_caches: dict[str, torch.Tensor]
50
+ ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec],
51
+ OffloadingHandler]]:
52
+ if not self._handler:
53
+ if not current_platform.is_cuda():
54
+ raise Exception("CPU Offloading is currently only supported"
55
+ " on CUDA GPUs")
56
+
57
+ layer_names = list(kv_caches.keys())
58
+ layers = get_layers_from_vllm_config(self.vllm_config,
59
+ AttentionLayerBase,
60
+ layer_names)
61
+ attn_backends = {
62
+ layer_name: layers[layer_name].get_attn_backend()
63
+ for layer_name in layer_names
64
+ }
65
+
66
+ self._handler = CpuGpuOffloadingHandler(
67
+ attn_backends=attn_backends,
68
+ gpu_block_size=self.gpu_block_size,
69
+ cpu_block_size=self.offloaded_block_size,
70
+ num_cpu_blocks=self.num_cpu_blocks,
71
+ gpu_caches=kv_caches)
72
+
73
+ assert self._handler is not None
74
+ yield GPULoadStoreSpec, CPULoadStoreSpec, self._handler
75
+ yield CPULoadStoreSpec, GPULoadStoreSpec, self._handler
vllm/v1/kv_offload/factory.py
+3
@@ -51,3 +51,6 @@ class OffloadingSpecFactory:
51
52
53
# Register various specs here.
54
+OffloadingSpecFactory.register_spec("CPUOffloadingSpec",
55
+ "vllm.v1.kv_offload.cpu",
56
+ "CPUOffloadingSpec")