@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 95995bbef8
Fork 0

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, 15:00 UTC

[ROCm][Engine] Fix GPU memory leaks in engine shutdown and test workaround for async KV prefix cache reset (#38503)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Apr 25, 2026 at 00:25 UTC 95995bbef81292e3ee1ef0df5ca3989bb481bdd5
8 files changed +300 -13
docker/Dockerfile.rocm
+14 -5
@@ -124,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
124
125 # RIXL/UCX build stages
126 FROM base AS build_rixl
127 -ARG RIXL_BRANCH="f33a5599"
127 +ARG RIXL_BRANCH="bf4a7214"
128 ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
129 -ARG UCX_BRANCH="da3fac2a"
130 -ARG UCX_REPO="https://github.com/ROCm/ucx.git"
129 +ARG UCX_BRANCH="7009d7a1"
130 +ARG UCX_REPO="https://github.com/openucx/ucx.git"
131 ENV ROCM_PATH=/opt/rocm
132 ENV UCX_HOME=/usr/local/ucx
133 ENV RIXL_HOME=/usr/local/rixl
@@ -165,7 +165,7 @@ RUN cd /usr/local/src && \
165 --disable-doxygen-doc \
166 --enable-optimizations \
167 --enable-devel-headers \
168 - --with-rocm=/opt/rocm \
168 + --with-rocm=${ROCM_PATH} \
169 --with-verbs \
170 --with-dm \
171 --enable-mt && \
@@ -186,7 +186,12 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
186 ninja install
187
188 # Generate RIXL wheel
189 -RUN cd /opt/rixl && mkdir -p /app/install && \
189 +# Exclude libcore and libpull from auditwheel: transitive dependencies
190 +# that are not shipped in the wheel and vary across base images.
191 +RUN cd /opt/rixl && \
192 + sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
193 + contrib/build-wheel.sh && \
194 + mkdir -p /app/install && \
195 ./contrib/build-wheel.sh \
196 --output-dir /app/install \
197 --rocm-dir ${ROCM_PATH} \
@@ -431,6 +436,10 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
436 ENV MIOPEN_DEBUG_CONV_DIRECT=0
437 ENV MIOPEN_DEBUG_CONV_GEMM=0
438
439 +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
440 +# See: https://github.com/ROCm/rocm-libraries/issues/6266
441 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1
442 +
443 # Source code is used in the `python_only_compile.sh` test
444 # We hide it inside `src/` so that this source code
445 # will not be imported by other tests
tests/entrypoints/openai/completion/test_shutdown.py
+1 -1
@@ -311,7 +311,7 @@ async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float):
311 pytest.fail("Process did not exit after SIGTERM with abort timeout")
312
313 exit_time = time.time() - start_time
314 - assert exit_time < 2, f"Default shutdown took too long: {exit_time:.1f}s"
314 + assert exit_time < 2.1, f"Default shutdown took too long: {exit_time:.1f}s"
315 assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
316
317 await _assert_children_cleaned_up(child_pids)
tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py new
+222
@@ -0,0 +1,222 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Verify that GPU memory is fully released after RixlConnector shutdown on ROCm.
4 +
5 +Regression test for ROCm/ucx#33: UCX rocm_ipc transport permanently pinned
6 +GPU memory via hsa_amd_ipc_memory_create during ucp_mem_map, causing
7 +GPU memory to be unrecoverable after engine shutdown.
8 +"""
9 +
10 +import gc
11 +
12 +import pytest
13 +import torch
14 +
15 +from vllm.platforms import current_platform
16 +
17 +pytestmark = pytest.mark.skipif(
18 + not current_platform.is_rocm(),
19 + reason="ROCm platform required",
20 +)
21 +
22 +
23 +def _mb(b: int) -> float:
24 + return b / (1024 * 1024)
25 +
26 +
27 +def _gpu_snapshot(tag: str, prev_alloc: float = 0.0) -> dict:
28 + """Print and return current GPU memory stats."""
29 + torch.accelerator.synchronize()
30 + alloc = torch.accelerator.memory_allocated()
31 + reserved = torch.accelerator.memory_reserved()
32 + # mem_get_info is not available on torch.accelerator
33 + try:
34 + drv_free, drv_total = torch.cuda.mem_get_info()
35 + drv_used = drv_total - drv_free
36 + drv_pct = drv_used / drv_total * 100
37 + except Exception:
38 + drv_used = drv_total = drv_pct = 0
39 + alloc_mb = _mb(alloc)
40 + drv_used_mb = _mb(drv_used)
41 + delta = alloc_mb - prev_alloc
42 + print(
43 + f" {tag:<40s} | {alloc_mb:>9.1f} alloc | "
44 + f"{_mb(reserved):>9.1f} rsrvd | "
45 + f"{drv_used_mb:>9.1f} driver ({drv_pct:.1f}%) | "
46 + f"delta {delta:>+9.1f}"
47 + )
48 + return {
49 + "tag": tag,
50 + "alloc_mb": alloc_mb,
51 + "drv_used_mb": drv_used_mb,
52 + "drv_pct": drv_pct,
53 + }
54 +
55 +
56 +def _full_gpu_cleanup():
57 + """gc.collect + torch empty_cache, multiple rounds."""
58 + gc.unfreeze()
59 + for _ in range(3):
60 + if gc.collect() == 0:
61 + break
62 + torch.accelerator.empty_cache()
63 +
64 +
65 +@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
66 +def test_gpu_memory_rixl_hma(model_name, sw_size):
67 + """Track GPU memory through NixlConnector create/infer/shutdown cycle."""
68 + from vllm import LLM, SamplingParams
69 + from vllm.config import KVTransferConfig
70 + from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
71 +
72 + llm_kwargs = {
73 + "model": model_name,
74 + "enforce_eager": True,
75 + "gpu_memory_utilization": 0.5,
76 + "kv_transfer_config": KVTransferConfig(
77 + kv_connector="NixlConnector",
78 + kv_role="kv_both",
79 + ),
80 + "max_model_len": 2048,
81 + "disable_hybrid_kv_cache_manager": False,
82 + "max_num_batched_tokens": 1024,
83 + "enable_prefix_caching": False,
84 + "block_size": 16,
85 + }
86 +
87 + print("\n" + "=" * 90)
88 + print("GPU MEMORY -- RIXL NixlConnector HMA (ROCm)")
89 + print("=" * 90)
90 + gc.collect()
91 + torch.accelerator.empty_cache()
92 + torch.accelerator.reset_peak_memory_stats()
93 + snap0 = _gpu_snapshot("0. baseline", 0.0)
94 +
95 + # create + infer
96 + llm = LLM(**llm_kwargs)
97 + snap1 = _gpu_snapshot("1. after LLM()", snap0["alloc_mb"])
98 +
99 + llm.generate(
100 + ["hi" * 1401],
101 + SamplingParams(
102 + temperature=0.0,
103 + max_tokens=1,
104 + extra_args={
105 + "kv_transfer_params": {
106 + "do_remote_decode": True,
107 + "do_remote_prefill": False,
108 + "remote_engine_id": None,
109 + "remote_block_ids": None,
110 + "remote_host": None,
111 + "remote_port": None,
112 + }
113 + },
114 + ),
115 + )
116 + snap2 = _gpu_snapshot("2. after generate()", snap1["alloc_mb"])
117 +
118 + # shutdown + cleanup
119 + print("\n--- shutdown ---")
120 + llm.llm_engine.engine_core.shutdown()
121 + _gpu_snapshot("3. after shutdown()", snap2["alloc_mb"])
122 +
123 + del llm
124 + _full_gpu_cleanup()
125 + cleanup_dist_env_and_memory()
126 + _full_gpu_cleanup()
127 + torch._dynamo.reset()
128 + gc.collect()
129 + torch.accelerator.empty_cache()
130 + snap_final = _gpu_snapshot("4. final", snap2["alloc_mb"])
131 +
132 + # summary
133 + print("\n" + "=" * 90)
134 + baseline = snap0["alloc_mb"]
135 + final = snap_final["alloc_mb"]
136 + peak = snap2["alloc_mb"]
137 + total_alloc = peak - baseline
138 +
139 + print(
140 + f" PyTorch: baseline={baseline:.0f} peak={peak:.0f} "
141 + f"final={final:.0f} "
142 + f"leaked={final - baseline:.0f} MB"
143 + + (
144 + f" ({(final - baseline) / total_alloc * 100:.1f}%)"
145 + if total_alloc > 0
146 + else ""
147 + )
148 + )
149 +
150 + drv_base = snap0["drv_used_mb"]
151 + drv_final = snap_final["drv_used_mb"]
152 + drv_leaked = drv_final - drv_base
153 + print(
154 + f" Driver: baseline={drv_base:.0f} ({snap0['drv_pct']:.1f}%) "
155 + f"peak={snap2['drv_used_mb']:.0f} ({snap2['drv_pct']:.1f}%) "
156 + f"final={drv_final:.0f} ({snap_final['drv_pct']:.1f}%) "
157 + f"leaked={drv_leaked:.0f} MB"
158 + )
159 + print("=" * 90)
160 +
161 + # Peak driver memory used above baseline
162 + drv_peak = snap2["drv_used_mb"] - drv_base
163 + leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
164 + max_leak_pct = 10
165 + assert leak_pct <= max_leak_pct, (
166 + f"{drv_leaked:.0f} MB ({leak_pct:.1f}%) of driver-level GPU memory "
167 + f"not freed after NixlConnector shutdown "
168 + f"(peak allocation: {drv_peak:.0f} MB, threshold: {max_leak_pct}%)"
169 + )
170 +
171 +
172 +@pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"])
173 +def test_gpu_memory_no_rixl_baseline(model_name):
174 + """Same workload without NixlConnector. Comparing driver-level memory
175 + between this and test_gpu_memory_rixl_hma isolates UCX/RIXL impact."""
176 + from vllm import LLM, SamplingParams
177 + from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
178 +
179 + print("\n" + "=" * 90)
180 + print("CONTROL -- same model, no RIXL connector")
181 + print("=" * 90)
182 + gc.collect()
183 + torch.accelerator.empty_cache()
184 + snap0 = _gpu_snapshot("baseline", 0.0)
185 +
186 + llm = LLM(
187 + model=model_name,
188 + enforce_eager=True,
189 + gpu_memory_utilization=0.5,
190 + max_model_len=2048,
191 + max_num_batched_tokens=1024,
192 + enable_prefix_caching=False,
193 + block_size=16,
194 + )
195 + _gpu_snapshot("after LLM()", snap0["alloc_mb"])
196 +
197 + llm.generate(["hi " * 500], SamplingParams(max_tokens=1))
198 + snap_peak = _gpu_snapshot("after generate()", snap0["alloc_mb"])
199 +
200 + llm.llm_engine.engine_core.shutdown()
201 + del llm
202 + _full_gpu_cleanup()
203 + cleanup_dist_env_and_memory()
204 + _full_gpu_cleanup()
205 + torch._dynamo.reset()
206 + gc.collect()
207 + torch.accelerator.empty_cache()
208 + snap_final = _gpu_snapshot("final", snap0["alloc_mb"])
209 +
210 + drv_base = snap0["drv_used_mb"]
211 + drv_leaked = snap_final["drv_used_mb"] - drv_base
212 + drv_peak = snap_peak["drv_used_mb"] - drv_base
213 + print(f"\n Driver leaked (no rixl): {drv_leaked:.0f} MB")
214 + print("=" * 90)
215 +
216 + leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
217 + max_leak_pct = 10
218 + assert leak_pct <= max_leak_pct, (
219 + f"{drv_leaked:.0f} MB ({leak_pct:.1f}%) of driver-level GPU memory "
220 + f"not freed after baseline shutdown "
221 + f"(peak allocation: {drv_peak:.0f} MB, threshold: {max_leak_pct}%)"
222 + )
tests/v1/kv_offload/test_cpu_offloading.py
+12 -3
@@ -87,10 +87,13 @@ class MockSubscriber:
87 def _wait_for_prefix_cache_reset(llm: LLM) -> None:
88 """Wait for async offload transfers to finish so prefix cache can reset.
89
90 - The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks
90 + The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks
91 are still held by the offload worker, ``reset_prefix_cache`` returns
92 - ``False``. Retry with a short sleep until it succeeds or we time out.
92 + ``False``. Between retries we send a dummy single-token prefill to force
93 + the engine to step, which polls the worker for completed transfers and
94 + frees GPU blocks.
95 """
96 + _dummy_params = SamplingParams(max_tokens=1)
97 deadline = time.monotonic() + _RESET_CACHE_TIMEOUT
98 while not llm.reset_prefix_cache():
99 if time.monotonic() > deadline:
@@ -98,7 +101,13 @@ def _wait_for_prefix_cache_reset(llm: LLM) -> None:
101 "reset_prefix_cache did not succeed within "
102 f"{_RESET_CACHE_TIMEOUT}s - async offload may be stuck"
103 )
101 - time.sleep(0.1)
104 + # Force an engine step so the scheduler polls get_finished()
105 + # and releases GPU blocks held by in-flight async stores.
106 + llm.generate(
107 + [TokensPrompt(prompt_token_ids=[0])],
108 + _dummy_params,
109 + use_tqdm=False,
110 + )
111
112
113 def _latency_test(llm: LLM, subscriber: MockSubscriber):
vllm/utils/func_utils.py
+25 -4
@@ -45,16 +45,14 @@ def run_once(f: Callable[P, None]) -> Callable[P, None]:
45
46
47 @lru_cache
48 -def supports_kw(
48 +def _supports_kw(
49 callable: Callable[..., object],
50 kw_name: str,
51 *,
52 requires_kw_only: bool = False,
53 allow_var_kwargs: bool = True,
54 ) -> bool:
55 - """Check if a keyword is a valid kwarg for a callable; if requires_kw_only
56 - disallows kwargs names that can also be positional arguments.
57 - """
55 + """Internal cached implementation of supports_kw."""
56 params = inspect.signature(callable).parameters
57 if not params:
58 return False
@@ -99,6 +97,29 @@ def supports_kw(
97 return False
98
99
100 +def supports_kw(
101 + callable: Callable[..., object],
102 + kw_name: str,
103 + *,
104 + requires_kw_only: bool = False,
105 + allow_var_kwargs: bool = True,
106 +) -> bool:
107 + """Check if a keyword is a valid kwarg for a callable; if requires_kw_only
108 + disallows kwargs names that can also be positional arguments.
109 + """
110 + # Unwrap bound methods so that the lru_cache key is the underlying
111 + # function, not the instance. Caching bound methods pins the object
112 + # (and all its GPU tensors) for the lifetime of the cache.
113 + if hasattr(callable, "__func__"):
114 + callable = callable.__func__
115 + return _supports_kw(
116 + callable,
117 + kw_name,
118 + requires_kw_only=requires_kw_only,
119 + allow_var_kwargs=allow_var_kwargs,
120 + )
121 +
122 +
123 def get_allowed_kwarg_only_overrides(
124 callable: Callable[..., object],
125 overrides: Mapping[str, object] | None,
vllm/v1/engine/core.py
+7
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import gc
4 import os
5 import queue
6 import signal
@@ -574,6 +575,12 @@ class EngineCore:
575 if self.scheduler:
576 self.scheduler.shutdown()
577
578 + # Undo the gc.freeze() from __init__ so that the objects allocated
579 + # during engine startup (model weights, KV caches, etc.) become
580 + # visible to the garbage collector again. Without this, deleting
581 + # the engine in-process (e.g. unit tests) leaks GPU memory.
582 + gc.unfreeze()
583 +
584 def profile(self, is_start: bool = True, profile_prefix: str | None = None):
585 self.model_executor.profile(is_start, profile_prefix)
586
vllm/v1/worker/gpu_model_runner.py
+14
@@ -5886,6 +5886,20 @@ class GPUModelRunner(
5886 gc.unfreeze()
5887 gc.collect()
5888
5889 + def shutdown(self) -> None:
5890 + """Release GPU tensors (model weights, KV caches, workspace) so that
5891 + memory is reclaimable when running in the same process."""
5892 + from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT
5893 + from vllm.v1.worker.workspace import reset_workspace_manager
5894 +
5895 + # Calls torch.accelerator.synchronize()
5896 + self._cleanup_profiling_kv_cache()
5897 + self.compilation_config.static_forward_context.clear()
5898 + self.model = None # type: ignore[assignment]
5899 + _ROPE_DICT.clear()
5900 +
5901 + reset_workspace_manager()
5902 +
5903 def _cleanup_profiling_kv_cache(self) -> None:
5904 torch.accelerator.synchronize()
5905 if hasattr(self, "kv_caches") and self.kv_caches:
vllm/v1/worker/gpu_worker.py
+5
@@ -1015,6 +1015,11 @@ class Worker(WorkerBase):
1015 if weight_transfer_engine := getattr(self, "weight_transfer_engine", None):
1016 weight_transfer_engine.shutdown()
1017
1018 + # Release GPU resources held by the model runner so that memory
1019 + # can be reclaimed when running in-process
1020 + if model_runner := getattr(self, "model_runner", None):
1021 + model_runner.shutdown()
1022 +
1023 def elastic_ep_execute(self, execute_method: str, *args, **kwargs):
1024 return self.elastic_ep_executor.execute(execute_method, *args, **kwargs)
1025