Skip to content

vllm.device_allocator.alloc_conf

Helpers for reading and editing PyTorch's allocator configuration.

Torch's allocator config is parsed as a whole: parseArgs resets every option that is not explicitly present in the string it is handed. Writing a single field therefore silently drops the rest of the user's configuration, for the remainder of the process:

>>> # PYTORCH_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512
>>> torch._C._accelerator_setAllocatorSettings("expandable_segments:False")
>>> # max_split_size is now SIZE_MAX, not 512 MiB -- and writing
>>> # "expandable_segments:True" back does not bring it home either.

So anything that wants to toggle one option has to read the current config, flip just that field, and write the whole string back.

Functions:

alloc_conf_from_env(environ=None)

The allocator config string the process was started with, if any.

Source code in vllm/device_allocator/alloc_conf.py 
34
35
36
37
38
39
40
41
42
43
def alloc_conf_from_env(
    environ: "os._Environ[str] | dict[str, str] | None" = None,
) -> str:
    """The allocator config string the process was started with, if any."""
    env = os.environ if environ is None else environ
    for name in ALLOC_CONF_ENV_VARS:
        conf = env.get(name)
        if conf:
            return conf
    return ""

conf_flag_enabled(conf, key)

Parse <key>:<bool> out of an allocator config string.

Source code in vllm/device_allocator/alloc_conf.py 
70
71
72
73
74
75
76
def conf_flag_enabled(conf: str, key: str) -> bool:
    """Parse ``<key>:<bool>`` out of an allocator config string."""
    for field in conf.split(","):
        name, _, value = field.partition(":")
        if name.strip() == key:
            return value.strip().lower() in ("true", "1")
    return False

current_alloc_conf()

The live allocator config, falling back to the environment.

Source code in vllm/device_allocator/alloc_conf.py 
64
65
66
67
def current_alloc_conf() -> str:
    """The live allocator config, falling back to the environment."""
    live = live_alloc_conf()
    return live if live is not None else alloc_conf_from_env()

expandable_segments_enabled()

Live expandable_segments state, or None if it cannot be read.

None is deliberately distinct from False: the environment can never reflect a runtime write, so a caller verifying a write it just issued must not read "cannot tell" as "the write did not take".

Source code in vllm/device_allocator/alloc_conf.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def expandable_segments_enabled() -> bool | None:
    """Live ``expandable_segments`` state, or ``None`` if it cannot be read.

    ``None`` is deliberately distinct from ``False``: the environment can never
    reflect a runtime write, so a caller verifying a write it just issued must
    not read "cannot tell" as "the write did not take".
    """
    try:
        settings = torch.cuda.memory._snapshot().get("allocator_settings")
        if settings is not None and EXPANDABLE_SEGMENTS in settings:
            return bool(settings[EXPANDABLE_SEGMENTS])
    except Exception:  # noqa: BLE001
        pass
    return None

live_alloc_conf()

The allocator config string currently in effect, or None.

Read from the live allocator rather than the environment, so that a runtime write by another component is visible.

Source code in vllm/device_allocator/alloc_conf.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def live_alloc_conf() -> str | None:
    """The allocator config string currently in effect, or ``None``.

    Read from the live allocator rather than the environment, so that a
    *runtime* write by another component is visible.
    """
    try:
        settings = torch.cuda.memory._snapshot().get("allocator_settings")
        if settings is not None:
            for name in ALLOC_CONF_ENV_VARS:
                conf = settings.get(name)
                if conf is not None:
                    return conf
    except Exception:  # noqa: BLE001
        pass
    return None

set_alloc_conf(conf)

Write a complete allocator config string.

Source code in vllm/device_allocator/alloc_conf.py 
113
114
115
116
117
118
119
def set_alloc_conf(conf: str) -> None:
    """Write a complete allocator config string."""
    setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None)
    if setter is not None:
        setter(conf)
        return
    torch.cuda.memory._set_allocator_settings(conf)

with_conf_flag(conf, key, enabled)

Return conf with only key flipped, every other field verbatim.

Source code in vllm/device_allocator/alloc_conf.py 
79
80
81
82
83
84
85
86
87
88
def with_conf_flag(conf: str, key: str, enabled: bool) -> str:
    """Return ``conf`` with only ``key`` flipped, every other field verbatim."""
    want = f"{key}:{'True' if enabled else 'False'}"
    fields = [f for f in (f.strip() for f in conf.split(",")) if f]
    for i, field in enumerate(fields):
        if field.partition(":")[0].strip() == key:
            fields[i] = want
            return ",".join(fields)
    fields.append(want)
    return ",".join(fields)