[Bugfix] Set worker runtime threads before profiling and compilation (#55891)
Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Andreas Karatzas committed
Sep 22, 2026 at 21:12 UTC
94f4170df37fe29d68bd2f7e5a496501b78d55ae
2 files changed
+42
-4
tests/v1/worker/test_gpu_worker.py
+37
@@ -6,6 +6,7 @@ from types import SimpleNamespace
6
from unittest.mock import patch
7
8
import pytest
9
+import torch
10
11
from vllm.utils.mem_constants import GiB_bytes
12
from vllm.v1.worker import gpu_worker, startup_plan
@@ -15,6 +16,42 @@ from vllm.v1.worker.startup_plan import (
16
maybe_save_startup_plan,
17
)
18
19
+
20
+def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch):
21
+ """Profiling must use serving's thread count to keep Dynamo guards valid."""
22
+ from torch._dynamo.testing import CompileCounter
23
+
24
+ monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
25
+ monkeypatch.setattr(gpu_worker, "has_ec_transfer", lambda: False)
26
+ monkeypatch.setattr(
27
+ gpu_worker, "set_current_vllm_config", lambda config: nullcontext()
28
+ )
29
+ loading_threads = []
30
+ worker = SimpleNamespace(
31
+ vllm_config=SimpleNamespace(weight_transfer_config=None),
32
+ model_runner=SimpleNamespace(
33
+ load_model=lambda **kwargs: loading_threads.append(torch.get_num_threads())
34
+ ),
35
+ _maybe_get_memory_pool_context=lambda **kwargs: nullcontext(),
36
+ _scoped_allocator_max_split=lambda **kwargs: nullcontext(),
37
+ )
38
+ original_threads = torch.get_num_threads()
39
+ try:
40
+ torch.set_num_threads(2)
41
+ gpu_worker.Worker.load_model(worker)
42
+ assert loading_threads == [2]
43
+
44
+ counter = CompileCounter()
45
+ compiled = torch.compile(lambda x: x + 1, backend=counter, fullgraph=True)
46
+ x = torch.ones(2)
47
+ compiled(x)
48
+ gpu_worker.set_torch_threads_for_runtime()
49
+ torch.testing.assert_close(compiled(x), x + 1)
50
+ assert counter.frame_count == 1
51
+ finally:
52
+ torch.set_num_threads(original_threads)
53
+
54
+
55
# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
56
# saved by Worker.determine_available_memory / compile_or_warm_up_model.
57
vllm/v1/worker/gpu_worker.py
+5
-4
@@ -550,6 +550,11 @@ class Worker(WorkerBase):
550
self.model_runner.get_model(),
551
)
552
553
+ # Preserve parallel weight loading, then use serving's thread count
554
+ # for profiling and compilation so Dynamo's global-state guards remain
555
+ # valid when requests arrive.
556
+ set_torch_threads_for_runtime()
557
+
558
def update_config(self, overrides: dict[str, Any]) -> None:
559
self.model_runner.update_config(overrides)
560
@@ -978,10 +983,6 @@ class Worker(WorkerBase):
983
# gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
984
enable_gpu_sync_check()
985
981
- # Startup is done; steady-state serving gets no benefit from torch
982
- # intra-op parallelism.
983
- set_torch_threads_for_runtime()
984
-
986
return CompilationTimes(
987
language_model=self.compilation_config.compilation_time,
988
encoder=self.compilation_config.encoder_compilation_time,