| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import nullcontext |
| 5 | from types import SimpleNamespace |
| 6 | from unittest.mock import patch |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from tests.utils import create_new_process_for_each_test |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.mem_constants import GiB_bytes |
| 14 | from vllm.v1.worker import gpu_worker, startup_plan |
| 15 | from vllm.v1.worker.gpu_worker import maybe_rocm_profiling_fallback |
| 16 | from vllm.v1.worker.startup_plan import ( |
| 17 | maybe_apply_startup_plan, |
| 18 | maybe_save_startup_plan, |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch): |
| 23 | """Profiling must use serving's thread count to keep Dynamo guards valid.""" |
| 24 | from torch._dynamo.testing import CompileCounter |
| 25 | |
| 26 | monkeypatch.delenv("OMP_NUM_THREADS", raising=False) |
| 27 | monkeypatch.setattr(gpu_worker, "has_ec_transfer", lambda: False) |
| 28 | monkeypatch.setattr( |
| 29 | gpu_worker, "set_current_vllm_config", lambda config: nullcontext() |
| 30 | ) |
| 31 | loading_threads = [] |
| 32 | worker = SimpleNamespace( |
| 33 | vllm_config=SimpleNamespace(weight_transfer_config=None), |
| 34 | model_runner=SimpleNamespace( |
| 35 | load_model=lambda **kwargs: loading_threads.append(torch.get_num_threads()) |
| 36 | ), |
| 37 | _maybe_get_memory_pool_context=lambda **kwargs: nullcontext(), |
| 38 | _scoped_allocator_max_split=lambda **kwargs: nullcontext(), |
| 39 | ) |
| 40 | original_threads = torch.get_num_threads() |
| 41 | try: |
| 42 | torch.set_num_threads(2) |
| 43 | gpu_worker.Worker.load_model(worker) |
| 44 | assert loading_threads == [2] |
| 45 | |
| 46 | counter = CompileCounter() |
| 47 | compiled = torch.compile(lambda x: x + 1, backend=counter, fullgraph=True) |
| 48 | x = torch.ones(2) |
| 49 | compiled(x) |
| 50 | gpu_worker.set_torch_threads_for_runtime() |
| 51 | torch.testing.assert_close(compiled(x), x + 1) |
| 52 | assert counter.frame_count == 1 |
| 53 | finally: |
| 54 | torch.set_num_threads(original_threads) |
| 55 | |
| 56 | |
| 57 | @pytest.mark.skipif( |
| 58 | not current_platform.is_cuda_alike() |
| 59 | or not torch.accelerator.is_available() |
| 60 | or torch.cuda.memory.get_allocator_backend() != "native", |
| 61 | reason="needs the native CUDA or ROCm allocator", |
| 62 | ) |
| 63 | # A fresh allocator: blocks cached by earlier tests could serve the large buffer. |
| 64 | @create_new_process_for_each_test("spawn") |
| 65 | def test_scoped_max_split_keeps_freed_large_blocks_releasable(): |
| 66 | """A small allocation made after a large buffer is freed must not pin the |
| 67 | buffer's segment: the profiling run (determine_available_memory) grows |
| 68 | workspaces this way, and a pinned segment survives empty_cache() and is |
| 69 | counted as consumed memory.""" |
| 70 | large = 512 * 1024 * 1024 |
| 71 | small = 2 * 1024 * 1024 # large pool, so it is served by splitting cached blocks |
| 72 | |
| 73 | def reserved_while_small_is_live(scope) -> int: |
| 74 | torch.accelerator.empty_cache() |
| 75 | with scope: |
| 76 | buf = torch.empty(large, dtype=torch.uint8, device="cuda") |
| 77 | del buf |
| 78 | tensor = torch.empty(small, dtype=torch.uint8, device="cuda") |
| 79 | torch.accelerator.empty_cache() |
| 80 | reserved = torch.accelerator.memory_reserved() |
| 81 | del tensor |
| 82 | torch.accelerator.empty_cache() |
| 83 | return reserved |
| 84 | |
| 85 | baseline = torch.accelerator.memory_reserved() |
| 86 | # Without the limit the small tensor is split off the freed block and pins it. |
| 87 | assert reserved_while_small_is_live(nullcontext()) - baseline >= large |
| 88 | scoped = gpu_worker.Worker._scoped_allocator_max_split( |
| 89 | SimpleNamespace(), max_split_size_mb=20 |
| 90 | ) |
| 91 | assert reserved_while_small_is_live(scoped) - baseline < large |
| 92 | |
| 93 | |
| 94 | @pytest.mark.skipif( |
| 95 | not current_platform.is_cuda_alike() |
| 96 | or not torch.accelerator.is_available() |
| 97 | or torch.cuda.memory.get_allocator_backend() != "native", |
| 98 | reason="needs the native CUDA or ROCm allocator", |
| 99 | ) |
| 100 | @pytest.mark.parametrize("fail", [False, True]) |
| 101 | @pytest.mark.parametrize("snapshot_fallback", [False, True]) |
| 102 | @pytest.mark.parametrize("suffix", ["", ", "]) |
| 103 | def test_scoped_max_split_preserves_allocator_settings( |
| 104 | monkeypatch, fail, snapshot_fallback, suffix |
| 105 | ): |
| 106 | """Preserve allocator settings across successful and failed profiling runs.""" |
| 107 | if snapshot_fallback: |
| 108 | monkeypatch.delattr( |
| 109 | torch._C, "_accelerator_getAllocatorSettings", raising=False |
| 110 | ) |
| 111 | |
| 112 | def settings(): |
| 113 | return torch.cuda.memory._snapshot()["allocator_settings"] |
| 114 | |
| 115 | original = settings()["PYTORCH_CUDA_ALLOC_CONF"] |
| 116 | configured = ( |
| 117 | "max_split_size_mb:128,garbage_collection_threshold:0.8," |
| 118 | "roundup_power2_divisions:[256:1,512:2,>:4],max_non_split_rounding_mb:32" |
| 119 | ) |
| 120 | try: |
| 121 | torch._C._accelerator_setAllocatorSettings(configured + suffix) |
| 122 | before = settings() |
| 123 | expected_error = pytest.raises(RuntimeError, match="profiling failed") |
| 124 | with ( |
| 125 | expected_error if fail else nullcontext(), |
| 126 | gpu_worker.Worker._scoped_allocator_max_split(SimpleNamespace(), 20), |
| 127 | ): |
| 128 | scoped = settings() |
| 129 | assert scoped["max_split_size"] == 20 * 1024 * 1024 |
| 130 | for key in ("garbage_collection_threshold", "roundup_power2_divisions"): |
| 131 | assert scoped[key] == before[key] |
| 132 | if fail: |
| 133 | raise RuntimeError("profiling failed") |
| 134 | assert settings() == before |
| 135 | finally: |
| 136 | torch._C._accelerator_setAllocatorSettings(original) |
| 137 | |
| 138 | |
| 139 | def test_scoped_max_split_ignores_async_allocator(monkeypatch): |
| 140 | """Async allocators ignore max_split and may not support memory snapshots.""" |
| 141 | monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) |
| 142 | monkeypatch.setattr( |
| 143 | torch.cuda.memory, "get_allocator_backend", lambda: "cudaMallocAsync" |
| 144 | ) |
| 145 | with ( |
| 146 | patch.object(torch.cuda.memory, "_snapshot") as snapshot, |
| 147 | patch.object(torch._C, "_accelerator_setAllocatorSettings") as set_settings, |
| 148 | gpu_worker.Worker._scoped_allocator_max_split(SimpleNamespace(), 20), |
| 149 | ): |
| 150 | pass |
| 151 | snapshot.assert_not_called() |
| 152 | set_settings.assert_not_called() |
| 153 | |
| 154 | |
| 155 | # Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and |
| 156 | # saved by Worker.determine_available_memory / compile_or_warm_up_model. |
| 157 | |
| 158 | |
| 159 | def _plan_worker(config_hash="abc123", free_memory=78 * GiB_bytes, kv_bytes=None): |
| 160 | """The minimal Worker surface the startup-plan entry points touch.""" |
| 161 | return SimpleNamespace( |
| 162 | vllm_config=SimpleNamespace(compute_hash=lambda: config_hash), |
| 163 | rank=0, |
| 164 | parallel_config=SimpleNamespace(world_size=1), |
| 165 | init_snapshot=SimpleNamespace(free_memory=free_memory), |
| 166 | cache_config=SimpleNamespace(kv_cache_memory_bytes=kv_bytes), |
| 167 | ) |
| 168 | |
| 169 | |
| 170 | def _plan_platform(name="NVIDIA H100 PCIe"): |
| 171 | return SimpleNamespace( |
| 172 | get_device_name=lambda device_id=0: name, |
| 173 | get_device_total_memory=lambda device_id=0: 80 * GiB_bytes, |
| 174 | get_device_capability=lambda device_id=0: (9, 0), |
| 175 | ) |
| 176 | |
| 177 | |
| 178 | @pytest.fixture |
| 179 | def plan_env(monkeypatch: pytest.MonkeyPatch, tmp_path): |
| 180 | """Enable the startup plan, isolated under a tmp cache root.""" |
| 181 | monkeypatch.setenv("VLLM_ENABLE_STARTUP_PLAN", "1") |
| 182 | monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) |
| 183 | with patch.object(startup_plan, "current_platform", _plan_platform()): |
| 184 | yield |
| 185 | |
| 186 | |
| 187 | def test_startup_plan_fingerprint_sensitivity(plan_env): |
| 188 | """The fingerprint is the OOM-safety key: stable for identical inputs, |
| 189 | different for anything the profiled value depends on.""" |
| 190 | fp = startup_plan.compute_plan_fingerprint |
| 191 | base = fp(_plan_worker().vllm_config, 0, 1) |
| 192 | assert base == fp(_plan_worker().vllm_config, 0, 1) |
| 193 | assert base != fp(_plan_worker("other").vllm_config, 0, 1) |
| 194 | assert base != fp(_plan_worker().vllm_config, 1, 2) |
| 195 | with patch.object(startup_plan, "current_platform", _plan_platform("NVIDIA A100")): |
| 196 | assert base != fp(_plan_worker().vllm_config, 0, 1) |
| 197 | with patch("vllm.__version__", "0.0.0+plan-test"): |
| 198 | assert base != fp(_plan_worker().vllm_config, 0, 1) |
| 199 | |
| 200 | |
| 201 | def test_startup_plan_apply_gate(plan_env): |
| 202 | """Only a fingerprint-matching, memory-safe plan is ever applied.""" |
| 203 | maybe_save_startup_plan(_plan_worker(), 50 * GiB_bytes) |
| 204 | |
| 205 | applied = _plan_worker() |
| 206 | maybe_apply_startup_plan(applied) |
| 207 | assert applied.cache_config.kv_cache_memory_bytes == 50 * GiB_bytes |
| 208 | |
| 209 | less_memory = _plan_worker(free_memory=60 * GiB_bytes) |
| 210 | other_config = _plan_worker(config_hash="zzz999") |
| 211 | for refused in (less_memory, other_config): |
| 212 | maybe_apply_startup_plan(refused) |
| 213 | assert refused.cache_config.kv_cache_memory_bytes is None |
| 214 | |
| 215 | # An explicit --kv-cache-memory is never overridden. |
| 216 | explicit = _plan_worker(kv_bytes=7 * GiB_bytes) |
| 217 | maybe_apply_startup_plan(explicit) |
| 218 | assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes |
| 219 | |
| 220 | |
| 221 | # Memory accounting of the profiling run (Worker.determine_available_memory). |
| 222 | |
| 223 | # The fallback reads only the sign of the measured drop and this process's torch |
| 224 | # reservation; free memory is only logged, so no amount here is a device size. |
| 225 | ANY_FREE_MEMORY = 8 * GiB_bytes |
| 226 | MEASURED_DROP = 4 * GiB_bytes |
| 227 | TORCH_RESERVED = 3 * GiB_bytes |
| 228 | RELEASED_BY_OTHERS = 2 * GiB_bytes |
| 229 | |
| 230 | |
| 231 | def _snapshot(free_memory, torch_memory=0): |
| 232 | return SimpleNamespace(free_memory=free_memory, torch_memory=torch_memory) |
| 233 | |
| 234 | |
| 235 | def _profile_result(consumed, reserved_before=0, reserved_after=0): |
| 236 | """A result whose free-memory readings agree with `consumed`, which |
| 237 | `memory_profiling` derives as the drop in free memory, negative when it grew.""" |
| 238 | return SimpleNamespace( |
| 239 | total_consumed=consumed, |
| 240 | transient_peak_headroom=0, |
| 241 | before_create=_snapshot(ANY_FREE_MEMORY, reserved_before), |
| 242 | after_profile=_snapshot(ANY_FREE_MEMORY - consumed, reserved_after), |
| 243 | ) |
| 244 | |
| 245 | |
| 246 | @pytest.fixture |
| 247 | def rocm(request): |
| 248 | with patch.object( |
| 249 | gpu_worker, "current_platform", SimpleNamespace(is_rocm=lambda: request.param) |
| 250 | ): |
| 251 | yield request.param |
| 252 | |
| 253 | |
| 254 | @pytest.mark.parametrize("rocm", [True, False], indirect=True) |
| 255 | def test_profiling_fallback_declines_when_free_memory_dropped(rocm): |
| 256 | """The profiling measurement is kept as-is whenever free memory dropped.""" |
| 257 | result = _profile_result(consumed=MEASURED_DROP) |
| 258 | |
| 259 | assert maybe_rocm_profiling_fallback(result) is None |
| 260 | |
| 261 | |
| 262 | @pytest.mark.parametrize("rocm", [True], indirect=True) |
| 263 | def test_profiling_fallback_replaces_a_released_measurement(rocm): |
| 264 | """A negative measurement describes the rest of the device, so it is replaced |
| 265 | by this process's reservation, which the rest of the device cannot move.""" |
| 266 | result = _profile_result( |
| 267 | consumed=-RELEASED_BY_OTHERS, |
| 268 | reserved_after=TORCH_RESERVED, |
| 269 | ) |
| 270 | |
| 271 | assert maybe_rocm_profiling_fallback(result) == TORCH_RESERVED |
| 272 | |
| 273 | |
| 274 | @pytest.mark.parametrize("rocm", [True], indirect=True) |
| 275 | def test_profiling_fallback_never_returns_a_negative_amount(rocm): |
| 276 | """A reservation that shrank across the run cannot become negative usage.""" |
| 277 | result = _profile_result( |
| 278 | consumed=-RELEASED_BY_OTHERS, |
| 279 | reserved_before=TORCH_RESERVED, |
| 280 | reserved_after=0, |
| 281 | ) |
| 282 | |
| 283 | assert maybe_rocm_profiling_fallback(result) == 0 |
| 284 | |
| 285 | |
| 286 | @pytest.mark.parametrize("rocm", [False], indirect=True) |
| 287 | def test_profiling_fallback_declines_off_rocm(rocm): |
| 288 | """Platforms that account frees eagerly keep reporting the error, so the |
| 289 | caller's assertion stays reachable there.""" |
| 290 | result = _profile_result(consumed=-RELEASED_BY_OTHERS) |
| 291 | |
| 292 | assert maybe_rocm_profiling_fallback(result) is None |
| 293 | |
| 294 | |
| 295 | class _OrderedHandle: |
| 296 | """Send handle that logs when it is waited.""" |
| 297 | |
| 298 | def __init__(self, log: list[str], name: str): |
| 299 | self.log = log |
| 300 | self.name = name |
| 301 | |
| 302 | def is_completed(self) -> bool: |
| 303 | return True |
| 304 | |
| 305 | def wait(self) -> None: |
| 306 | self.log.append(f"wait:{self.name}") |
| 307 | |
| 308 | |
| 309 | def test_execute_model_waits_previous_pp_send_before_forward( |
| 310 | monkeypatch: pytest.MonkeyPatch, |
| 311 | ): |
| 312 | """Previous device handles are waited before the forward pass; the |
| 313 | metadata handle is left to the GroupCoordinator's reaper.""" |
| 314 | import torch |
| 315 | |
| 316 | from vllm.sequence import IntermediateTensors |
| 317 | |
| 318 | log: list[str] = [] |
| 319 | previous_tensor_send = _OrderedHandle(log, "prev-tensor") |
| 320 | metadata_handle = _OrderedHandle(log, "meta") |
| 321 | tensor_handle = _OrderedHandle(log, "tensor") |
| 322 | |
| 323 | def isend_tensor_dict(tensors, all_gather_group=None, all_gather_tensors=None): |
| 324 | log.append("isend") |
| 325 | return [metadata_handle, tensor_handle] |
| 326 | |
| 327 | pp_group = SimpleNamespace( |
| 328 | is_first_rank=True, |
| 329 | is_last_rank=False, |
| 330 | isend_tensor_dict=isend_tensor_dict, |
| 331 | ) |
| 332 | monkeypatch.setattr(gpu_worker, "get_pp_group", lambda: pp_group) |
| 333 | monkeypatch.setattr(gpu_worker, "get_tp_group", lambda: SimpleNamespace()) |
| 334 | |
| 335 | def run_model(scheduler_output, intermediate_tensors): |
| 336 | log.append("forward") |
| 337 | return IntermediateTensors({"hidden_states": torch.zeros(1)}) |
| 338 | |
| 339 | worker = SimpleNamespace( |
| 340 | vllm_config=SimpleNamespace( |
| 341 | compilation_config=SimpleNamespace( |
| 342 | pass_config=SimpleNamespace(enable_sp=False) |
| 343 | ), |
| 344 | parallel_config=SimpleNamespace( |
| 345 | pipeline_parallel_size=2, distributed_executor_backend="mp" |
| 346 | ), |
| 347 | ), |
| 348 | use_v2_model_runner=False, |
| 349 | model_runner=SimpleNamespace(execute_model=run_model), |
| 350 | annotate_profile=lambda scheduler_output: nullcontext(), |
| 351 | _pp_send_work=[previous_tensor_send], |
| 352 | ) |
| 353 | scheduler_output = SimpleNamespace( |
| 354 | total_num_scheduled_tokens=4, num_scheduled_tokens={"r0": 4} |
| 355 | ) |
| 356 | |
| 357 | assert gpu_worker.Worker.execute_model(worker, scheduler_output) is None |
| 358 | |
| 359 | assert log == ["wait:prev-tensor", "forward", "isend"] |
| 360 | assert worker._pp_send_work == [tensor_handle] |
| 361 | |
| 362 | |
| 363 | def test_jit_monitor_activation_follows_enable_jit_warmup( |
| 364 | monkeypatch: pytest.MonkeyPatch, |
| 365 | ): |
| 366 | """The post-warmup JIT monitor must stay off when JIT warmup is disabled |
| 367 | (e.g. by enforce_eager): runtime compilation is then expected, and |
| 368 | warning/erroring on it would be noise.""" |
| 369 | from vllm.utils import jit_monitor |
| 370 | |
| 371 | calls = [] |
| 372 | monkeypatch.setattr(jit_monitor, "activate", lambda **kwargs: calls.append(kwargs)) |
| 373 | |
| 374 | def worker(enable_jit_warmup): |
| 375 | return SimpleNamespace( |
| 376 | vllm_config=SimpleNamespace( |
| 377 | kernel_config=SimpleNamespace(enable_jit_warmup=enable_jit_warmup) |
| 378 | ), |
| 379 | observability_config=SimpleNamespace( |
| 380 | jit_monitor_mode="warn", jit_monitor_verbose=False |
| 381 | ), |
| 382 | ) |
| 383 | |
| 384 | gpu_worker.Worker._maybe_activate_jit_monitor(worker(True)) |
| 385 | assert calls == [{"mode": "warn", "verbose": False}] |
| 386 | |
| 387 | calls.clear() |
| 388 | gpu_worker.Worker._maybe_activate_jit_monitor(worker(False)) |
| 389 | assert calls == [] |