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

[Bugfix] Stop allocator fragmentation from shrinking the KV cache during memory profiling (#58430)

Browse Source 
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
Signed-off-by: Robert Shaw <robertgshaw2-redhat@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
Signed-off-by: Robert Shaw <robshaw@redhat.com>
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: Robert Shaw <robertgshaw2-redhat@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
Co-authored-by: Robert Shaw <robshaw@redhat.com>
Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Robert Shaw
2026-09-25 12:16:27 -04:00
committed by GitHub
co-authored by Robert Shaw Robert Shaw Andreas Karatzas Andreas Karatzas
parent a44d7b5151
commit 6491f481a7
2 changed files with 133 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_gpu_worker.py
+100
View File
@@ -8,6 +8,8 @@ from unittest.mock import patch
import pytest
import torch
from tests.utils import create_new_process_for_each_test
from vllm.platforms import current_platform
from vllm.utils.mem_constants import GiB_bytes
from vllm.v1.worker import gpu_worker, startup_plan
from vllm.v1.worker.gpu_worker import maybe_rocm_profiling_fallback
@@ -52,6 +54,104 @@ def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch):
torch.set_num_threads(original_threads)
@pytest.mark.skipif(
not current_platform.is_cuda_alike()
or not torch.accelerator.is_available()
or torch.cuda.memory.get_allocator_backend() != "native",
reason="needs the native CUDA or ROCm allocator",
)
# A fresh allocator: blocks cached by earlier tests could serve the large buffer.
@create_new_process_for_each_test("spawn")
def test_scoped_max_split_keeps_freed_large_blocks_releasable():
"""A small allocation made after a large buffer is freed must not pin the
buffer's segment: the profiling run (determine_available_memory) grows
workspaces this way, and a pinned segment survives empty_cache() and is
counted as consumed memory."""
large = 512 * 1024 * 1024
small = 2 * 1024 * 1024 # large pool, so it is served by splitting cached blocks
def reserved_while_small_is_live(scope) -> int:
torch.accelerator.empty_cache()
with scope:
buf = torch.empty(large, dtype=torch.uint8, device="cuda")
del buf
tensor = torch.empty(small, dtype=torch.uint8, device="cuda")
torch.accelerator.empty_cache()
reserved = torch.accelerator.memory_reserved()
del tensor
torch.accelerator.empty_cache()
return reserved
baseline = torch.accelerator.memory_reserved()
# Without the limit the small tensor is split off the freed block and pins it.
assert reserved_while_small_is_live(nullcontext()) - baseline >= large
scoped = gpu_worker.Worker._scoped_allocator_max_split(
SimpleNamespace(), max_split_size_mb=20
)
assert reserved_while_small_is_live(scoped) - baseline < large
@pytest.mark.skipif(
not current_platform.is_cuda_alike()
or not torch.accelerator.is_available()
or torch.cuda.memory.get_allocator_backend() != "native",
reason="needs the native CUDA or ROCm allocator",
)
@pytest.mark.parametrize("fail", [False, True])
@pytest.mark.parametrize("snapshot_fallback", [False, True])
@pytest.mark.parametrize("suffix", ["", ", "])
def test_scoped_max_split_preserves_allocator_settings(
monkeypatch, fail, snapshot_fallback, suffix
):
"""Preserve allocator settings across successful and failed profiling runs."""
if snapshot_fallback:
monkeypatch.delattr(
torch._C, "_accelerator_getAllocatorSettings", raising=False
)
def settings():
return torch.cuda.memory._snapshot()["allocator_settings"]
original = settings()["PYTORCH_CUDA_ALLOC_CONF"]
configured = (
"max_split_size_mb:128,garbage_collection_threshold:0.8,"
"roundup_power2_divisions:[256:1,512:2,>:4],max_non_split_rounding_mb:32"
)
try:
torch._C._accelerator_setAllocatorSettings(configured + suffix)
before = settings()
expected_error = pytest.raises(RuntimeError, match="profiling failed")
with (
expected_error if fail else nullcontext(),
gpu_worker.Worker._scoped_allocator_max_split(SimpleNamespace(), 20),
):
scoped = settings()
assert scoped["max_split_size"] == 20 * 1024 * 1024
for key in ("garbage_collection_threshold", "roundup_power2_divisions"):
assert scoped[key] == before[key]
if fail:
raise RuntimeError("profiling failed")
assert settings() == before
finally:
torch._C._accelerator_setAllocatorSettings(original)
def test_scoped_max_split_ignores_async_allocator(monkeypatch):
"""Async allocators ignore max_split and may not support memory snapshots."""
monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
monkeypatch.setattr(
torch.cuda.memory, "get_allocator_backend", lambda: "cudaMallocAsync"
)
with (
patch.object(torch.cuda.memory, "_snapshot") as snapshot,
patch.object(torch._C, "_accelerator_setAllocatorSettings") as set_settings,
gpu_worker.Worker._scoped_allocator_max_split(SimpleNamespace(), 20),
):
pass
snapshot.assert_not_called()
set_settings.assert_not_called()
# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
# saved by Worker.determine_available_memory / compile_or_warm_up_model.
vllm/v1/worker/gpu_worker.py
+33 -19
View File
@@ -13,7 +13,6 @@ from types import NoneType
from typing import TYPE_CHECKING, Any, cast
import numpy as np
import regex as re
import torch
import torch.nn as nn
@@ -367,27 +366,32 @@ class Worker(WorkerBase):
@contextmanager
def _scoped_allocator_max_split(self, max_split_size_mb: int):
"""Temporarily set max_split_size_mb to reduce allocator fragmentation at the
cost of more cudaMalloc calls (negligible in practice). Restores the original
value on exit."""
if not current_platform.is_cuda():
"""Limit CUDA/ROCm allocator splitting, restoring settings on exit."""
if (
not current_platform.is_cuda_alike()
or torch.cuda.memory.get_allocator_backend() != "native"
):
yield
return
conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
match = re.search(r"max_split_size_mb:(\d+)", conf)
original_value = match.group(1) if match else None
get_settings = getattr(torch._C, "_accelerator_getAllocatorSettings", None)
if get_settings is not None:
original_settings = get_settings()
else:
original_settings = torch.cuda.memory._snapshot()["allocator_settings"][
"PYTORCH_CUDA_ALLOC_CONF"
]
torch._C._accelerator_setAllocatorSettings(
f"max_split_size_mb:{max_split_size_mb}"
)
# The setter resets other allocator options unless they are included.
prefix = original_settings.strip().rstrip(",")
settings = f"max_split_size_mb:{max_split_size_mb}"
if prefix:
settings = f"{prefix},{settings}"
try:
torch._C._accelerator_setAllocatorSettings(settings)
yield
finally:
# PyTorch defaults to SIZE_MAX (no limit).
_SIZE_MAX_MB = (2**64 - 1) // (1024 * 1024)
restore = original_value if original_value else str(_SIZE_MAX_MB)
torch._C._accelerator_setAllocatorSettings(f"max_split_size_mb:{restore}")
torch._C._accelerator_setAllocatorSettings(original_settings)
@instrument(span_name="Init device")
def init_device(self):
@@ -605,10 +609,20 @@ class Worker(WorkerBase):
# Execute a forward pass with dummy inputs to profile the memory usage
# of the model.
with memory_profiling(
self.init_snapshot,
weights_memory=int(self.model_runner.model_memory_usage),
) as profile_result:
with (
memory_profiling(
self.init_snapshot,
weights_memory=int(self.model_runner.model_memory_usage),
) as profile_result,
# Workspaces (e.g. the MoE workspace) grow in steps during this pass,
# freeing each smaller buffer. Without a split limit, a later small
# allocation can be carved out of a freed multi-GiB block and pin the
# whole segment past the empty_cache() in memory_profiling, so it is
# counted as consumed and taken from the KV cache. Blocks above the
# limit are never split, so they stay releasable. Exits before
# memory_profiling measures, restoring the original limit.
self._scoped_allocator_max_split(max_split_size_mb=20),
):
self.model_runner.profile_run()
# Profile CUDA graph memory if graphs will be captured.