[Bugfix][MRV2] Release layer-bound KV cache memory in shutdown() (#54246)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Kimi Code <noreply@moonshot.ai>
Nick Hill committed
Aug 28, 2026 at 14:27 UTC
96242aa50d196993735854f03bd717496e88b7c3
4 files changed
+42
-13
tests/conftest.py
+4
-2
@@ -1371,8 +1371,10 @@ class VllmRunner:
1371
shutdown_timeout = 60.0 if current_platform.is_rocm() else None
1372
self.llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout)
1373
except Exception:
1374
- # Ignore shutdown errors as cleanup will still proceed
1375
- pass
1374
+ # Don't fail the test on shutdown errors since cleanup will still
1375
+ # proceed, but don't hide them either: a failure here usually
1376
+ # means the engine's GPU memory was never released.
1377
+ logger.exception("Engine core shutdown raised; GPU memory may leak")
1378
del self.llm
1379
torch._dynamo.reset()
1380
cleanup_dist_env_and_memory()
vllm/v1/worker/gpu/cudagraph_utils.py
+9
-10
@@ -1,8 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import gc
4
+import itertools
5
from collections import defaultdict
5
-from collections.abc import Callable
6
+from collections.abc import Callable, Iterable
7
from dataclasses import dataclass
8
from itertools import groupby, product
9
from typing import TYPE_CHECKING, Any, NamedTuple, Protocol
@@ -38,7 +39,7 @@ from vllm.v1.worker.gpu.block_table import BlockTables
39
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
40
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
41
from vllm.v1.worker.gpu.model_states.interface import ModelState
41
-from vllm.v1.worker.utils import AttentionGroup
42
+from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches
43
44
if TYPE_CHECKING:
45
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
@@ -872,14 +873,12 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
873
# capture_model() re-captures them.
874
if runner.model_state.supports_mm_inputs:
875
runner.model_state.encoder_runner.clear()
875
- # Detach profiling KV tensors held by attention layers.
876
- for layer in runner.compilation_config.static_forward_context.values():
877
- if hasattr(layer, "kv_cache"):
878
- kv_cache = layer.kv_cache
879
- layer.kv_cache = (
880
- torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []
881
- )
882
- del kv_cache
876
+ # Detach profiling KV tensors held by attention layers. The layers live
877
+ # in the static forward context for compiled models.
878
+ layers: Iterable[Any] = runner.compilation_config.static_forward_context.values()
879
+ if (model := getattr(runner, "model", None)) is not None:
880
+ layers = itertools.chain(layers, model.modules())
881
+ clear_layer_kv_caches(layers)
882
runner.cache_config.num_gpu_blocks = None
883
runner.maybe_remove_all_loras(runner.lora_config)
884
gc.collect()
vllm/v1/worker/gpu/model_runner.py
+8
-1
@@ -156,6 +156,7 @@ from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
156
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
157
from vllm.v1.worker.utils import (
158
KVBlockZeroer,
159
+ clear_layer_kv_caches,
160
copy_kv_cache_blocks_inplace,
161
get_uniform_decode_token_count,
162
)
@@ -2026,9 +2027,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2027
free_before_shutdown(self.vllm_config)
2028
if hasattr(self, "model_state"):
2029
del self.model_state
2029
- if getattr(self, "speculator", None) is not None:
2030
+ # Detach the layer-level KV/state cache tensors before dropping the
2031
+ # models; the model objects can outlive this runner.
2032
+ speculator = getattr(self, "speculator", None)
2033
+ if speculator is not None:
2034
+ if draft_model := getattr(speculator, "model", None):
2035
+ clear_layer_kv_caches(draft_model.modules())
2036
self.speculator = None
2037
if hasattr(self, "model"):
2038
+ clear_layer_kv_caches(self.model.modules())
2039
del self.model
2040
2041
gc.collect()
vllm/v1/worker/utils.py
+21
@@ -622,6 +622,27 @@ def bind_kv_cache(
622
forward_context[layer_name].bind_kv_cache(kv_cache)
623
624
625
+def clear_layer_kv_caches(layers: Iterable[Any]) -> None:
626
+ """Detach the KV/state cache tensors installed by bind_kv_cache().
627
+
628
+ The model object can outlive the runner (e.g. LLMEngine's finalizer keeps
629
+ it reachable until engine deletion), so dropping the runner's references
630
+ alone does not release the KV cache memory on teardown paths.
631
+ """
632
+ for layer in layers:
633
+ if not hasattr(layer, "kv_cache"):
634
+ continue
635
+ kv_cache = layer.kv_cache
636
+ layer.kv_cache = torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []
637
+ # Clean up quantized KV cache scale views
638
+ # (int8_per_token_head, fp8_per_token_head)
639
+ if hasattr(layer, "impl"):
640
+ if hasattr(layer.impl, "_k_scale_cache"):
641
+ layer.impl._k_scale_cache = None
642
+ if hasattr(layer.impl, "_v_scale_cache"):
643
+ layer.impl._v_scale_cache = None
644
+
645
+
646
def copy_kv_cache_blocks_inplace(
647
kv_caches: Iterable[torch.Tensor],
648
num_blocks: int,