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
[test][RL] Add sleep level 2 test and fix reload with sleep mode (#23521)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
22quinn committed
Aug 25, 2025 at 09:25 UTC
2a167b2eeb993638c198db49f3927bae5d55508b
2 files changed
+32
-2
tests/basic_correctness/test_cumem.py
+31
@@ -177,3 +177,34 @@ def test_end_to_end(monkeypatch: pytest.MonkeyPatch, model: str, use_v1: bool):
177
178
# cmp output
179
assert output[0].outputs[0].text == output3[0].outputs[0].text
180
+
181
+
182
+@create_new_process_for_each_test()
183
+def test_deep_sleep():
184
+ model = "Qwen/Qwen3-0.6B"
185
+ free, total = torch.cuda.mem_get_info()
186
+ used_bytes_baseline = total - free # in case other process is running
187
+ llm = LLM(model, enable_sleep_mode=True)
188
+ prompt = "How are you?"
189
+ sampling_params = SamplingParams(temperature=0, max_tokens=10)
190
+ output = llm.generate(prompt, sampling_params)
191
+
192
+ # Put the engine to deep sleep
193
+ llm.sleep(level=2)
194
+
195
+ free_gpu_bytes_after_sleep, total = torch.cuda.mem_get_info()
196
+ used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
197
+ assert used_bytes < 3 * GiB_bytes
198
+
199
+ llm.wake_up(tags=["weights"])
200
+ llm.collective_rpc("reload_weights")
201
+ free_gpu_bytes_wake_up_w, total = torch.cuda.mem_get_info()
202
+ used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
203
+ assert used_bytes < 4 * GiB_bytes
204
+
205
+ # now allocate kv cache and cuda graph memory
206
+ llm.wake_up(tags=["kv_cache"])
207
+ output2 = llm.generate(prompt, sampling_params)
208
+
209
+ # cmp output
210
+ assert output[0].outputs[0].text == output2[0].outputs[0].text
vllm/v1/worker/gpu_worker.py
+1
-2
@@ -216,8 +216,7 @@ class Worker(WorkerBase):
216
self.model_runner.update_config(overrides)
217
218
def reload_weights(self) -> None:
219
- with self._maybe_get_memory_pool_context(tag="weights"):
220
- self.model_runner.reload_weights()
219
+ self.model_runner.reload_weights()
220
221
@torch.inference_mode()
222
def determine_available_memory(self) -> int: