Fix MHC import: don't import .torch from layers/mhc.py
The layers/mhc.py was trying to import kernels.mhc.torch which failed because our __init__.py was breaking the package. Instead, just import our mhc_torch_ops which has everything we need. Also fix __init__.py to explicitly import mhc_pre_torch and mhc_post_torch from .torch instead of using import *.
This commit is contained in:
@@ -8,8 +8,6 @@ from vllm.model_executor.custom_op import CustomOp
|
||||
|
||||
# Import our torch implementations (registers torch.ops.vllm.mhc_pre, etc.)
|
||||
import vllm.model_executor.kernels.mhc.mhc_torch_ops as _mhc_torch # noqa: F401
|
||||
# Also import the original torch impls (mhc_pre_torch, mhc_post_torch)
|
||||
import vllm.model_executor.kernels.mhc.torch as mhc_kernels # noqa: F401
|
||||
|
||||
|
||||
@CustomOp.register("mhc_pre")
|
||||
|
||||
Reference in New Issue
Block a user