From e97c3cb303f9ba54ed1c73b14f880760cfdfd344 Mon Sep 17 00:00:00 2001 From: Nils Matteson Date: Tue, 7 Jul 2026 18:53:02 -0600 Subject: [PATCH] [Core] Persist and reuse the memory-profiling result across boots (opt-in) (#47388) Signed-off-by: Nils Matteson Signed-off-by: Nils Matteson Co-authored-by: Nils Matteson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_gpu_worker.py | 72 +++++++++++ vllm/envs.py | 13 ++ vllm/v1/worker/gpu_worker.py | 8 ++ vllm/v1/worker/startup_plan.py | 191 +++++++++++++++++++++++++++++ 4 files changed, 284 insertions(+) create mode 100644 vllm/v1/worker/startup_plan.py diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 31be4a8402f..cdaa644b62e 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest.mock import patch import pytest @@ -13,7 +14,12 @@ from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, ) from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.worker import startup_plan from vllm.v1.worker.gpu_worker import Worker +from vllm.v1.worker.startup_plan import ( + maybe_apply_startup_plan, + maybe_save_startup_plan, +) def _worker_with_mm_config( @@ -114,3 +120,69 @@ def test_reserve_mm_ipc_gpu_memory_scales_pynvvideocodec_budget_by_api_servers( assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3) ) + + +# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and +# saved by Worker.determine_available_memory / compile_or_warm_up_model. + + +def _plan_worker(config_hash="abc123", free_memory=78 * GiB_bytes, kv_bytes=None): + """The minimal Worker surface the startup-plan entry points touch.""" + return SimpleNamespace( + vllm_config=SimpleNamespace(compute_hash=lambda: config_hash), + rank=0, + parallel_config=SimpleNamespace(world_size=1), + init_snapshot=SimpleNamespace(free_memory=free_memory), + cache_config=SimpleNamespace(kv_cache_memory_bytes=kv_bytes), + ) + + +def _plan_platform(name="NVIDIA H100 PCIe"): + return SimpleNamespace( + get_device_name=lambda device_id=0: name, + get_device_total_memory=lambda device_id=0: 80 * GiB_bytes, + get_device_capability=lambda device_id=0: (9, 0), + ) + + +@pytest.fixture +def plan_env(monkeypatch: pytest.MonkeyPatch, tmp_path): + """Enable the startup plan, isolated under a tmp cache root.""" + monkeypatch.setenv("VLLM_ENABLE_STARTUP_PLAN", "1") + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + with patch.object(startup_plan, "current_platform", _plan_platform()): + yield + + +def test_startup_plan_fingerprint_sensitivity(plan_env): + """The fingerprint is the OOM-safety key: stable for identical inputs, + different for anything the profiled value depends on.""" + fp = startup_plan.compute_plan_fingerprint + base = fp(_plan_worker().vllm_config, 0, 1) + assert base == fp(_plan_worker().vllm_config, 0, 1) + assert base != fp(_plan_worker("other").vllm_config, 0, 1) + assert base != fp(_plan_worker().vllm_config, 1, 2) + with patch.object(startup_plan, "current_platform", _plan_platform("NVIDIA A100")): + assert base != fp(_plan_worker().vllm_config, 0, 1) + with patch("vllm.__version__", "0.0.0+plan-test"): + assert base != fp(_plan_worker().vllm_config, 0, 1) + + +def test_startup_plan_apply_gate(plan_env): + """Only a fingerprint-matching, memory-safe plan is ever applied.""" + maybe_save_startup_plan(_plan_worker(), 50 * GiB_bytes) + + applied = _plan_worker() + maybe_apply_startup_plan(applied) + assert applied.cache_config.kv_cache_memory_bytes == 50 * GiB_bytes + + less_memory = _plan_worker(free_memory=60 * GiB_bytes) + other_config = _plan_worker(config_hash="zzz999") + for refused in (less_memory, other_config): + maybe_apply_startup_plan(refused) + assert refused.cache_config.kv_cache_memory_bytes is None + + # An explicit --kv-cache-memory is never overridden. + explicit = _plan_worker(kv_bytes=7 * GiB_bytes) + maybe_apply_startup_plan(explicit) + assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes diff --git a/vllm/envs.py b/vllm/envs.py index 8f40f8cfa9a..03082a1f902 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -234,6 +234,7 @@ if TYPE_CHECKING: VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True VLLM_ALLREDUCE_USE_FLASHINFER: bool = False VLLM_TUNED_CONFIG_FOLDER: str | None = None + VLLM_ENABLE_STARTUP_PLAN: bool = False VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set() VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False @@ -1728,6 +1729,16 @@ environment_variables: dict[str, Callable[[], Any]] = { # Each component first checks this folder, then the configs shipped with # vLLM (if any). If no JSON matches, it uses a hard-coded heuristic. "VLLM_TUNED_CONFIG_FOLDER": lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None), + # Opt-in persistence of the startup plan. When enabled, each worker + # saves the memory-profiling result (the suggested --kv-cache-memory value + # and the free-memory baseline) under VLLM_CACHE_ROOT/startup_plan/, + # keyed by a hardware+config fingerprint, and later boots auto-apply it + # -- skipping memory profiling -- when the fingerprint matches and + # current free memory >= the recorded baseline. + # See vllm/v1/worker/startup_plan.py. + "VLLM_ENABLE_STARTUP_PLAN": lambda: bool( + int(os.getenv("VLLM_ENABLE_STARTUP_PLAN", "0")) + ), # Valid values are container,code_interpreter,web_search_preview # ex VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS=container,code_interpreter # 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]: "VLLM_DEBUG_DUMP_PATH", "VLLM_PORT", "VLLM_CACHE_ROOT", + # Runtime memory-plan persistence; does not affect compiled graphs. + "VLLM_ENABLE_STARTUP_PLAN", "LD_LIBRARY_PATH", "VLLM_SERVER_DEV_MODE", "VLLM_DP_MASTER_IP", diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 341607b1b3a..182476e2533 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -75,6 +75,10 @@ from vllm.v1.outputs import ( ModelRunnerOutput, ) from vllm.v1.utils import compute_iteration_details, report_usage_stats +from vllm.v1.worker.startup_plan import ( + maybe_apply_startup_plan, + maybe_save_startup_plan, +) from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager @@ -439,6 +443,8 @@ class Worker(WorkerBase): You may limit the usage of GPU memory by adjusting the `gpu_memory_utilization` parameter. """ + maybe_apply_startup_plan(self) + if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: # still need a profile run which compiles the model for # max_num_batched_tokens @@ -834,6 +840,8 @@ class Worker(WorkerBase): logger.info(msg) + maybe_save_startup_plan(self, kv_cache_memory_bytes_to_requested_limit) + if self.use_v2_model_runner: # V2: Run full execute_model + sample_tokens to JIT compile triton kernels. warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens) diff --git a/vllm/v1/worker/startup_plan.py b/vllm/v1/worker/startup_plan.py new file mode 100644 index 00000000000..2c494207070 --- /dev/null +++ b/vllm/v1/worker/startup_plan.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Persist and reuse the memory-profiling result across engine boots. + +On startup, vLLM measures how much GPU memory the KV cache can use and +computes the ``--kv-cache-memory`` value that reproduces that allocation. +For a fixed (model, config, hardware, library) combination the result is +deterministic, yet it is re-measured on every boot. + +When ``VLLM_ENABLE_STARTUP_PLAN=1``, each worker persists that value under +``{VLLM_CACHE_ROOT}/startup_plan/`` (regenerable derived state, alongside +the torch.compile cache), keyed by a fingerprint of everything the value +depends on, and later boots apply it automatically -- skipping the +memory-profiling measurement and the CUDA-graph memory estimation pass -- +if and only if the fingerprint matches and the device has at least as much +free memory as when the plan was recorded. On any mismatch the worker +falls back to full profiling, so a stale plan costs nothing and is never +trusted. +""" + +import hashlib +import json +import os +from typing import TYPE_CHECKING + +import torch + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +PLAN_SCHEMA_VERSION = 1 + + +def compute_plan_fingerprint( + vllm_config: VllmConfig, rank: int, world_size: int +) -> str: + """Hash everything the profiled KV-cache memory value depends on. + + ``VllmConfig.compute_hash()`` covers the vLLM version and the model, + cache, parallel, and compilation configs, but deliberately contains no + device identity (``DeviceConfig.compute_hash`` is empty), so device + name, total memory, compute capability, and the torch/CUDA build are + added here. The vLLM version is also pinned as an explicit factor so + version invalidation holds no matter how ``compute_hash`` evolves. + Rank is included because per-rank memory use differs under TP/PP. + Driver-only changes are not part of the key; the free-memory gate at + apply time bounds the residual risk. + """ + # Imported here (as VllmConfig.compute_hash does) to avoid a cycle with + # the top-level vllm package. + from vllm import __version__ as vllm_version + + capability = current_platform.get_device_capability() + factors = { + "schema": PLAN_SCHEMA_VERSION, + "vllm": vllm_version, + "vllm_config": vllm_config.compute_hash(), + "device_name": current_platform.get_device_name(), + "device_total_memory": current_platform.get_device_total_memory(), + "device_capability": str(capability) if capability else "", + "torch": torch.__version__, + "cuda": torch.version.cuda or "", + "rank": rank, + "world_size": world_size, + } + digest = hashlib.sha256(json.dumps(factors, sort_keys=True).encode()).hexdigest() + return digest[:16] + + +def _plan_path(fingerprint: str) -> str: + """Plans are regenerable derived state, so they live under the standard + vLLM cache root (like the torch.compile cache) and relocate with + ``VLLM_CACHE_ROOT`` instead of needing a location knob of their own.""" + # VLLM_CACHE_ROOT is already user-expanded by envs.py. + return os.path.join( + envs.VLLM_CACHE_ROOT, "startup_plan", f"startup_plan_{fingerprint}.json" + ) + + +def _load_plan(fingerprint: str) -> dict | None: + """Load a plan for this fingerprint; None if absent or unreadable.""" + path = _plan_path(fingerprint) + try: + with open(path) as f: + plan = json.load(f) + except FileNotFoundError: + return None + except (OSError, json.JSONDecodeError) as e: + logger.warning("Ignoring unreadable startup plan %s: %s", path, e) + return None + if ( + plan.get("schema") != PLAN_SCHEMA_VERSION + or plan.get("fingerprint") != fingerprint + ): + return None + return plan + + +def _applicable_kv_cache_memory_bytes( + plan: dict, current_free_memory: int +) -> int | None: + """The apply-time OOM-safety gate. + + The recorded value is only valid if the device has at least as much + free memory now as when the plan was measured (co-tenants, leaked + allocations, or MIG changes all reduce it). Outside that envelope, + return None and let the caller re-profile. + """ + kv_bytes = plan.get("kv_cache_memory_bytes") + baseline = plan.get("free_memory_baseline") + if not isinstance(kv_bytes, int) or not isinstance(baseline, int): + return None + if kv_bytes <= 0: + return None + if current_free_memory < baseline: + logger.info( + "Startup plan not applied: current free memory (%.2f GiB) is " + "below the recorded baseline (%.2f GiB); falling back to full " + "memory profiling.", + current_free_memory / (1 << 30), + baseline / (1 << 30), + ) + return None + return kv_bytes + + +def maybe_apply_startup_plan(worker: "Worker") -> None: + """If enabled and ``--kv-cache-memory`` was not set explicitly, apply a + persisted plan by setting ``worker.cache_config.kv_cache_memory_bytes``. + No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``.""" + if ( + not envs.VLLM_ENABLE_STARTUP_PLAN + or worker.cache_config.kv_cache_memory_bytes is not None + ): + return + fingerprint = compute_plan_fingerprint( + worker.vllm_config, worker.rank, worker.parallel_config.world_size + ) + plan = _load_plan(fingerprint) + if plan is None: + return + current_free_memory = worker.init_snapshot.free_memory + kv_bytes = _applicable_kv_cache_memory_bytes(plan, current_free_memory) + if kv_bytes is None: + return + logger.info( + "Applying persisted startup plan (fingerprint %s): " + "kv_cache_memory_bytes=%d (%.2f GiB), recorded free-memory " + "baseline %.2f GiB, current %.2f GiB. Memory profiling will " + "be skipped.", + fingerprint, + kv_bytes, + kv_bytes / (1 << 30), + plan["free_memory_baseline"] / (1 << 30), + current_free_memory / (1 << 30), + ) + worker.cache_config.kv_cache_memory_bytes = kv_bytes + + +def maybe_save_startup_plan(worker: "Worker", kv_cache_memory_bytes: int) -> None: + """Atomically persist this boot's profiling result for future boots. + No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``; failures are logged, + never raised.""" + if not envs.VLLM_ENABLE_STARTUP_PLAN: + return + fingerprint = compute_plan_fingerprint( + worker.vllm_config, worker.rank, worker.parallel_config.world_size + ) + path = _plan_path(fingerprint) + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + payload = { + "schema": PLAN_SCHEMA_VERSION, + "fingerprint": fingerprint, + "kv_cache_memory_bytes": int(kv_cache_memory_bytes), + "free_memory_baseline": int(worker.init_snapshot.free_memory), + } + tmp = f"{path}.tmp.{os.getpid()}" + with open(tmp, "w") as f: + json.dump(payload, f) + os.replace(tmp, path) + logger.info("Saved startup plan to %s", path) + except OSError as e: + logger.warning("Failed to save startup plan to %s: %s", path, e)