diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index f8494111d..f3ffbe4e8 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -68,10 +68,13 @@ class IrOpPriorityConfig: def set_priority(self): """ Context manager to set the IR op priority for all op members. - It also imports vllm.kernels to ensure all implementations are made available. + It also imports IR kernel implementations for the current platform + to ensure all implementations are made available. """ - import vllm.kernels # noqa: F401, registers IR op implementations from vllm.ir.op import IrOp + from vllm.platforms import current_platform + + current_platform.import_ir_kernels() with contextlib.ExitStack() as stack: for field in fields(self): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 1f54004f7..72b714df1 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -207,6 +207,15 @@ class Platform: """ return cls.simple_compile_backend + @classmethod + def import_ir_kernels(cls) -> None: + """ + The default implementation imports ``vllm.kernels``, which registers + the built-in IR op implementations. Out-of-tree (OOT) platforms should + override this method to import their own kernel modules. + """ + import vllm.kernels # noqa: F401 + @classmethod def device_id_to_physical_device_id(cls, device_id: int): # Treat empty device control env var as unset. This is a valid