frozenleaves/vllm
1
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2025-10-20 23:03:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix PiecewiseCompileInterpreter (#17338)

Browse Source 
Signed-off-by: rzou <zou3519@gmail.com>
This commit is contained in:
Richard Zou
2025-05-28 04:40:53 -04:00
committed by GitHub
parent de65fc8e1e
commit aa42561e40
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/compilation/backends.py
View File

@ -10,6 +10,7 @@ from typing import Any, Callable, Optional
import torch import torch
import torch.fx as fx import torch.fx as fx
from torch._dispatch.python import enable_python_dispatcher
import vllm.envs as envs import vllm.envs as envs
from vllm.config import CompilationConfig, VllmConfig 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 self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
for t in args for t in args
] ]
with self.fake_mode: with self.fake_mode, enable_python_dispatcher():
return super().run(*fake_args) return super().run(*fake_args)
def call_module(self, target: torch.fx.node.Target, def call_module(self, target: torch.fx.node.Target,