[ROCm][CI] Query HIP device memory for test GPU teardown waits. (#57450)
Signed-off-by: Sheral Kumar <shekumar@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Sheral Kumar committed
Sep 20, 2026 at 16:20 UTC
01f1f58f10c616850a5f5ade129a7b58d62075fb
4 files changed
+56
-53
tests/compile/passes/conftest.py
new
+22
@@ -0,0 +1,22 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager
9
+
10
+
11
+@pytest.fixture(autouse=True)
12
+def _workspace_manager_for_compile_passes(cleanup_fixture):
13
+ if not current_platform.is_rocm():
14
+ yield
15
+ return
16
+
17
+ # Release workspace tensors before cleanup_fixture flushes the allocator.
18
+ reset_workspace_manager()
19
+ if torch.accelerator.is_available():
20
+ init_workspace_manager(torch.device(0))
21
+ yield
22
+ reset_workspace_manager()
tests/conftest.py
+2
-5
@@ -901,13 +901,10 @@ class HfRunner:
901
# tests/basic_correctness/test_basic_correctness.py::test_models_distributed
902
# where vllm worker processes are still alive and holding GPU
903
# memory when hf_runner.__exit__ is called.
904
- from tests.utils import (
905
- get_physical_device_indices,
906
- record_gpu_memory_usage_stats,
907
- )
904
+ from tests.utils import record_gpu_memory_usage_stats
905
906
if (device_count := current_platform.device_count()) > 0:
910
- devices = get_physical_device_indices(devices=list(range(device_count)))
907
+ devices = list(range(device_count))
908
mem_usage_stats = record_gpu_memory_usage_stats(devices=devices)
909
self.threshold_ratios = {
910
device: 0.05 + mem_used / mem_tot
tests/kernels/moe/conftest.py
+2
@@ -3,6 +3,7 @@
3
import pytest
4
5
from vllm.distributed import cleanup_dist_env_and_memory
6
+from vllm.v1.worker.workspace import reset_workspace_manager
7
8
NEEDS_CLEAN_ENTRY = frozenset(
9
{
@@ -58,4 +59,5 @@ def pytest_runtest_setup(item):
59
# hook, so tests about to be skipped never reach this, while fixture setup
60
# still runs afterwards.
61
if item.path.name in NEEDS_CLEAN_ENTRY:
62
+ reset_workspace_manager()
63
cleanup_dist_env_and_memory()
tests/utils.py
+30
-48
@@ -87,25 +87,7 @@ def prewarm_hf_cache(assets: list[tuple[str, str]]) -> None:
87
)
88
89
90
-if current_platform.is_rocm():
91
- from amdsmi import (
92
- amdsmi_get_gpu_vram_usage,
93
- amdsmi_get_processor_handles,
94
- amdsmi_init,
95
- amdsmi_shut_down,
96
- )
97
-
98
- _amdsmi_lock = threading.Lock()
99
-
100
- @contextmanager
101
- def _nvml():
102
- with _amdsmi_lock:
103
- try:
104
- amdsmi_init()
105
- yield
106
- finally:
107
- amdsmi_shut_down()
108
-elif current_platform.is_cuda():
90
+if current_platform.is_cuda():
91
from vllm.third_party.pynvml import (
92
nvmlDeviceGetHandleByIndex,
93
nvmlDeviceGetHandleByUUID,
@@ -595,27 +577,20 @@ class RemoteVLLMServer:
577
return members
578
579
def _get_gpu_memory_used(self) -> float | None:
598
- """Get total GPU memory used across all visible devices in bytes."""
580
+ """Get device-wide usage across visible devices in bytes.
581
+
582
+ On ROCm this initializes HIP in the caller; subsequent multiprocessing
583
+ GPU workers must use spawn, as the ROCm server helpers already do.
584
+ """
585
try:
586
if current_platform.is_rocm():
601
- with _nvml():
602
- handles = amdsmi_get_processor_handles()
603
- devices = get_physical_device_indices(
604
- list(range(current_platform.device_count()))
605
- )
606
- total_used_mib = 0
607
- for device in devices:
608
- handle = handles[device]
609
- vram_info = amdsmi_get_gpu_vram_usage(handle)
610
- total_used_mib += vram_info["vram_used"]
611
- # amdsmi reports VRAM in MiB; convert to bytes so this
612
- # matches the CUDA/nvml branch (already bytes) and the
613
- # byte-based target in _wait_for_gpu_memory_release. Without
614
- # this, that wait compares MiB against a ~2e9-byte target,
615
- # is always satisfied instantly, and returns "released to
616
- # 0.00 GB" while the previous server's VRAM is still
617
- # resident -- OOMing the next server's startup on ROCm.
618
- return total_used_mib * 1024 * 1024
587
+ # Use HIP logical devices. On MI355 DPX/NPS2, the primary
588
+ # AMD SMI device can expose whole-card memory counters.
589
+ total_used = 0
590
+ for i in range(current_platform.device_count()):
591
+ free, total = torch.accelerator.get_memory_info(i)
592
+ total_used += total - free
593
+ return float(total_used)
594
elif current_platform.is_cuda():
595
with _nvml():
596
total_used = 0
@@ -653,9 +628,12 @@ class RemoteVLLMServer:
628
# Can't query GPU memory - nothing to do
629
return
630
656
- # Allow up to 2 GiB overhead above baseline for driver/context state
657
- # that may persist between server instances.
658
- headroom_bytes = 2 * 1024 * 1024 * 1024
631
+ # Allow aggregate driver/context growth above the baseline.
632
+ headroom_bytes = (
633
+ 4 * 1024 * 1024 * 1024
634
+ if current_platform.is_rocm()
635
+ else 2 * 1024 * 1024 * 1024
636
+ )
637
target = baseline + headroom_bytes
638
639
start = time.time()
@@ -1616,13 +1594,17 @@ def record_gpu_memory_usage_stats(
1594
*,
1595
devices: list[int],
1596
) -> dict[int, tuple[float, float]]:
1597
+ """Return device-wide used/total GiB; ROCm IDs are HIP logical devices.
1598
+
1599
+ ROCm queries initialize HIP in the caller and include other processes'
1600
+ allocations on the same memory partition.
1601
+ """
1602
output: dict[int, tuple[float, float]] = {}
1603
for device in devices:
1604
if current_platform.is_rocm():
1622
- dev_handle = amdsmi_get_processor_handles()[device]
1623
- mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
1624
- gb_used = mem_info["vram_used"] / 2**10
1625
- gb_total = mem_info["vram_total"] / 2**10
1605
+ free, total = torch.accelerator.get_memory_info(device)
1606
+ gb_used = (total - free) / 2**30
1607
+ gb_total = total / 2**30
1608
elif current_platform.is_xpu():
1609
# nvml/amdsmi are unavailable on XPU. Query device memory through
1610
# torch.accelerator.get_memory_info, which the XPU platform patches
@@ -1650,7 +1632,9 @@ def wait_for_gpu_memory_to_clear(
1632
poll_interval_s: float = 5,
1633
) -> None:
1634
assert threshold_bytes is not None or threshold_ratio is not None
1653
- devices = get_physical_device_indices(devices)
1635
+ # HIP already applies device visibility; keep logical IDs and threshold keys.
1636
+ if not current_platform.is_rocm():
1637
+ devices = get_physical_device_indices(devices)
1638
if isinstance(threshold_bytes, int):
1639
threshold_bytes = {device: threshold_bytes for device in devices}
1640
elif isinstance(threshold_bytes, dict):
@@ -1672,8 +1656,6 @@ def wait_for_gpu_memory_to_clear(
1656
threshold_bytes.get(device, 0), min_threshold_b if ratio < 0.05 else 0
1657
)
1658
1675
- # Use nvml instead of pytorch to reduce measurement error from torch cuda
1676
- # context.
1659
start_time = time.time()
1660
stable_since: float | None = None
1661
stable_used_bytes: dict[int, int] | None = None