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