iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][CI] Keep startup profiling from aborting when free memory grows (#53591)

Browse Source 
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
This commit is contained in:
stefankoncarevic
2026-08-28 00:02:40 -07:00
committed by GitHub
parent c01b50e390
commit 9236159bb6
2 changed files with 130 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_gpu_worker.py
+76 -1
View File
@@ -7,7 +7,8 @@ from unittest.mock import patch
import pytest
from vllm.utils.mem_constants import GiB_bytes
from vllm.v1.worker import startup_plan
from vllm.v1.worker import gpu_worker, startup_plan
from vllm.v1.worker.gpu_worker import maybe_rocm_profiling_fallback
from vllm.v1.worker.startup_plan import (
maybe_apply_startup_plan,
maybe_save_startup_plan,
@@ -77,3 +78,77 @@ def test_startup_plan_apply_gate(plan_env):
explicit = _plan_worker(kv_bytes=7 * GiB_bytes)
maybe_apply_startup_plan(explicit)
assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes
# Memory accounting of the profiling run (Worker.determine_available_memory).
# The fallback reads only the sign of the measured drop and this process's torch
# reservation; free memory is only logged, so no amount here is a device size.
ANY_FREE_MEMORY = 8 * GiB_bytes
MEASURED_DROP = 4 * GiB_bytes
TORCH_RESERVED = 3 * GiB_bytes
RELEASED_BY_OTHERS = 2 * GiB_bytes
def _snapshot(free_memory, torch_memory=0):
return SimpleNamespace(free_memory=free_memory, torch_memory=torch_memory)
def _profile_result(consumed, reserved_before=0, reserved_after=0):
"""A result whose free-memory readings agree with `consumed`, which
`memory_profiling` derives as the drop in free memory, negative when it grew."""
return SimpleNamespace(
total_consumed=consumed,
transient_peak_headroom=0,
before_create=_snapshot(ANY_FREE_MEMORY, reserved_before),
after_profile=_snapshot(ANY_FREE_MEMORY - consumed, reserved_after),
)
@pytest.fixture
def rocm(request):
with patch.object(
gpu_worker, "current_platform", SimpleNamespace(is_rocm=lambda: request.param)
):
yield request.param
@pytest.mark.parametrize("rocm", [True, False], indirect=True)
def test_profiling_fallback_declines_when_free_memory_dropped(rocm):
"""The profiling measurement is kept as-is whenever free memory dropped."""
result = _profile_result(consumed=MEASURED_DROP)
assert maybe_rocm_profiling_fallback(result) is None
@pytest.mark.parametrize("rocm", [True], indirect=True)
def test_profiling_fallback_replaces_a_released_measurement(rocm):
"""A negative measurement describes the rest of the device, so it is replaced
by this process's reservation, which the rest of the device cannot move."""
result = _profile_result(
consumed=-RELEASED_BY_OTHERS,
reserved_after=TORCH_RESERVED,
)
assert maybe_rocm_profiling_fallback(result) == TORCH_RESERVED
@pytest.mark.parametrize("rocm", [True], indirect=True)
def test_profiling_fallback_never_returns_a_negative_amount(rocm):
"""A reservation that shrank across the run cannot become negative usage."""
result = _profile_result(
consumed=-RELEASED_BY_OTHERS,
reserved_before=TORCH_RESERVED,
reserved_after=0,
)
assert maybe_rocm_profiling_fallback(result) == 0
@pytest.mark.parametrize("rocm", [False], indirect=True)
def test_profiling_fallback_declines_off_rocm(rocm):
"""Platforms that account frees eagerly keep reporting the error, so the
caller's assertion stays reachable there."""
result = _profile_result(consumed=-RELEASED_BY_OTHERS)
assert maybe_rocm_profiling_fallback(result) is None
vllm/v1/worker/gpu_worker.py
+54 -13
View File
@@ -66,7 +66,12 @@ from vllm.tracing import instrument
from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
from vllm.utils.mem_utils import (
MemoryProfilingResult,
MemorySnapshot,
format_gib,
memory_profiling,
)
from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime
from vllm.v1.attention.backends.utils import record_kv_cache_layout
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) ->
)
def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int | None:
"""Memory in bytes to size the KV cache from when profiling measured a
release on ROCm, or None to keep what profiling measured.
The measurement is the drop in free device memory across startup, so it comes
out negative when anything else on the device released memory meanwhile. Torch
reserved is process local, so it stands in as a lower bound instead of refusing
to start. Kept to ROCm, where the AMD CI groups hit this, following #45490.
"""
if profile_result.total_consumed >= 0 or not current_platform.is_rocm():
return None
baseline = profile_result.before_create
after_profile = profile_result.after_profile
torch_reserved = max(after_profile.torch_memory - baseline.torch_memory, 0)
logger.warning(
"Free memory grew by %s GiB during profiling (initial %s GiB, current "
"%s GiB). Sizing the KV cache from the %s GiB torch reserved instead of "
"the profiling difference.",
format_gib(-profile_result.total_consumed),
format_gib(baseline.free_memory),
format_gib(after_profile.free_memory),
format_gib(torch_reserved),
)
return torch_reserved
if TYPE_CHECKING:
from vllm.device_allocator.sleep_mode_backend import SleepModeBackend
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
@@ -543,24 +575,33 @@ class Worker(WorkerBase):
else 0
)
init_free_memory = self.init_snapshot.free_memory
free_gpu_memory = profile_result.after_profile.free_memory
rocm_fallback = maybe_rocm_profiling_fallback(profile_result)
if rocm_fallback is None:
# NOTE(woosuk): Here we assume that the other processes using the same
# GPU did not change their memory usage during the profiling.
assert init_free_memory >= free_gpu_memory, (
"Error in memory profiling. "
f"Initial free memory {format_gib(init_free_memory)} GiB, "
f"current free memory {format_gib(free_gpu_memory)} GiB. "
"This happens when other processes sharing the same container "
"release GPU memory while vLLM is profiling during initialization. "
"To fix this, ensure consistent GPU memory allocation or "
"isolate vLLM in its own container."
)
else:
profile_result.total_consumed = rocm_fallback
profile_result.non_kv_cache_memory = (
profile_result.total_consumed + profile_result.transient_peak_headroom
)
self.total_consumed = profile_result.total_consumed
self.peak_activation_memory = (
profile_result.transient_peak_headroom + cudagraph_memory_estimate_applied
)
self.cudagraph_memory_estimate = cudagraph_memory_estimate
free_gpu_memory = profile_result.after_profile.free_memory
# NOTE(woosuk): Here we assume that the other processes using the same
# GPU did not change their memory usage during the profiling.
assert self.init_snapshot.free_memory >= free_gpu_memory, (
"Error in memory profiling. "
f"Initial free memory {format_gib(self.init_snapshot.free_memory)} GiB, "
f"current free memory {format_gib(free_gpu_memory)} GiB. "
"This happens when other processes sharing the same container "
"release GPU memory while vLLM is profiling during initialization. "
"To fix this, ensure consistent GPU memory allocation or "
"isolate vLLM in its own container."
)
self.available_kv_cache_memory_bytes = (
self.requested_memory
- profile_result.non_kv_cache_memory