A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 17:30 UTC
[Bugfix] tpu_model_runner: set vllm config context when calling reset_dynamo_cache() (#30331)
Signed-off-by: Daniele Trifirò <dtrifiro@redhat.com>
Daniele committed
Dec 10, 2025 at 13:58 UTC
53d2420b4447fbcab572dc23d2c3bb9224a8a561
1 file changed
+3
-2
vllm/v1/worker/tpu_worker.py
+3
-2
@@ -10,7 +10,7 @@ import torch
10
import torch.nn as nn
11
12
import vllm.envs as envs
13
-from vllm.config import VllmConfig
13
+from vllm.config import VllmConfig, set_current_vllm_config
14
from vllm.distributed import (
15
ensure_model_parallel_initialized,
16
init_distributed_environment,
@@ -207,7 +207,8 @@ class TPUWorker:
207
# one compiled bytecode. Having one FX graph/cached bytecode per
208
# compiled model is required for `support_torch_compile` decorator to
209
# skip dynamo guard.
210
- self.model_runner.reset_dynamo_cache()
210
+ with set_current_vllm_config(self.vllm_config):
211
+ self.model_runner.reset_dynamo_cache()
212
213
# Get the maximum amount of memory used by the model weights and
214
# intermediate activations.