[Bugfix][MRV2] Release model and KV cache on in-process engine shutdown (#54162)
Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Oxana Korzh committed
Aug 28, 2026 at 23:09 UTC
fb6802513850d387e13ec231a0d3aa11d2754414
4 files changed
+142
-3
tests/models/test_language_model_cache_is_weak.py
new
+59
@@ -0,0 +1,59 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The `get_language_model` cache must not own the models it keys on.
4
+
5
+A strong-keyed module-level dict pins every model ever loaded in the process
6
+for the life of the interpreter. In-process teardown
7
+(``VLLM_ENABLE_V1_MULTIPROCESSING=0``) then keeps the weights -- and the KV
8
+cache aliased onto the model's layers -- resident after engine shutdown.
9
+"""
10
+
11
+import gc
12
+import weakref
13
+
14
+import pytest
15
+import torch.nn as nn
16
+
17
+from vllm.model_executor.models.interfaces import _language_model_by_module
18
+
19
+pytestmark = pytest.mark.cpu_test
20
+
21
+
22
+class _LanguageModel(nn.Module):
23
+ def embed_input_ids(self, input_ids): # pragma: no cover - never called
24
+ raise NotImplementedError
25
+
26
+
27
+class _MultiModalModel(nn.Module):
28
+ def __init__(self):
29
+ super().__init__()
30
+ self.language_model = _LanguageModel()
31
+
32
+
33
+def test_cache_is_weak_keyed():
34
+ assert isinstance(_language_model_by_module, weakref.WeakKeyDictionary)
35
+
36
+
37
+def test_cached_entry_does_not_pin_the_model():
38
+ model = _MultiModalModel()
39
+ _language_model_by_module[model] = model.language_model
40
+ ref = weakref.ref(model)
41
+
42
+ del model
43
+ gc.collect()
44
+
45
+ assert ref() is None, (
46
+ "the get_language_model cache is keeping the model alive; in-process "
47
+ "engine shutdown will not release its weights or KV cache"
48
+ )
49
+
50
+
51
+def test_cached_entry_is_dropped_with_its_key():
52
+ model = _MultiModalModel()
53
+ _language_model_by_module[model] = model.language_model
54
+ assert model in _language_model_by_module
55
+
56
+ del model
57
+ gc.collect()
58
+
59
+ assert not any(isinstance(k, _MultiModalModel) for k in _language_model_by_module)
tests/v1/engine/test_llm_engine_finalizer_is_weak.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""`LLMEngine`'s cleanup finalizer must not own the model.
4
+
5
+``weakref.finalize`` keeps its ``args`` strongly in a class-level registry
6
+until it fires, and it only fires once the ``LLMEngine`` itself is collected.
7
+Handing it the model directly therefore pins the weights -- and the KV cache
8
+aliased onto the model's layers -- for as long as the engine is reachable,
9
+which in-process is past engine shutdown.
10
+"""
11
+
12
+import gc
13
+import weakref
14
+
15
+import pytest
16
+import torch.nn as nn
17
+
18
+from vllm.v1.engine.llm_engine import LLMEngine
19
+
20
+pytestmark = pytest.mark.cpu_test
21
+
22
+
23
+class _Model(nn.Module):
24
+ pass
25
+
26
+
27
+def test_finalize_args_are_held_strongly():
28
+ """The premise: this is why the model must be passed weakly."""
29
+
30
+ class Owner:
31
+ pass
32
+
33
+ owner, model = Owner(), _Model()
34
+ weakref.finalize(owner, lambda _m: None, model)
35
+ ref = weakref.ref(model)
36
+
37
+ del model
38
+ gc.collect()
39
+
40
+ assert ref() is not None
41
+ del owner
42
+ gc.collect()
43
+ assert ref() is None
44
+
45
+
46
+def test_cleanup_tolerates_a_dead_model():
47
+ model = _Model()
48
+ ref = weakref.ref(model)
49
+ del model
50
+ gc.collect()
51
+ assert ref() is None
52
+
53
+ LLMEngine._cleanup_instance_caches(ref)
54
+
55
+
56
+def test_cleanup_still_walks_a_live_model(monkeypatch):
57
+ from vllm.compilation import wrapper
58
+
59
+ cleaned = []
60
+
61
+ class _FakeWrapper(nn.Module):
62
+ def cleanup(self):
63
+ cleaned.append(self)
64
+
65
+ monkeypatch.setattr(
66
+ wrapper, "TorchCompileWithNoGuardsWrapper", _FakeWrapper, raising=False
67
+ )
68
+
69
+ model = _Model()
70
+ model.inner = _FakeWrapper()
71
+
72
+ LLMEngine._cleanup_instance_caches(weakref.ref(model))
73
+
74
+ assert cleaned == [model.inner]
vllm/model_executor/models/interfaces.py
+4
-1
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import weakref
6
from collections.abc import (
7
AsyncGenerator,
8
Callable,
@@ -121,7 +122,9 @@ def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor:
122
123
124
# Cache results of `SupportsMultiModal.get_language_model`
124
-_language_model_by_module = dict[nn.Module, "VllmModel"]()
125
+_language_model_by_module: "weakref.WeakKeyDictionary[nn.Module, VllmModel]" = (
126
+ weakref.WeakKeyDictionary()
127
+)
128
129
130
@runtime_checkable
vllm/v1/engine/llm_engine.py
+5
-2
@@ -129,7 +129,7 @@ class LLMEngine:
129
model = self._get_driver_model_for_cleanup()
130
if model is not None:
131
self._finalizer = weakref.finalize(
132
- self, LLMEngine._cleanup_instance_caches, model
132
+ self, LLMEngine._cleanup_instance_caches, weakref.ref(model)
133
)
134
135
if self.external_launcher_dp:
@@ -444,10 +444,13 @@ class LLMEngine:
444
return getattr(model_runner, "model", None)
445
446
@staticmethod
447
- def _cleanup_instance_caches(model) -> None:
447
+ def _cleanup_instance_caches(model_ref: "weakref.ref[nn.Module]") -> None:
448
"""Remove the bytecode hooks that pin the compiled model."""
449
from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
450
451
+ model = model_ref()
452
+ if model is None:
453
+ return
454
for module in model.modules():
455
if isinstance(module, TorchCompileWithNoGuardsWrapper):
456
module.cleanup()