ddb94c2605 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/basic_correctness/test_cumem.py

+19 -1

@@ -155,6 +155,24 @@ def test_end_to_end(monkeypatch: pytest.MonkeyPatch, model: str, use_v1: bool):

155

156 llm.wake_up()

157 output2 = llm.generate(prompt, sampling_params)

158 -

158 # cmp output

159 assert output[0].outputs[0].text == output2[0].outputs[0].text

160 +

161 + llm.sleep(level=1)

162 + llm.wake_up(tags=["weights"])

163 +

164 + free_gpu_bytes_wake_up_w, total = torch.cuda.mem_get_info()

165 + used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline

166 +

167 + # should just reallocate memory for weights (1B model, ~2GiB weights)

168 + if use_v1:

169 + assert used_bytes < 10 * GiB_bytes

170 + else:

171 + assert used_bytes < 6 * GiB_bytes

172 +

173 + # now allocate kv cache memory

174 + llm.wake_up(tags=["kv_cache"])

175 + output3 = llm.generate(prompt, sampling_params)

176 +

177 + # cmp output

178 + assert output[0].outputs[0].text == output3[0].outputs[0].text

tests/entrypoints/openai/test_sleep.py

+25 -4

@@ -25,16 +25,37 @@ def test_sleep_mode():

25 "VLLM_SERVER_DEV_MODE": "1",

26 "CUDA_VISIBLE_DEVICES": "0"

27 }) as remote_server:

28 + response = requests.post(remote_server.url_for("sleep"),

29 + params={"level": "1"})

30 + assert response.status_code == 200

31 + response = requests.get(remote_server.url_for("is_sleeping"))

32 + assert response.status_code == 200

33 + assert response.json().get("is_sleeping") is True

34 +

35 + response = requests.post(remote_server.url_for("wake_up"))

36 + assert response.status_code == 200

37 + response = requests.get(remote_server.url_for("is_sleeping"))

38 + assert response.status_code == 200

39 + assert response.json().get("is_sleeping") is False

40

