[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>

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

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",

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)

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)