[Bugfix] Stop allocator fragmentation from shrinking the KV cache during memory profiling (#58430)
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>
Robert Shaw committed
Sep 25, 2026 at 12:16 UTC
6491f481a7c0fb3aa77bbc6649584b545c65c871
2 files changed
+134
-20
tests/v1/worker/test_gpu_worker.py
+100
@@ -8,6 +8,8 @@ from unittest.mock import patch
8
import pytest
9
import torch
10
11
+from tests.utils import create_new_process_for_each_test
12
+from vllm.platforms import current_platform
13
from vllm.utils.mem_constants import GiB_bytes
14
from vllm.v1.worker import gpu_worker, startup_plan
15
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):
54
torch.set_num_threads(original_threads)
55
56
57
+@pytest.mark.skipif(
58
+ not current_platform.is_cuda_alike()
59
+ or not torch.accelerator.is_available()
60
+ or torch.cuda.memory.get_allocator_backend() != "native",
61
+ reason="needs the native CUDA or ROCm allocator",
62
+)
63
+# A fresh allocator: blocks cached by earlier tests could serve the large buffer.
64
+@create_new_process_for_each_test("spawn")
65
+def test_scoped_max_split_keeps_freed_large_blocks_releasable():
66
+ """A small allocation made after a large buffer is freed must not pin the
67
+ buffer's segment: the profiling run (determine_available_memory) grows
68
+ workspaces this way, and a pinned segment survives empty_cache() and is
69
+ counted as consumed memory."""
70
+ large = 512 * 1024 * 1024
71
+ small = 2 * 1024 * 1024 # large pool, so it is served by splitting cached blocks
72
+
73
+ def reserved_while_small_is_live(scope) -> int:
74
+ torch.accelerator.empty_cache()
75
+ with scope:
76
+ buf = torch.empty(large, dtype=torch.uint8, device="cuda")
77
+ del buf
78
+ tensor = torch.empty(small, dtype=torch.uint8, device="cuda")
79
+ torch.accelerator.empty_cache()
80
+ reserved = torch.accelerator.memory_reserved()
81
+ del tensor
82
+ torch.accelerator.empty_cache()
83
+ return reserved
84
+
85
+ baseline = torch.accelerator.memory_reserved()
86
+ # Without the limit the small tensor is split off the freed block and pins it.
87
+ assert reserved_while_small_is_live(nullcontext()) - baseline >= large
88
+ scoped = gpu_worker.Worker._scoped_allocator_max_split(
89
+ SimpleNamespace(), max_split_size_mb=20
90
+ )
91
+ assert reserved_while_small_is_live(scoped) - baseline < large
92
+
93
+
94
+@pytest.mark.skipif(
95
+ not current_platform.is_cuda_alike()
96
+ or not torch.accelerator.is_available()
97
+ or torch.cuda.memory.get_allocator_backend() != "native",
98
+ reason="needs the native CUDA or ROCm allocator",
99
+)
100
+@pytest.mark.parametrize("fail", [False, True])
101
+@pytest.mark.parametrize("snapshot_fallback", [False, True])
102
+@pytest.mark.parametrize("suffix", ["", ", "])
103
+def test_scoped_max_split_preserves_allocator_settings(
104
+ monkeypatch, fail, snapshot_fallback, suffix
105
+):
106
+ """Preserve allocator settings across successful and failed profiling runs."""
107
+ if snapshot_fallback:
108
+ monkeypatch.delattr(
109
+ torch._C, "_accelerator_getAllocatorSettings", raising=False
110
+ )
111
+
112
+ def settings():
113
+ return torch.cuda.memory._snapshot()["allocator_settings"]
114
+
115
+ original = settings()["PYTORCH_CUDA_ALLOC_CONF"]
116
+ configured = (
117
+ "max_split_size_mb:128,garbage_collection_threshold:0.8,"
118
+ "roundup_power2_divisions:[256:1,512:2,>:4],max_non_split_rounding_mb:32"
119
+ )
120
+ try:
121
+ torch._C._accelerator_setAllocatorSettings(configured + suffix)
122
+ before = settings()
123
+ expected_error = pytest.raises(RuntimeError, match="profiling failed")
124
+ with (
125
+ expected_error if fail else nullcontext(),
126
+ gpu_worker.Worker._scoped_allocator_max_split(SimpleNamespace(), 20),
127
+ ):
128
+ scoped = settings()
129
+ assert scoped["max_split_size"] == 20 * 1024 * 1024
130
+ for key in ("garbage_collection_threshold", "roundup_power2_divisions"):
131
+ assert scoped[key] == before[key]
132
+ if fail:
133
+ raise RuntimeError("profiling failed")
134
+ assert settings() == before
135
+ finally:
136
+ torch._C._accelerator_setAllocatorSettings(original)
137
+
138
+
139
+def test_scoped_max_split_ignores_async_allocator(monkeypatch):
140
+ """Async allocators ignore max_split and may not support memory snapshots."""
141
+ monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
142
+ monkeypatch.setattr(
143
+ torch.cuda.memory, "get_allocator_backend", lambda: "cudaMallocAsync"
144
+ )
145
+ with (
146
+ patch.object(torch.cuda.memory, "_snapshot") as snapshot,
147
+ patch.object(torch._C, "_accelerator_setAllocatorSettings") as set_settings,
148
+ gpu_worker.Worker._scoped_allocator_max_split(SimpleNamespace(), 20),
149
+ ):
150
+ pass
151
+ snapshot.assert_not_called()
152
+ set_settings.assert_not_called()
153
+
154
+
155
# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
156
# saved by Worker.determine_available_memory / compile_or_warm_up_model.
157
vllm/v1/worker/gpu_worker.py
+34
-20
@@ -13,7 +13,6 @@ from types import NoneType
13
from typing import TYPE_CHECKING, Any, cast
14
15
import numpy as np
16
-import regex as re
16
import torch
17
import torch.nn as nn
18
@@ -367,27 +366,32 @@ class Worker(WorkerBase):
366
367
@contextmanager
368
def _scoped_allocator_max_split(self, max_split_size_mb: int):
370
- """Temporarily set max_split_size_mb to reduce allocator fragmentation at the
371
- cost of more cudaMalloc calls (negligible in practice). Restores the original
372
- value on exit."""
373
- if not current_platform.is_cuda():
369
+ """Limit CUDA/ROCm allocator splitting, restoring settings on exit."""
370
+ if (
371
+ not current_platform.is_cuda_alike()
372
+ or torch.cuda.memory.get_allocator_backend() != "native"
373
+ ):
374
yield
375
return
376
377
- conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
378
- match = re.search(r"max_split_size_mb:(\d+)", conf)
379
- original_value = match.group(1) if match else None
380
-
381
- torch._C._accelerator_setAllocatorSettings(
382
- f"max_split_size_mb:{max_split_size_mb}"
383
- )
377
+ get_settings = getattr(torch._C, "_accelerator_getAllocatorSettings", None)
378
+ if get_settings is not None:
379
+ original_settings = get_settings()
380
+ else:
381
+ original_settings = torch.cuda.memory._snapshot()["allocator_settings"][
382
+ "PYTORCH_CUDA_ALLOC_CONF"
383
+ ]
384
+
385
+ # The setter resets other allocator options unless they are included.
386
+ prefix = original_settings.strip().rstrip(",")
387
+ settings = f"max_split_size_mb:{max_split_size_mb}"
388
+ if prefix:
389
+ settings = f"{prefix},{settings}"
390
try:
391
+ torch._C._accelerator_setAllocatorSettings(settings)
392
yield
393
finally:
387
- # PyTorch defaults to SIZE_MAX (no limit).
388
- _SIZE_MAX_MB = (2**64 - 1) // (1024 * 1024)
389
- restore = original_value if original_value else str(_SIZE_MAX_MB)
390
- torch._C._accelerator_setAllocatorSettings(f"max_split_size_mb:{restore}")
394
+ torch._C._accelerator_setAllocatorSettings(original_settings)
395
396
@instrument(span_name="Init device")
397
def init_device(self):
@@ -605,10 +609,20 @@ class Worker(WorkerBase):
609
610
# Execute a forward pass with dummy inputs to profile the memory usage
611
# of the model.
608
- with memory_profiling(
609
- self.init_snapshot,
610
- weights_memory=int(self.model_runner.model_memory_usage),
611
- ) as profile_result:
612
+ with (
613
+ memory_profiling(
614
+ self.init_snapshot,
615
+ weights_memory=int(self.model_runner.model_memory_usage),
616
+ ) as profile_result,
617
+ # Workspaces (e.g. the MoE workspace) grow in steps during this pass,
618
+ # freeing each smaller buffer. Without a split limit, a later small
619
+ # allocation can be carved out of a freed multi-GiB block and pin the
620
+ # whole segment past the empty_cache() in memory_profiling, so it is
621
+ # counted as consumed and taken from the KV cache. Blocks above the
622
+ # limit are never split, so they stay releasable. Exits before
623
+ # memory_profiling measures, restoring the original limit.
624
+ self._scoped_allocator_max_split(max_split_size_mb=20),
625
+ ):
626
self.model_runner.profile_run()
627
628
# Profile CUDA graph memory if graphs will be captured.