From aa42561e4054f43f41bf0ea564369f5ea3147316 Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Wed, 28 May 2025 04:40:53 -0400 Subject: [PATCH] Fix PiecewiseCompileInterpreter (#17338) Signed-off-by: rzou --- vllm/compilation/backends.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 8114cddcd9fa..0358c9d0d1b5 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -10,6 +10,7 @@ from typing import Any, Callable, Optional import torch import torch.fx as fx +from torch._dispatch.python import enable_python_dispatcher import vllm.envs as envs from vllm.config import CompilationConfig, VllmConfig @@ -269,7 +270,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t for t in args ] - with self.fake_mode: + with self.fake_mode, enable_python_dispatcher(): return super().run(*fake_args) def call_module(self, target: torch.fx.node.Target,