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, 00:30 UTC
[KV offload] Enable CPU KV offload on CUDA alike Platforms (#27770)
Signed-off-by: zhewenli <zhewenli@meta.com>
Zhewen Li committed
Oct 30, 2025 at 07:10 UTC
0fe01404082744c955d135c3634e17de1404b00c
2 files changed
+2
-6
tests/v1/kv_offload/test_cpu_offloading.py
-4
@@ -12,7 +12,6 @@ from tqdm import tqdm
12
from vllm import LLM, SamplingParams, TokensPrompt
13
from vllm.config import KVEventsConfig, KVTransferConfig
14
from vllm.distributed.kv_events import BlockStored, KVEventBatch
15
-from vllm.platforms import current_platform
15
16
CPU_BLOCK_SIZES = [16, 48]
17
@@ -64,9 +63,6 @@ class MockSubscriber:
63
self.sub.close()
64
65
67
-@pytest.mark.skipif(
68
- not current_platform.is_cuda(), reason="CPU offloading only supported on CUDA"
69
-)
66
@pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
67
def test_cpu_offloading(cpu_block_size: int) -> None:
68
"""
vllm/v1/kv_offload/cpu.py
+2
-2
@@ -51,9 +51,9 @@ class CPUOffloadingSpec(OffloadingSpec):
51
self, kv_caches: dict[str, torch.Tensor]
52
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
53
if not self._handler:
54
- if not current_platform.is_cuda():
54
+ if not current_platform.is_cuda_alike():
55
raise Exception(
56
- "CPU Offloading is currently only supported on CUDA GPUs"
56
+ "CPU Offloading is currently only supported on CUDA-alike GPUs"
57
)
58
59
layer_names = list(kv_caches.keys())