From 6491f481a7c0fb3aa77bbc6649584b545c65c871 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Fri, 25 Sep 2026 12:16:27 -0400 Subject: [PATCH] [Bugfix] Stop allocator fragmentation from shrinking the KV cache during memory profiling (#58430) Signed-off-by: Robert Shaw Signed-off-by: Robert Shaw Signed-off-by: Robert Shaw Signed-off-by: Andreas Karatzas Co-authored-by: Robert Shaw Co-authored-by: Robert Shaw Co-authored-by: Andreas Karatzas Co-authored-by: Andreas Karatzas --- tests/v1/worker/test_gpu_worker.py | 100 +++++++++++++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 52 +++++++++------ 2 files changed, 133 insertions(+), 19 deletions(-) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 722932636b..b63681deca 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -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. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 35c986b131..def92f601d 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -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.