"lib/vscode:/vscode.git/clone" did not exist on "9bc6d98d58f10f9671b71c72178fa46affa58c8e"
Unverified Commit f8c3af2d authored by Xiaoshuang Wang's avatar Xiaoshuang Wang Committed by GitHub
Browse files

[vLLM IR] add `import_ir_kernels()` to support OOT platforms (#38807)


Signed-off-by: default avatarIcey <1790571317@qq.com>
parent 50cd5674
...@@ -68,10 +68,13 @@ class IrOpPriorityConfig: ...@@ -68,10 +68,13 @@ class IrOpPriorityConfig:
def set_priority(self): def set_priority(self):
""" """
Context manager to set the IR op priority for all op members. 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.ir.op import IrOp
from vllm.platforms import current_platform
current_platform.import_ir_kernels()
with contextlib.ExitStack() as stack: with contextlib.ExitStack() as stack:
for field in fields(self): for field in fields(self):
......
...@@ -207,6 +207,15 @@ class Platform: ...@@ -207,6 +207,15 @@ class Platform:
""" """
return cls.simple_compile_backend 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 @classmethod
def device_id_to_physical_device_id(cls, device_id: int): def device_id_to_physical_device_id(cls, device_id: int):
# Treat empty device control env var as unset. This is a valid # Treat empty device control env var as unset. This is a valid
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment