mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-12 21:15:00 +08:00
Fix PiecewiseCompileInterpreter (#17338)
Signed-off-by: rzou <zou3519@gmail.com>
This commit is contained in:
parent
de65fc8e1e
commit
aa42561e40
@ -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,
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user