[Bugfix] Don't drop the rest of the allocator config when toggling expandable segments (#57982)
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Oxana Korzh committed
Sep 25, 2026 at 16:34 UTC
e55d076f89fd01a0538a3e496d8ff20bf7980100
4 files changed
+270
-8
.buildkite/test_areas/misc.yaml
+3
@@ -495,6 +495,7 @@ steps:
495
source_file_dependencies: &cpu-async-engine-deps
496
- vllm/assets/
497
- vllm/config/
498
+ - vllm/device_allocator/
499
- vllm/engine/arg_utils.py
500
- vllm/entrypoints/chat_utils.py
501
- vllm/entrypoints/mcp/
@@ -536,9 +537,11 @@ steps:
537
- tests/parser
538
- tests/transformers_utils
539
- tests/config
540
+ - tests/device_allocator
541
device: cpu-small
542
commands:
543
- python3 standalone_tests/lazy_imports.py
544
+ - pytest -v -s device_allocator
545
- pytest -v -s test_envs.py
546
- pytest -v -s test_outputs.py
547
- pytest -v -s test_pcp_dp.py
tests/device_allocator/test_alloc_conf.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.device_allocator import alloc_conf
7
+
8
+pytestmark = pytest.mark.cpu_test
9
+
10
+ES = alloc_conf.EXPANDABLE_SEGMENTS
11
+
12
+
13
+@pytest.mark.parametrize(
14
+ "environ,expected",
15
+ [
16
+ ({}, False),
17
+ ({"PYTORCH_ALLOC_CONF": "expandable_segments:True"}, True),
18
+ ({"PYTORCH_ALLOC_CONF": "expandable_segments:true"}, True),
19
+ ({"PYTORCH_ALLOC_CONF": "expandable_segments:1"}, True),
20
+ ({"PYTORCH_ALLOC_CONF": "expandable_segments:False"}, False),
21
+ ({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}, True),
22
+ # ROCm users set this one; it was previously not looked at.
23
+ ({"PYTORCH_HIP_ALLOC_CONF": "expandable_segments:True"}, True),
24
+ (
25
+ {"PYTORCH_ALLOC_CONF": "max_split_size_mb:512,expandable_segments:True"},
26
+ True,
27
+ ),
28
+ ({"PYTORCH_ALLOC_CONF": "max_split_size_mb:512"}, False),
29
+ ({"PYTORCH_ALLOC_CONF": "not_expandable_segments_really:True"}, False),
30
+ ({"PYTORCH_ALLOC_CONF": ""}, False),
31
+ # torch's precedence: the accelerator-agnostic variable wins.
32
+ (
33
+ {
34
+ "PYTORCH_ALLOC_CONF": "expandable_segments:False",
35
+ "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
36
+ },
37
+ False,
38
+ ),
39
+ ],
40
+)
41
+def test_expandable_segments_from_env(environ, expected):
42
+ assert alloc_conf.expandable_segments_enabled_from_env(environ) is expected
43
+
44
+
45
+@pytest.mark.parametrize(
46
+ "conf,key,expected",
47
+ [
48
+ ("", ES, False),
49
+ ("expandable_segments:True", ES, True),
50
+ ("expandable_segments:TRUE", ES, True),
51
+ ("expandable_segments:1", ES, True),
52
+ ("expandable_segments:0", ES, False),
53
+ (" expandable_segments : True ", ES, True),
54
+ ("max_split_size_mb:512", ES, False),
55
+ ("max_split_size_mb:512", "max_split_size_mb", False),
56
+ ("garbage_collection_threshold:0.9", ES, False),
57
+ ],
58
+)
59
+def test_conf_flag_enabled(conf, key, expected):
60
+ assert alloc_conf.conf_flag_enabled(conf, key) is expected
61
+
62
+
63
+@pytest.mark.parametrize(
64
+ "conf,enabled,expected",
65
+ [
66
+ ("", False, "expandable_segments:False"),
67
+ ("expandable_segments:True", False, "expandable_segments:False"),
68
+ ("expandable_segments:False", True, "expandable_segments:True"),
69
+ (
70
+ "max_split_size_mb:512,expandable_segments:True",
71
+ False,
72
+ "max_split_size_mb:512,expandable_segments:False",
73
+ ),
74
+ (
75
+ "garbage_collection_threshold:0.9,max_split_size_mb:512",
76
+ False,
77
+ "garbage_collection_threshold:0.9,max_split_size_mb:512,"
78
+ "expandable_segments:False",
79
+ ),
80
+ ],
81
+)
82
+def test_with_conf_flag_preserves_other_fields(conf, enabled, expected):
83
+ assert alloc_conf.with_conf_flag(conf, ES, enabled) == expected
84
+
85
+
86
+def test_with_conf_flag_round_trips():
87
+ """The regression this guards: torch resets every option absent from the
88
+ string it is handed, so a one-field write drops the rest permanently."""
89
+ conf = (
90
+ "max_split_size_mb:512,expandable_segments:True,"
91
+ "garbage_collection_threshold:0.9"
92
+ )
93
+ off = alloc_conf.with_conf_flag(conf, ES, False)
94
+ assert "max_split_size_mb:512" in off
95
+ assert "garbage_collection_threshold:0.9" in off
96
+ assert alloc_conf.with_conf_flag(off, ES, True) == conf
97
+
98
+
99
+def test_alloc_conf_from_env_precedence():
100
+ assert (
101
+ alloc_conf.alloc_conf_from_env(
102
+ {
103
+ "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:256",
104
+ "PYTORCH_HIP_ALLOC_CONF": "max_split_size_mb:128",
105
+ }
106
+ )
107
+ == "max_split_size_mb:256"
108
+ )
109
+ assert alloc_conf.alloc_conf_from_env({}) == ""
110
+
111
+
112
+def test_current_alloc_conf_prefers_live_state(monkeypatch):
113
+ monkeypatch.setattr(alloc_conf, "live_alloc_conf", lambda: "max_split_size_mb:64")
114
+ monkeypatch.setattr(alloc_conf, "alloc_conf_from_env", lambda: "stale:True")
115
+ assert alloc_conf.current_alloc_conf() == "max_split_size_mb:64"
116
+
117
+ monkeypatch.setattr(alloc_conf, "live_alloc_conf", lambda: None)
118
+ assert alloc_conf.current_alloc_conf() == "stale:True"
119
+
120
+
121
+def test_expandable_segments_enabled_returns_none_when_unreadable(monkeypatch):
122
+ """'cannot read' must stay distinct from 'disabled'."""
123
+
124
+ class _NoSnapshot:
125
+ @staticmethod
126
+ def _snapshot():
127
+ raise RuntimeError("no allocator here")
128
+
129
+ monkeypatch.setattr(alloc_conf.torch.cuda, "memory", _NoSnapshot)
130
+ assert alloc_conf.expandable_segments_enabled() is None
vllm/device_allocator/alloc_conf.py
new
+119
@@ -0,0 +1,119 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Helpers for reading and editing PyTorch's allocator configuration.
4
+
5
+Torch's allocator config is parsed as a whole: ``parseArgs`` resets every
6
+option that is not explicitly present in the string it is handed. Writing a
7
+single field therefore silently drops the rest of the user's configuration,
8
+for the remainder of the process:
9
+
10
+ >>> # PYTORCH_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512
11
+ >>> torch._C._accelerator_setAllocatorSettings("expandable_segments:False")
12
+ >>> # max_split_size is now SIZE_MAX, not 512 MiB -- and writing
13
+ >>> # "expandable_segments:True" back does not bring it home either.
14
+
15
+So anything that wants to toggle one option has to read the current config,
16
+flip just that field, and write the whole string back.
17
+"""
18
+
19
+import os
20
+
21
+import torch
22
+
23
+# Checked in torch's own precedence order: the accelerator-agnostic variable
24
+# wins, then the CUDA one, then the HIP one (which is what ROCm users set).
25
+ALLOC_CONF_ENV_VARS = (
26
+ "PYTORCH_ALLOC_CONF",
27
+ "PYTORCH_CUDA_ALLOC_CONF",
28
+ "PYTORCH_HIP_ALLOC_CONF",
29
+)
30
+
31
+EXPANDABLE_SEGMENTS = "expandable_segments"
32
+
33
+
34
+def alloc_conf_from_env(
35
+ environ: "os._Environ[str] | dict[str, str] | None" = None,
36
+) -> str:
37
+ """The allocator config string the process was started with, if any."""
38
+ env = os.environ if environ is None else environ
39
+ for name in ALLOC_CONF_ENV_VARS:
40
+ conf = env.get(name)
41
+ if conf:
42
+ return conf
43
+ return ""
44
+
45
+
46
+def live_alloc_conf() -> str | None:
47
+ """The allocator config string currently in effect, or ``None``.
48
+
49
+ Read from the live allocator rather than the environment, so that a
50
+ *runtime* write by another component is visible.
51
+ """
52
+ try:
53
+ settings = torch.cuda.memory._snapshot().get("allocator_settings")
54
+ if settings is not None:
55
+ for name in ALLOC_CONF_ENV_VARS:
56
+ conf = settings.get(name)
57
+ if conf is not None:
58
+ return conf
59
+ except Exception: # noqa: BLE001
60
+ pass
61
+ return None
62
+
63
+
64
+def current_alloc_conf() -> str:
65
+ """The live allocator config, falling back to the environment."""
66
+ live = live_alloc_conf()
67
+ return live if live is not None else alloc_conf_from_env()
68
+
69
+
70
+def conf_flag_enabled(conf: str, key: str) -> bool:
71
+ """Parse ``<key>:<bool>`` out of an allocator config string."""
72
+ for field in conf.split(","):
73
+ name, _, value = field.partition(":")
74
+ if name.strip() == key:
75
+ return value.strip().lower() in ("true", "1")
76
+ return False
77
+
78
+
79
+def with_conf_flag(conf: str, key: str, enabled: bool) -> str:
80
+ """Return ``conf`` with only ``key`` flipped, every other field verbatim."""
81
+ want = f"{key}:{'True' if enabled else 'False'}"
82
+ fields = [f for f in (f.strip() for f in conf.split(",")) if f]
83
+ for i, field in enumerate(fields):
84
+ if field.partition(":")[0].strip() == key:
85
+ fields[i] = want
86
+ return ",".join(fields)
87
+ fields.append(want)
88
+ return ",".join(fields)
89
+
90
+
91
+def expandable_segments_enabled_from_env(
92
+ environ: "os._Environ[str] | dict[str, str] | None" = None,
93
+) -> bool:
94
+ return conf_flag_enabled(alloc_conf_from_env(environ), EXPANDABLE_SEGMENTS)
95
+
96
+
97
+def expandable_segments_enabled() -> bool | None:
98
+ """Live ``expandable_segments`` state, or ``None`` if it cannot be read.
99
+
100
+ ``None`` is deliberately distinct from ``False``: the environment can never
101
+ reflect a runtime write, so a caller verifying a write it just issued must
102
+ not read "cannot tell" as "the write did not take".
103
+ """
104
+ try:
105
+ settings = torch.cuda.memory._snapshot().get("allocator_settings")
106
+ if settings is not None and EXPANDABLE_SEGMENTS in settings:
107
+ return bool(settings[EXPANDABLE_SEGMENTS])
108
+ except Exception: # noqa: BLE001
109
+ pass
110
+ return None
111
+
112
+
113
+def set_alloc_conf(conf: str) -> None:
114
+ """Write a complete allocator config string."""
115
+ setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None)
116
+ if setter is not None:
117
+ setter(conf)
118
+ return
119
+ torch.cuda.memory._set_allocator_settings(conf)
vllm/device_allocator/cumem.py
+18
-8
@@ -10,7 +10,6 @@
10
# the only successful approach is to call cuda driver API in C.
11
import atexit
12
import gc
13
-import os
13
from collections.abc import Callable, Iterator
14
from contextlib import contextmanager
15
from typing import Any
@@ -18,6 +17,13 @@ from typing import Any
17
import torch
18
19
from vllm.device_allocator import AllocationData, HandleType
20
+from vllm.device_allocator.alloc_conf import (
21
+ EXPANDABLE_SEGMENTS,
22
+ conf_flag_enabled,
23
+ current_alloc_conf,
24
+ set_alloc_conf,
25
+ with_conf_flag,
26
+)
27
from vllm.logger import init_logger
28
from vllm.platforms import current_platform
29
from vllm.utils.system_utils import find_loaded_library
@@ -371,13 +377,17 @@ class CuMemAllocator:
377
378
# Expandable segments are incompatible with the memory pool used for
379
# sleep mode (see https://github.com/pytorch/pytorch/issues/147851).
374
- # If the user has enabled expandable segments via
375
- # PYTORCH_CUDA_ALLOC_CONF, temporarily disable them for the duration
376
- # of the memory pool context and restore on exit.
377
- conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
378
- expandable_was_enabled = "expandable_segments:True" in conf
380
+ # If the user has enabled expandable segments, temporarily disable
381
+ # them for the duration of the memory pool context and restore on
382
+ # exit. The whole config string is rewritten, not just this one
383
+ # field: torch resets every option that is absent from the string it
384
+ # is handed, so a bare "expandable_segments:False" would also drop
385
+ # max_split_size_mb, garbage_collection_threshold and
386
+ # roundup_power2_divisions for the rest of the process.
387
+ prev_conf = current_alloc_conf()
388
+ expandable_was_enabled = conf_flag_enabled(prev_conf, EXPANDABLE_SEGMENTS)
389
if expandable_was_enabled:
380
- torch.cuda.memory._set_allocator_settings("expandable_segments:False")
390
+ set_alloc_conf(with_conf_flag(prev_conf, EXPANDABLE_SEGMENTS, False))
391
392
old_tag = self.current_tag
393
self.current_tag = tag
@@ -411,7 +421,7 @@ class CuMemAllocator:
421
finally:
422
self.current_tag = old_tag
423
if expandable_was_enabled:
414
- torch.cuda.memory._set_allocator_settings("expandable_segments:True")
424
+ set_alloc_conf(prev_conf)
425
426
def get_current_usage(self) -> int:
427
"""Get the total number of bytes allocated in the memory pool."""