[Core] Persist and reuse the memory-profiling result across boots (opt-in) (#47388)
Signed-off-by: Nils Matteson <nils@thaw.sh> Signed-off-by: Nils Matteson <nilsmatteson@icloud.com> Co-authored-by: Nils Matteson <nils@thaw.sh> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Nils Matteson committed
Jul 7, 2026 at 18:53 UTC
e97c3cb303f9ba54ed1c73b14f880760cfdfd344
4 files changed
+284
tests/v1/worker/test_gpu_worker.py
+72
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from types import SimpleNamespace
5
+from unittest.mock import patch
6
7
import pytest
8
@@ -13,7 +14,12 @@ from vllm.multimodal.video import (
14
PYNVVIDEOCODEC_VIDEO_BACKEND,
15
)
16
from vllm.utils.mem_constants import GiB_bytes
17
+from vllm.v1.worker import startup_plan
18
from vllm.v1.worker.gpu_worker import Worker
19
+from vllm.v1.worker.startup_plan import (
20
+ maybe_apply_startup_plan,
21
+ maybe_save_startup_plan,
22
+)
23
24
25
def _worker_with_mm_config(
@@ -114,3 +120,69 @@ def test_reserve_mm_ipc_gpu_memory_scales_pynvvideocodec_budget_by_api_servers(
120
assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == (
121
available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3)
122
)
123
+
124
+
125
+# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
126
+# saved by Worker.determine_available_memory / compile_or_warm_up_model.
127
+
128
+
129
+def _plan_worker(config_hash="abc123", free_memory=78 * GiB_bytes, kv_bytes=None):
130
+ """The minimal Worker surface the startup-plan entry points touch."""
131
+ return SimpleNamespace(
132
+ vllm_config=SimpleNamespace(compute_hash=lambda: config_hash),
133
+ rank=0,
134
+ parallel_config=SimpleNamespace(world_size=1),
135
+ init_snapshot=SimpleNamespace(free_memory=free_memory),
136
+ cache_config=SimpleNamespace(kv_cache_memory_bytes=kv_bytes),
137
+ )
138
+
139
+
140
+def _plan_platform(name="NVIDIA H100 PCIe"):
141
+ return SimpleNamespace(
142
+ get_device_name=lambda device_id=0: name,
143
+ get_device_total_memory=lambda device_id=0: 80 * GiB_bytes,
144
+ get_device_capability=lambda device_id=0: (9, 0),
145
+ )
146
+
147
+
148
+@pytest.fixture
149
+def plan_env(monkeypatch: pytest.MonkeyPatch, tmp_path):
150
+ """Enable the startup plan, isolated under a tmp cache root."""
151
+ monkeypatch.setenv("VLLM_ENABLE_STARTUP_PLAN", "1")
152
+ monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path))
153
+ with patch.object(startup_plan, "current_platform", _plan_platform()):
154
+ yield
155
+
156
+
157
+def test_startup_plan_fingerprint_sensitivity(plan_env):
158
+ """The fingerprint is the OOM-safety key: stable for identical inputs,
159
+ different for anything the profiled value depends on."""
160
+ fp = startup_plan.compute_plan_fingerprint
161
+ base = fp(_plan_worker().vllm_config, 0, 1)
162
+ assert base == fp(_plan_worker().vllm_config, 0, 1)
163
+ assert base != fp(_plan_worker("other").vllm_config, 0, 1)
164
+ assert base != fp(_plan_worker().vllm_config, 1, 2)
165
+ with patch.object(startup_plan, "current_platform", _plan_platform("NVIDIA A100")):
166
+ assert base != fp(_plan_worker().vllm_config, 0, 1)
167
+ with patch("vllm.__version__", "0.0.0+plan-test"):
168
+ assert base != fp(_plan_worker().vllm_config, 0, 1)
169
+
170
+
171
+def test_startup_plan_apply_gate(plan_env):
172
+ """Only a fingerprint-matching, memory-safe plan is ever applied."""
173
+ maybe_save_startup_plan(_plan_worker(), 50 * GiB_bytes)
174
+
175
+ applied = _plan_worker()
176
+ maybe_apply_startup_plan(applied)
177
+ assert applied.cache_config.kv_cache_memory_bytes == 50 * GiB_bytes
178
+
179
+ less_memory = _plan_worker(free_memory=60 * GiB_bytes)
180
+ other_config = _plan_worker(config_hash="zzz999")
181
+ for refused in (less_memory, other_config):
182
+ maybe_apply_startup_plan(refused)
183
+ assert refused.cache_config.kv_cache_memory_bytes is None
184
+
185
+ # An explicit --kv-cache-memory is never overridden.
186
+ explicit = _plan_worker(kv_bytes=7 * GiB_bytes)
187
+ maybe_apply_startup_plan(explicit)
188
+ assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes
vllm/envs.py
+13
@@ -234,6 +234,7 @@ if TYPE_CHECKING:
234
VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
235
VLLM_ALLREDUCE_USE_FLASHINFER: bool = False
236
VLLM_TUNED_CONFIG_FOLDER: str | None = None
237
+ VLLM_ENABLE_STARTUP_PLAN: bool = False
238
VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
239
VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False
240
VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False
@@ -1728,6 +1729,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
1729
# Each component first checks this folder, then the configs shipped with
1730
# vLLM (if any). If no JSON matches, it uses a hard-coded heuristic.
1731
"VLLM_TUNED_CONFIG_FOLDER": lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None),
1732
+ # Opt-in persistence of the startup plan. When enabled, each worker
1733
+ # saves the memory-profiling result (the suggested --kv-cache-memory value
1734
+ # and the free-memory baseline) under VLLM_CACHE_ROOT/startup_plan/,
1735
+ # keyed by a hardware+config fingerprint, and later boots auto-apply it
1736
+ # -- skipping memory profiling -- when the fingerprint matches and
1737
+ # current free memory >= the recorded baseline.
1738
+ # See vllm/v1/worker/startup_plan.py.
1739
+ "VLLM_ENABLE_STARTUP_PLAN": lambda: bool(
1740
+ int(os.getenv("VLLM_ENABLE_STARTUP_PLAN", "0"))
1741
+ ),
1742
# Valid values are container,code_interpreter,web_search_preview
1743
# ex VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS=container,code_interpreter
1744
# If the server_label of your mcp tool is not in this list it will
@@ -2062,6 +2073,8 @@ def compile_factors() -> dict[str, object]:
2073
"VLLM_DEBUG_DUMP_PATH",
2074
"VLLM_PORT",
2075
"VLLM_CACHE_ROOT",
2076
+ # Runtime memory-plan persistence; does not affect compiled graphs.
2077
+ "VLLM_ENABLE_STARTUP_PLAN",
2078
"LD_LIBRARY_PATH",
2079
"VLLM_SERVER_DEV_MODE",
2080
"VLLM_DP_MASTER_IP",
vllm/v1/worker/gpu_worker.py
+8
@@ -75,6 +75,10 @@ from vllm.v1.outputs import (
75
ModelRunnerOutput,
76
)
77
from vllm.v1.utils import compute_iteration_details, report_usage_stats
78
+from vllm.v1.worker.startup_plan import (
79
+ maybe_apply_startup_plan,
80
+ maybe_save_startup_plan,
81
+)
82
from vllm.v1.worker.utils import is_residual_scattered_for_sp
83
from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase
84
from vllm.v1.worker.workspace import init_workspace_manager
@@ -439,6 +443,8 @@ class Worker(WorkerBase):
443
You may limit the usage of GPU memory
444
by adjusting the `gpu_memory_utilization` parameter.
445
"""
446
+ maybe_apply_startup_plan(self)
447
+
448
if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
449
# still need a profile run which compiles the model for
450
# max_num_batched_tokens
@@ -834,6 +840,8 @@ class Worker(WorkerBase):
840
841
logger.info(msg)
842
843
+ maybe_save_startup_plan(self, kv_cache_memory_bytes_to_requested_limit)
844
+
845
if self.use_v2_model_runner:
846
# V2: Run full execute_model + sample_tokens to JIT compile triton kernels.
847
warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)
vllm/v1/worker/startup_plan.py
new
+191
@@ -0,0 +1,191 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Persist and reuse the memory-profiling result across engine boots.
4
+
5
+On startup, vLLM measures how much GPU memory the KV cache can use and
6
+computes the ``--kv-cache-memory`` value that reproduces that allocation.
7
+For a fixed (model, config, hardware, library) combination the result is
8
+deterministic, yet it is re-measured on every boot.
9
+
10
+When ``VLLM_ENABLE_STARTUP_PLAN=1``, each worker persists that value under
11
+``{VLLM_CACHE_ROOT}/startup_plan/`` (regenerable derived state, alongside
12
+the torch.compile cache), keyed by a fingerprint of everything the value
13
+depends on, and later boots apply it automatically -- skipping the
14
+memory-profiling measurement and the CUDA-graph memory estimation pass --
15
+if and only if the fingerprint matches and the device has at least as much
16
+free memory as when the plan was recorded. On any mismatch the worker
17
+falls back to full profiling, so a stale plan costs nothing and is never
18
+trusted.
19
+"""
20
+
21
+import hashlib
22
+import json
23
+import os
24
+from typing import TYPE_CHECKING
25
+
26
+import torch
27
+
28
+import vllm.envs as envs
29
+from vllm.config import VllmConfig
30
+from vllm.logger import init_logger
31
+from vllm.platforms import current_platform
32
+
33
+if TYPE_CHECKING:
34
+ from vllm.v1.worker.gpu_worker import Worker
35
+
36
+logger = init_logger(__name__)
37
+
38
+PLAN_SCHEMA_VERSION = 1
39
+
40
+
41
+def compute_plan_fingerprint(
42
+ vllm_config: VllmConfig, rank: int, world_size: int
43
+) -> str:
44
+ """Hash everything the profiled KV-cache memory value depends on.
45
+
46
+ ``VllmConfig.compute_hash()`` covers the vLLM version and the model,
47
+ cache, parallel, and compilation configs, but deliberately contains no
48
+ device identity (``DeviceConfig.compute_hash`` is empty), so device
49
+ name, total memory, compute capability, and the torch/CUDA build are
50
+ added here. The vLLM version is also pinned as an explicit factor so
51
+ version invalidation holds no matter how ``compute_hash`` evolves.
52
+ Rank is included because per-rank memory use differs under TP/PP.
53
+ Driver-only changes are not part of the key; the free-memory gate at
54
+ apply time bounds the residual risk.
55
+ """
56
+ # Imported here (as VllmConfig.compute_hash does) to avoid a cycle with
57
+ # the top-level vllm package.
58
+ from vllm import __version__ as vllm_version
59
+
60
+ capability = current_platform.get_device_capability()
61
+ factors = {
62
+ "schema": PLAN_SCHEMA_VERSION,
63
+ "vllm": vllm_version,
64
+ "vllm_config": vllm_config.compute_hash(),
65
+ "device_name": current_platform.get_device_name(),
66
+ "device_total_memory": current_platform.get_device_total_memory(),
67
+ "device_capability": str(capability) if capability else "",
68
+ "torch": torch.__version__,
69
+ "cuda": torch.version.cuda or "",
70
+ "rank": rank,
71
+ "world_size": world_size,
72
+ }
73
+ digest = hashlib.sha256(json.dumps(factors, sort_keys=True).encode()).hexdigest()
74
+ return digest[:16]
75
+
76
+
77
+def _plan_path(fingerprint: str) -> str:
78
+ """Plans are regenerable derived state, so they live under the standard
79
+ vLLM cache root (like the torch.compile cache) and relocate with
80
+ ``VLLM_CACHE_ROOT`` instead of needing a location knob of their own."""
81
+ # VLLM_CACHE_ROOT is already user-expanded by envs.py.
82
+ return os.path.join(
83
+ envs.VLLM_CACHE_ROOT, "startup_plan", f"startup_plan_{fingerprint}.json"
84
+ )
85
+
86
+
87
+def _load_plan(fingerprint: str) -> dict | None:
88
+ """Load a plan for this fingerprint; None if absent or unreadable."""
89
+ path = _plan_path(fingerprint)
90
+ try:
91
+ with open(path) as f:
92
+ plan = json.load(f)
93
+ except FileNotFoundError:
94
+ return None
95
+ except (OSError, json.JSONDecodeError) as e:
96
+ logger.warning("Ignoring unreadable startup plan %s: %s", path, e)
97
+ return None
98
+ if (
99
+ plan.get("schema") != PLAN_SCHEMA_VERSION
100
+ or plan.get("fingerprint") != fingerprint
101
+ ):
102
+ return None
103
+ return plan
104
+
105
+
106
+def _applicable_kv_cache_memory_bytes(
107
+ plan: dict, current_free_memory: int
108
+) -> int | None:
109
+ """The apply-time OOM-safety gate.
110
+
111
+ The recorded value is only valid if the device has at least as much
112
+ free memory now as when the plan was measured (co-tenants, leaked
113
+ allocations, or MIG changes all reduce it). Outside that envelope,
114
+ return None and let the caller re-profile.
115
+ """
116
+ kv_bytes = plan.get("kv_cache_memory_bytes")
117
+ baseline = plan.get("free_memory_baseline")
118
+ if not isinstance(kv_bytes, int) or not isinstance(baseline, int):
119
+ return None
120
+ if kv_bytes <= 0:
121
+ return None
122
+ if current_free_memory < baseline:
123
+ logger.info(
124
+ "Startup plan not applied: current free memory (%.2f GiB) is "
125
+ "below the recorded baseline (%.2f GiB); falling back to full "
126
+ "memory profiling.",
127
+ current_free_memory / (1 << 30),
128
+ baseline / (1 << 30),
129
+ )
130
+ return None
131
+ return kv_bytes
132
+
133
+
134
+def maybe_apply_startup_plan(worker: "Worker") -> None:
135
+ """If enabled and ``--kv-cache-memory`` was not set explicitly, apply a
136
+ persisted plan by setting ``worker.cache_config.kv_cache_memory_bytes``.
137
+ No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``."""
138
+ if (
139
+ not envs.VLLM_ENABLE_STARTUP_PLAN
140
+ or worker.cache_config.kv_cache_memory_bytes is not None
141
+ ):
142
+ return
143
+ fingerprint = compute_plan_fingerprint(
144
+ worker.vllm_config, worker.rank, worker.parallel_config.world_size
145
+ )
146
+ plan = _load_plan(fingerprint)
147
+ if plan is None:
148
+ return
149
+ current_free_memory = worker.init_snapshot.free_memory
150
+ kv_bytes = _applicable_kv_cache_memory_bytes(plan, current_free_memory)
151
+ if kv_bytes is None:
152
+ return
153
+ logger.info(
154
+ "Applying persisted startup plan (fingerprint %s): "
155
+ "kv_cache_memory_bytes=%d (%.2f GiB), recorded free-memory "
156
+ "baseline %.2f GiB, current %.2f GiB. Memory profiling will "
157
+ "be skipped.",
158
+ fingerprint,
159
+ kv_bytes,
160
+ kv_bytes / (1 << 30),
161
+ plan["free_memory_baseline"] / (1 << 30),
162
+ current_free_memory / (1 << 30),
163
+ )
164
+ worker.cache_config.kv_cache_memory_bytes = kv_bytes
165
+
166
+
167
+def maybe_save_startup_plan(worker: "Worker", kv_cache_memory_bytes: int) -> None:
168
+ """Atomically persist this boot's profiling result for future boots.
169
+ No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``; failures are logged,
170
+ never raised."""
171
+ if not envs.VLLM_ENABLE_STARTUP_PLAN:
172
+ return
173
+ fingerprint = compute_plan_fingerprint(
174
+ worker.vllm_config, worker.rank, worker.parallel_config.world_size
175
+ )
176
+ path = _plan_path(fingerprint)
177
+ try:
178
+ os.makedirs(os.path.dirname(path), exist_ok=True)
179
+ payload = {
180
+ "schema": PLAN_SCHEMA_VERSION,
181
+ "fingerprint": fingerprint,
182
+ "kv_cache_memory_bytes": int(kv_cache_memory_bytes),
183
+ "free_memory_baseline": int(worker.init_snapshot.free_memory),
184
+ }
185
+ tmp = f"{path}.tmp.{os.getpid()}"
186
+ with open(tmp, "w") as f:
187
+ json.dump(payload, f)
188
+ os.replace(tmp, path)
189
+ logger.info("Saved startup plan to %s", path)
190
+ except OSError as e:
191
+ logger.warning("Failed to save startup plan to %s: %s", path, e)