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 15, 08:30 UTC
Fix PiecewiseCompileInterpreter (#17338)
Signed-off-by: rzou <zou3519@gmail.com>
Richard Zou committed
May 28, 2025 at 04:40 UTC
aa42561e4054f43f41bf0ea564369f5ea3147316
1 file changed
+2
-1
vllm/compilation/backends.py
+2
-1
@@ -10,6 +10,7 @@ from typing import Any, Callable, Optional
10
11
import torch
12
import torch.fx as fx
13
+from torch._dispatch.python import enable_python_dispatcher
14
15
import vllm.envs as envs
16
from vllm.config import CompilationConfig, VllmConfig
@@ -269,7 +270,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
270
self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
271
for t in args
272
]
272
- with self.fake_mode:
273
+ with self.fake_mode, enable_python_dispatcher():
274
return super().run(*fake_args)
275
276
def call_module(self, target: torch.fx.node.Target,