29 - response = requests.post(remote_server.url_for("/sleep"),

41 + # test wake up with tags

42 + response = requests.post(remote_server.url_for("sleep"),

43 params={"level": "1"})

44 assert response.status_code == 200

32 - response = requests.get(remote_server.url_for("/is_sleeping"))

45 +

46 + response = requests.post(remote_server.url_for("wake_up"),

47 + params={"tags": ["weights"]})

48 + assert response.status_code == 200

49 +

50 + # is sleeping should be false after waking up any part of the engine

51 + response = requests.get(remote_server.url_for("is_sleeping"))

52 assert response.status_code == 200

53 assert response.json().get("is_sleeping") is True

54

36 - response = requests.post(remote_server.url_for("/wake_up"))

55 + response = requests.post(remote_server.url_for("wake_up"),

56 + params={"tags": ["kv_cache"]})

57 assert response.status_code == 200

38 - response = requests.get(remote_server.url_for("/is_sleeping"))

58 +

59 + response = requests.get(remote_server.url_for("is_sleeping"))

60 assert response.status_code == 200

61 assert response.json().get("is_sleeping") is False

vllm/device_allocator/cumem.py

+19 -13

@@ -208,22 +208,28 @@ class CuMemAllocator:

208 gc.collect()

209 torch.cuda.empty_cache()

210

211 - def wake_up(self):

211 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

212 """

213 Wake up the allocator from sleep mode.

214 - All data that is previously offloaded will be loaded back to GPU

215 - memory, and the rest of the data will have empty memory."""

214 + All data that is previously offloaded will be loaded back to GPU

215 + memory, and the rest of the data will have empty memory.

216 +

217 + :param tags: The tags of the memory allocation that will be loaded

218 + back to GPU memory. If None, all memory allocation will be loaded

219 + back to GPU memory.

220 + """

221 for ptr, data in self.pointer_to_data.items():

217 - handle = data.handle

218 - create_and_map(handle)

219 - if data.cpu_backup_tensor is not None:

220 - cpu_backup_tensor = data.cpu_backup_tensor

221 - if cpu_backup_tensor is not None:

222 - size_in_bytes = cpu_backup_tensor.numel(

223 - ) * cpu_backup_tensor.element_size()

224 - cpu_ptr = cpu_backup_tensor.data_ptr()

225 - libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes)

226 - data.cpu_backup_tensor = None

222 + if tags is None or data.tag in tags:

223 + handle = data.handle

224 + create_and_map(handle)

225 + if data.cpu_backup_tensor is not None:

226 + cpu_backup_tensor = data.cpu_backup_tensor

227 + if cpu_backup_tensor is not None:

228 + size_in_bytes = cpu_backup_tensor.numel(

229 + ) * cpu_backup_tensor.element_size()

230 + cpu_ptr = cpu_backup_tensor.data_ptr()

231 + libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes)

232 + data.cpu_backup_tensor = None

233

234 @contextmanager

235 def use_memory_pool(self, tag: Optional[str] = None):

vllm/engine/async_llm_engine.py

+2 -2

@@ -1225,8 +1225,8 @@ class AsyncLLMEngine(EngineClient):

1225 async def sleep(self, level: int = 1) -> None:

1226 self.engine.sleep(level)

1227

1228 - async def wake_up(self) -> None:

1229 - self.engine.wake_up()

1228 + async def wake_up(self, tags: Optional[list[str]] = None) -> None:

1229 + self.engine.wake_up(tags)

1230

1231 async def is_sleeping(self) -> bool:

1232 return self.engine.is_sleeping()

vllm/engine/llm_engine.py

+2 -2

@@ -1938,10 +1938,10 @@ class LLMEngine:

1938 "Sleep mode is not enabled in the model config")

1939 self.model_executor.sleep(level=level)

1940

1941 - def wake_up(self) -> None:

1941 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

1942 assert self.vllm_config.model_config.enable_sleep_mode, (

1943 "Sleep mode is not enabled in the model config")

1944 - self.model_executor.wake_up()

1944 + self.model_executor.wake_up(tags)

1945

1946 def is_sleeping(self) -> bool:

1947 return self.model_executor.is_sleeping

vllm/engine/multiprocessing/__init__.py

+3 -2

@@ -133,8 +133,9 @@ class RPCSleepRequest(Enum):

133 SLEEP_LEVEL_2 = 2

134

135

136 -class RPCWakeUpRequest(Enum):

137 - WAKE_UP = 1

136 +@dataclass

137 +class RPCWakeUpRequest:

138 + tags: Optional[list[str]] = None

139

140

141 @dataclass

vllm/engine/multiprocessing/client.py

+2 -2

@@ -697,10 +697,10 @@ class MQLLMEngineClient(EngineClient):

697 return await self._send_one_way_rpc_request(

698 request=RPCSleepRequest(level), socket=self.input_socket)

699

700 - async def wake_up(self) -> None:

700 + async def wake_up(self, tags: Optional[list[str]] = None) -> None:

701 """Wake up the engine"""

702 return await self._send_one_way_rpc_request(

703 - request=RPCWakeUpRequest.WAKE_UP, socket=self.input_socket)

703 + request=RPCWakeUpRequest(tags), socket=self.input_socket)

704

705 async def is_sleeping(self) -> bool:

706 """Check whether the engine is sleeping"""

vllm/engine/multiprocessing/engine.py

+3 -3

@@ -274,7 +274,7 @@ class MQLLMEngine:

274 elif isinstance(request, RPCSleepRequest):

275 self.sleep(request.value)

276 elif isinstance(request, RPCWakeUpRequest):

277 - self.wake_up()

277 + self.wake_up(request.tags)

278 elif isinstance(request, RPCIsSleepingRequest):

279 self._handle_is_sleeping_request(request)

280 else:

@@ -415,8 +415,8 @@ class MQLLMEngine:

415 def sleep(self, level: int = 1) -> None:

416 self.engine.sleep(level)

417

418 - def wake_up(self) -> None:

419 - self.engine.wake_up()

418 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

419 + self.engine.wake_up(tags)

420

421 def is_sleeping(self) -> bool:

422 return self.engine.is_sleeping()

vllm/engine/protocol.py

+1 -1

@@ -282,7 +282,7 @@ class EngineClient(ABC):

282 ...

283

284 @abstractmethod

285 - async def wake_up(self) -> None:

285 + async def wake_up(self, tags: Optional[list[str]] = None) -> None:

286 """Wake up the engine"""

287 ...

288

vllm/entrypoints/llm.py

+23 -14

@@ -1200,26 +1200,35 @@ class LLM:

1200 The caller should guarantee that no requests are being processed

1201 during the sleep period, before `wake_up` is called.

1202

1203 - :param level: The sleep level. Level 1 sleep will offload the model

1204 - weights and discard the kv cache. The content of kv cache is

1205 - forgotten. Level 1 sleep is good for sleeping and waking up the

1206 - engine to run the same model again. The model weights are backed

1207 - up in CPU memory. Please make sure there's enough CPU memory to

1208 - store the model weights. Level 2 sleep will discard both the model

1209 - weights and the kv cache. The content of both the model weights

1210 - and kv cache is forgotten. Level 2 sleep is good for sleeping and

1211 - waking up the engine to run a different model or update the model,

1212 - where previous model weights are not needed. It reduces CPU memory

1213 - pressure.

1203 + Args:

1204 + level: The sleep level. Level 1 sleep will offload the model

1205 + weights and discard the kv cache. The content of kv cache

1206 + is forgotten. Level 1 sleep is good for sleeping and waking

1207 + up the engine to run the same model again. The model weights

1208 + are backed up in CPU memory. Please make sure there's enough

1209 + CPU memory to store the model weights. Level 2 sleep will

1210 + discard both the model weights and the kv cache. The content

1211 + of both the model weights and kv cache is forgotten. Level 2

1212 + sleep is good for sleeping and waking up the engine to run a

1213 + different model or update the model, where previous model

1214 + weights are not needed. It reduces CPU memory pressure.

1215 """

1216 self.reset_prefix_cache()

1217 self.llm_engine.sleep(level=level)

1218

1218 - def wake_up(self):

1219 + def wake_up(self, tags: Optional[list[str]] = None):

1220 """

1221 Wake up the engine from sleep mode. See the :meth:`sleep` method

1221 - for more details."""

1222 - self.llm_engine.wake_up()

1222 + for more details.

1223 +

1224 + Args:

1225 + tags: An optional list of tags to reallocate the engine memory

1226 + for specific memory allocations. Values must be in

1227 + ("weights", "kv_cache",). If None, all memory is reallocated.

1228 + wake_up should be called with all tags (or None) before the

1229 + engine is used again.

1230 + """

1231 + self.llm_engine.wake_up(tags)

1232

1233 # LEGACY

1234 def _convert_v1_inputs(

vllm/entrypoints/openai/api_server.py

+6 -3

@@ -705,7 +705,6 @@ if envs.VLLM_SERVER_DEV_MODE:

705 async def sleep(raw_request: Request):

706 # get POST params

707 level = raw_request.query_params.get("level", "1")

708 - logger.info("sleep the engine with level %s", level)

708 await engine_client(raw_request).sleep(int(level))

709 # FIXME: in v0 with frontend multiprocessing, the sleep command

710 # is sent but does not finish yet when we return a response.

@@ -713,8 +712,12 @@ if envs.VLLM_SERVER_DEV_MODE:

712

713 @router.post("/wake_up")

714 async def wake_up(raw_request: Request):

716 - logger.info("wake up the engine")

717 - await engine_client(raw_request).wake_up()

715 + tags = raw_request.query_params.getlist("tags")

716 + if tags == []:

717 + # set to None to wake up all tags if no tags are provided

718 + tags = None

719 + logger.info("wake up the engine with tags: %s", tags)

720 + await engine_client(raw_request).wake_up(tags)

721 # FIXME: in v0 with frontend multiprocessing, the wake-up command

722 # is sent but does not finish yet when we return a response.

723 return Response(status_code=200)

vllm/executor/executor_base.py

+20 -5

@@ -51,6 +51,7 @@ class ExecutorBase(ABC):

51 self.observability_config = vllm_config.observability_config

52 self._init_executor()

53 self.is_sleeping = False

54 + self.sleeping_tags: set[str] = set()

55

56 @abstractmethod

57 def _init_executor(self) -> None:

@@ -204,20 +205,34 @@ class ExecutorBase(ABC):

205 time_before_sleep = time.perf_counter()

206 self.collective_rpc("sleep", kwargs=dict(level=level))

207 time_after_sleep = time.perf_counter()

208 + self.sleeping_tags = {"weights", "kv_cache"}

209 self.is_sleeping = True

210 logger.info("It took %.6f seconds to fall asleep.",

211 time_after_sleep - time_before_sleep)

212

211 - def wake_up(self):

213 + def wake_up(self, tags: Optional[list[str]] = None):

214 if not self.is_sleeping:

215 logger.warning("Executor is not sleeping.")

216 return

217 + if tags:

218 + for tag in tags:

219 + if tag not in self.sleeping_tags:

220 + logger.warning("Tag %s is not in sleeping tags %s", tag,

221 + self.sleeping_tags)

222 + return

223 time_before_wakeup = time.perf_counter()

216 - self.collective_rpc("wake_up")

224 + self.collective_rpc("wake_up", kwargs=dict(tags=tags))

225 time_after_wakeup = time.perf_counter()

218 - self.is_sleeping = False

219 - logger.info("It took %.6f seconds to wake up.",

220 - time_after_wakeup - time_before_wakeup)

226 + logger.info("It took %.6f seconds to wake up tags %s.",

227 + time_after_wakeup - time_before_wakeup,

228 + tags if tags is not None else self.sleeping_tags)

229 + if tags:

230 + for tag in tags:

231 + self.sleeping_tags.remove(tag)

232 + else:

233 + self.sleeping_tags.clear()

234 + if not self.sleeping_tags:

235 + self.is_sleeping = False

236

237 def save_sharded_state(

238 self,

vllm/v1/engine/async_llm.py

+2 -2

@@ -424,8 +424,8 @@ class AsyncLLM(EngineClient):

424 async def sleep(self, level: int = 1) -> None:

425 await self.engine_core.sleep_async(level)

426

427 - async def wake_up(self) -> None:

428 - await self.engine_core.wake_up_async()

427 + async def wake_up(self, tags: Optional[list[str]] = None) -> None:

428 + await self.engine_core.wake_up_async(tags)

429

430 async def is_sleeping(self) -> bool:

431 return await self.engine_core.is_sleeping_async()

vllm/v1/engine/core.py

+2 -2

@@ -264,8 +264,8 @@ class EngineCore:

264 def sleep(self, level: int = 1):

265 self.model_executor.sleep(level)

266

267 - def wake_up(self):

268 - self.model_executor.wake_up()

267 + def wake_up(self, tags: Optional[list[str]] = None):

268 + self.model_executor.wake_up(tags)

269

270 def is_sleeping(self) -> bool:

271 return self.model_executor.is_sleeping

vllm/v1/engine/core_client.py

+8 -8

@@ -92,7 +92,7 @@ class EngineCoreClient(ABC):

92 def sleep(self, level: int = 1) -> None:

93 raise NotImplementedError

94

95 - def wake_up(self) -> None:

95 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

96 raise NotImplementedError

97

98 def is_sleeping(self) -> bool:

@@ -141,7 +141,7 @@ class EngineCoreClient(ABC):

141 async def sleep_async(self, level: int = 1) -> None:

142 raise NotImplementedError

143

144 - async def wake_up_async(self) -> None:

144 + async def wake_up_async(self, tags: Optional[list[str]] = None) -> None:

145 raise NotImplementedError

146

147 async def is_sleeping_async(self) -> bool:

@@ -206,8 +206,8 @@ class InprocClient(EngineCoreClient):

206 def sleep(self, level: int = 1) -> None:

207 self.engine_core.sleep(level)

208

209 - def wake_up(self) -> None:

210 - self.engine_core.wake_up()

209 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

210 + self.engine_core.wake_up(tags)

211

212 def is_sleeping(self) -> bool:

213 return self.engine_core.is_sleeping()

@@ -520,8 +520,8 @@ class SyncMPClient(MPClient):

520 def sleep(self, level: int = 1) -> None:

521 self.call_utility("sleep", level)

522

523 - def wake_up(self) -> None:

524 - self.call_utility("wake_up")

523 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

524 + self.call_utility("wake_up", tags)

525

526 def is_sleeping(self) -> bool:

527 return self.call_utility("is_sleeping")

@@ -647,8 +647,8 @@ class AsyncMPClient(MPClient):

647 async def sleep_async(self, level: int = 1) -> None:

648 await self.call_utility_async("sleep", level)

649

650 - async def wake_up_async(self) -> None:

651 - await self.call_utility_async("wake_up")

650 + async def wake_up_async(self, tags: Optional[list[str]] = None) -> None:

651 + await self.call_utility_async("wake_up", tags)

652

653 async def is_sleeping_async(self) -> bool:

654 return await self.call_utility_async("is_sleeping")

vllm/v1/engine/llm_engine.py

+2 -2

@@ -245,8 +245,8 @@ class LLMEngine:

245 def sleep(self, level: int = 1):

246 self.engine_core.sleep(level)

247

248 - def wake_up(self):

249 - self.engine_core.wake_up()

248 + def wake_up(self, tags: Optional[list[str]] = None):

249 + self.engine_core.wake_up(tags)

250

251 def is_sleeping(self) -> bool:

252 return self.engine_core.is_sleeping()

vllm/v1/worker/gpu_worker.py

+2 -2

@@ -83,9 +83,9 @@ class Worker(WorkerBase):

83 "%.2f GiB memory is still in use.", freed_bytes / GiB_bytes,

84 used_bytes / GiB_bytes)

85

86 - def wake_up(self) -> None:

86 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

87 allocator = CuMemAllocator.get_instance()

88 - allocator.wake_up()

88 + allocator.wake_up(tags)

89

90 def init_device(self):

91 if self.device_config.device.type == "cuda":

vllm/worker/worker.py

+2 -2

@@ -135,9 +135,9 @@ class Worker(LocalOrDistributedWorkerBase):

135 "%.2f GiB memory is still in use.", freed_bytes / GiB_bytes,

136 used_bytes / GiB_bytes)

137

138 - def wake_up(self) -> None:

138 + def wake_up(self, tags: Optional[list[str]] = None) -> None:

139 allocator = CuMemAllocator.get_instance()

140 - allocator.wake_up()

140 + allocator.wake_up(tags=tags)

141

142 def init_device(self) -> None:

143 if self.device_config.device.type == "cuda":

[core] Add tags parameter to wake_up() (#15500)