@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 53d2420b44
Fork 0

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.