From 94f4170df37fe29d68bd2f7e5a496501b78d55ae Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 22 Sep 2026 21:12:20 -0500 Subject: [PATCH] [Bugfix] Set worker runtime threads before profiling and compilation (#55891) Signed-off-by: Andreas Karatzas Co-authored-by: Codex Co-authored-by: Nick Hill --- tests/v1/worker/test_gpu_worker.py | 37 ++++++++++++++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 9 ++++---- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 88cbee6802..375e43810f 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -6,6 +6,7 @@ from types import SimpleNamespace from unittest.mock import patch import pytest +import torch from vllm.utils.mem_constants import GiB_bytes from vllm.v1.worker import gpu_worker, startup_plan @@ -15,6 +16,42 @@ from vllm.v1.worker.startup_plan import ( maybe_save_startup_plan, ) + +def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch): + """Profiling must use serving's thread count to keep Dynamo guards valid.""" + from torch._dynamo.testing import CompileCounter + + monkeypatch.delenv("OMP_NUM_THREADS", raising=False) + monkeypatch.setattr(gpu_worker, "has_ec_transfer", lambda: False) + monkeypatch.setattr( + gpu_worker, "set_current_vllm_config", lambda config: nullcontext() + ) + loading_threads = [] + worker = SimpleNamespace( + vllm_config=SimpleNamespace(weight_transfer_config=None), + model_runner=SimpleNamespace( + load_model=lambda **kwargs: loading_threads.append(torch.get_num_threads()) + ), + _maybe_get_memory_pool_context=lambda **kwargs: nullcontext(), + _scoped_allocator_max_split=lambda **kwargs: nullcontext(), + ) + original_threads = torch.get_num_threads() + try: + torch.set_num_threads(2) + gpu_worker.Worker.load_model(worker) + assert loading_threads == [2] + + counter = CompileCounter() + compiled = torch.compile(lambda x: x + 1, backend=counter, fullgraph=True) + x = torch.ones(2) + compiled(x) + gpu_worker.set_torch_threads_for_runtime() + torch.testing.assert_close(compiled(x), x + 1) + assert counter.frame_count == 1 + finally: + torch.set_num_threads(original_threads) + + # Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and # saved by Worker.determine_available_memory / compile_or_warm_up_model. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3a3c5314b6..3d4c5f3d5c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -550,6 +550,11 @@ class Worker(WorkerBase): self.model_runner.get_model(), ) + # Preserve parallel weight loading, then use serving's thread count + # for profiling and compilation so Dynamo's global-state guards remain + # valid when requests arrive. + set_torch_threads_for_runtime() + def update_config(self, overrides: dict[str, Any]) -> None: self.model_runner.update_config(overrides) @@ -978,10 +983,6 @@ class Worker(WorkerBase): # gate so subsequent `execute_model` / `sample_tokens` calls enforce it. enable_gpu_sync_check() - # Startup is done; steady-state serving gets no benefit from torch - # intra-op parallelism. - set_torch_threads_for_runtime() - return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time,