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 15, 02:15 UTC
[BugFix][CPU] Fix CPU worker dependency on cumem_allocator (#20696)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Jul 10, 2025 at 07:24 UTC
59389c927b7fdfa9e58dd538df834c89ab23a0c5
1 file changed
+9
-1
vllm/v1/worker/gpu_worker.py
+9
-1
@@ -11,7 +11,6 @@ import torch.nn as nn
11
12
import vllm.envs as envs
13
from vllm.config import VllmConfig
14
-from vllm.device_allocator.cumem import CuMemAllocator
14
from vllm.distributed import (ensure_model_parallel_initialized,
15
init_distributed_environment,
16
set_custom_all_reduce)
@@ -79,6 +78,8 @@ class Worker(WorkerBase):
78
self.profiler = None
79
80
def sleep(self, level: int = 1) -> None:
81
+ from vllm.device_allocator.cumem import CuMemAllocator
82
+
83
free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
84
85
# Save the buffers before level 2 sleep
@@ -101,6 +102,8 @@ class Worker(WorkerBase):
102
used_bytes / GiB_bytes)
103
104
def wake_up(self, tags: Optional[list[str]] = None) -> None:
105
+ from vllm.device_allocator.cumem import CuMemAllocator
106
+
107
allocator = CuMemAllocator.get_instance()
108
allocator.wake_up(tags)
109
@@ -174,6 +177,8 @@ class Worker(WorkerBase):
177
# to hijack tensor allocation.
178
def load_model(self) -> None:
179
if self.vllm_config.model_config.enable_sleep_mode:
180
+ from vllm.device_allocator.cumem import CuMemAllocator
181
+
182
allocator = CuMemAllocator.get_instance()
183
assert allocator.get_current_usage() == 0, (
184
"Sleep mode can only be "
@@ -241,7 +246,10 @@ class Worker(WorkerBase):
246
247
def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
248
"""Allocate GPU KV cache with the specified kv_cache_config."""
249
+
250
if self.vllm_config.model_config.enable_sleep_mode:
251
+ from vllm.device_allocator.cumem import CuMemAllocator
252
+
253
allocator = CuMemAllocator.get_instance()
254
context = allocator.use_memory_pool(tag="kv_cache")
255
else: