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 4, 08:45 UTC
[MRV2] Add shutdown() method (#41297)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 2, 2026 at 21:06 UTC
e6ff3e9c83a6520c3793f4e0511ac8591a07c243
2 files changed
+39
vllm/v1/worker/gpu/model_runner.py
+19
@@ -91,6 +91,7 @@ from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
91
from vllm.v1.worker.gpu.sample.output import SamplerOutput
92
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
93
from vllm.v1.worker.gpu.sample.sampler import Sampler
94
+from vllm.v1.worker.gpu.shutdown import free_before_shutdown
95
from vllm.v1.worker.gpu.spec_decode import init_speculator
96
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
97
set_eagle3_aux_hidden_state_layers,
@@ -1339,6 +1340,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1340
input_batch.num_scheduled_tokens
1341
)
1342
1343
+ def shutdown(self) -> None:
1344
+ """Release GPU tensors (model weights, KV caches, workspace) so that
1345
+ memory is reclaimable when running in the same process."""
1346
+ torch.accelerator.synchronize()
1347
+ if hasattr(self, "kv_caches"):
1348
+ self.kv_caches.clear()
1349
+ if hasattr(self, "attn_groups"):
1350
+ self.attn_groups.clear()
1351
+ if hasattr(self, "kv_cache_config"):
1352
+ del self.kv_cache_config
1353
+ free_before_shutdown(self.vllm_config)
1354
+ if hasattr(self, "model"):
1355
+ del self.model
1356
+
1357
+ gc.collect()
1358
+ torch.accelerator.empty_cache()
1359
+ logger.debug("Cleaned up model weights, KV caches, and workspace")
1360
+
1361
########### EPLB methods start ###########
1362
@property
1363
def eplb_state(self):
vllm/v1/worker/gpu/shutdown.py
new
+20
@@ -0,0 +1,20 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from vllm.config import VllmConfig
4
+from vllm.logger import init_logger
5
+
6
+logger = init_logger(__name__)
7
+
8
+
9
+def free_before_shutdown(vllm_config: VllmConfig) -> None:
10
+ from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT
11
+ from vllm.v1.worker.workspace import reset_workspace_manager
12
+
13
+ cache_config = vllm_config.cache_config
14
+ cache_config.num_gpu_blocks = None
15
+
16
+ compilation_config = vllm_config.compilation_config
17
+ compilation_config.static_forward_context.clear()
18
+
19
+ _ROPE_DICT.clear()
20
+ reset_workspace_manager()