[ROCm][CI] Keep startup profiling from aborting when free memory grows (#53591)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
stefankoncarevic committed
Aug 28, 2026 at 09:02 UTC
9236159bb60e0a34dbf442a90f0f813780f2eb87
2 files changed
+130
-14
tests/v1/worker/test_gpu_worker.py
+76
-1
@@ -7,7 +7,8 @@ from unittest.mock import patch
7
import pytest
8
9
from vllm.utils.mem_constants import GiB_bytes
10
-from vllm.v1.worker import startup_plan
10
+from vllm.v1.worker import gpu_worker, startup_plan
11
+from vllm.v1.worker.gpu_worker import maybe_rocm_profiling_fallback
12
from vllm.v1.worker.startup_plan import (
13
maybe_apply_startup_plan,
14
maybe_save_startup_plan,
@@ -77,3 +78,77 @@ def test_startup_plan_apply_gate(plan_env):
78
explicit = _plan_worker(kv_bytes=7 * GiB_bytes)
79
maybe_apply_startup_plan(explicit)
80
assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes
81
+
82
+
83
+# Memory accounting of the profiling run (Worker.determine_available_memory).
84
+
85
+# The fallback reads only the sign of the measured drop and this process's torch
86
+# reservation; free memory is only logged, so no amount here is a device size.
87
+ANY_FREE_MEMORY = 8 * GiB_bytes
88
+MEASURED_DROP = 4 * GiB_bytes
89
+TORCH_RESERVED = 3 * GiB_bytes
90
+RELEASED_BY_OTHERS = 2 * GiB_bytes
91
+
92
+
93
+def _snapshot(free_memory, torch_memory=0):
94
+ return SimpleNamespace(free_memory=free_memory, torch_memory=torch_memory)
95
+
96
+
97
+def _profile_result(consumed, reserved_before=0, reserved_after=0):
98
+ """A result whose free-memory readings agree with `consumed`, which
99
+ `memory_profiling` derives as the drop in free memory, negative when it grew."""
100
+ return SimpleNamespace(
101
+ total_consumed=consumed,
102
+ transient_peak_headroom=0,
103
+ before_create=_snapshot(ANY_FREE_MEMORY, reserved_before),
104
+ after_profile=_snapshot(ANY_FREE_MEMORY - consumed, reserved_after),
105
+ )
106
+
107
+
108
+@pytest.fixture
109
+def rocm(request):
110
+ with patch.object(
111
+ gpu_worker, "current_platform", SimpleNamespace(is_rocm=lambda: request.param)
112
+ ):
113
+ yield request.param
114
+
115
+
116
+@pytest.mark.parametrize("rocm", [True, False], indirect=True)
117
+def test_profiling_fallback_declines_when_free_memory_dropped(rocm):
118
+ """The profiling measurement is kept as-is whenever free memory dropped."""
119
+ result = _profile_result(consumed=MEASURED_DROP)
120
+
121
+ assert maybe_rocm_profiling_fallback(result) is None
122
+
123
+
124
+@pytest.mark.parametrize("rocm", [True], indirect=True)
125
+def test_profiling_fallback_replaces_a_released_measurement(rocm):
126
+ """A negative measurement describes the rest of the device, so it is replaced
127
+ by this process's reservation, which the rest of the device cannot move."""
128
+ result = _profile_result(
129
+ consumed=-RELEASED_BY_OTHERS,
130
+ reserved_after=TORCH_RESERVED,
131
+ )
132
+
133
+ assert maybe_rocm_profiling_fallback(result) == TORCH_RESERVED
134
+
135
+
136
+@pytest.mark.parametrize("rocm", [True], indirect=True)
137
+def test_profiling_fallback_never_returns_a_negative_amount(rocm):
138
+ """A reservation that shrank across the run cannot become negative usage."""
139
+ result = _profile_result(
140
+ consumed=-RELEASED_BY_OTHERS,
141
+ reserved_before=TORCH_RESERVED,
142
+ reserved_after=0,
143
+ )
144
+
145
+ assert maybe_rocm_profiling_fallback(result) == 0
146
+
147
+
148
+@pytest.mark.parametrize("rocm", [False], indirect=True)
149
+def test_profiling_fallback_declines_off_rocm(rocm):
150
+ """Platforms that account frees eagerly keep reporting the error, so the
151
+ caller's assertion stays reachable there."""
152
+ result = _profile_result(consumed=-RELEASED_BY_OTHERS)
153
+
154
+ assert maybe_rocm_profiling_fallback(result) is None
vllm/v1/worker/gpu_worker.py
+54
-13
@@ -66,7 +66,12 @@ from vllm.tracing import instrument
66
from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
67
from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
68
from vllm.utils.mem_constants import GiB_bytes
69
-from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
69
+from vllm.utils.mem_utils import (
70
+ MemoryProfilingResult,
71
+ MemorySnapshot,
72
+ format_gib,
73
+ memory_profiling,
74
+)
75
from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime
76
from vllm.v1.attention.backends.utils import record_kv_cache_layout
77
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -102,6 +107,33 @@ def _num_workspace_lanes(vllm_config: VllmConfig, use_v2_model_runner: bool) ->
107
)
108
109
110
+def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int | None:
111
+ """Memory in bytes to size the KV cache from when profiling measured a
112
+ release on ROCm, or None to keep what profiling measured.
113
+
114
+ The measurement is the drop in free device memory across startup, so it comes
115
+ out negative when anything else on the device released memory meanwhile. Torch
116
+ reserved is process local, so it stands in as a lower bound instead of refusing
117
+ to start. Kept to ROCm, where the AMD CI groups hit this, following #45490.
118
+ """
119
+ if profile_result.total_consumed >= 0 or not current_platform.is_rocm():
120
+ return None
121
+
122
+ baseline = profile_result.before_create
123
+ after_profile = profile_result.after_profile
124
+ torch_reserved = max(after_profile.torch_memory - baseline.torch_memory, 0)
125
+ logger.warning(
126
+ "Free memory grew by %s GiB during profiling (initial %s GiB, current "
127
+ "%s GiB). Sizing the KV cache from the %s GiB torch reserved instead of "
128
+ "the profiling difference.",
129
+ format_gib(-profile_result.total_consumed),
130
+ format_gib(baseline.free_memory),
131
+ format_gib(after_profile.free_memory),
132
+ format_gib(torch_reserved),
133
+ )
134
+ return torch_reserved
135
+
136
+
137
if TYPE_CHECKING:
138
from vllm.device_allocator.sleep_mode_backend import SleepModeBackend
139
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
@@ -543,24 +575,33 @@ class Worker(WorkerBase):
575
else 0
576
)
577
578
+ init_free_memory = self.init_snapshot.free_memory
579
+ free_gpu_memory = profile_result.after_profile.free_memory
580
+ rocm_fallback = maybe_rocm_profiling_fallback(profile_result)
581
+ if rocm_fallback is None:
582
+ # NOTE(woosuk): Here we assume that the other processes using the same
583
+ # GPU did not change their memory usage during the profiling.
584
+ assert init_free_memory >= free_gpu_memory, (
585
+ "Error in memory profiling. "
586
+ f"Initial free memory {format_gib(init_free_memory)} GiB, "
587
+ f"current free memory {format_gib(free_gpu_memory)} GiB. "
588
+ "This happens when other processes sharing the same container "
589
+ "release GPU memory while vLLM is profiling during initialization. "
590
+ "To fix this, ensure consistent GPU memory allocation or "
591
+ "isolate vLLM in its own container."
592
+ )
593
+ else:
594
+ profile_result.total_consumed = rocm_fallback
595
+ profile_result.non_kv_cache_memory = (
596
+ profile_result.total_consumed + profile_result.transient_peak_headroom
597
+ )
598
+
599
self.total_consumed = profile_result.total_consumed
600
self.peak_activation_memory = (
601
profile_result.transient_peak_headroom + cudagraph_memory_estimate_applied
602
)
603
self.cudagraph_memory_estimate = cudagraph_memory_estimate
604
552
- free_gpu_memory = profile_result.after_profile.free_memory
553
- # NOTE(woosuk): Here we assume that the other processes using the same
554
- # GPU did not change their memory usage during the profiling.
555
- assert self.init_snapshot.free_memory >= free_gpu_memory, (
556
- "Error in memory profiling. "
557
- f"Initial free memory {format_gib(self.init_snapshot.free_memory)} GiB, "
558
- f"current free memory {format_gib(free_gpu_memory)} GiB. "
559
- "This happens when other processes sharing the same container "
560
- "release GPU memory while vLLM is profiling during initialization. "
561
- "To fix this, ensure consistent GPU memory allocation or "
562
- "isolate vLLM in its own container."
563
- )
605
self.available_kv_cache_memory_bytes = (
606
self.requested_memory
607
- profile_result.non_kv_cache_memory