[RFC][Plugin]: support loading kernels from other place

[ILikeIneine]

Motivation.

Currently, all the compiled kernels are loaded in _custom_ops.py by:

import vllm._C
import vllm._moe_C
...

This has several limits:

• It needs to guarantee that the dynamic libraries are under the vllm's directory (usually under site-package).
• If plugin(or out_of_tree) wanna reuse this file and adapted or registered the same kernels, they have to copy their kernels from plugin to vllm's folder during the installation.

Proposed Change.

We want to offer a platform interface to support load the platform-specific kernel operations:

class Platform:
    @classmethod
    def import_general_kernels(cls) -> None:
        """ Import any platform-specific C kernels. """
        import vllm._C  # noqa: F401
        # import vllm.some_other_kernels
    
    @classmethod
    def import_moe_kernels(cls) -> None:
        """
        Import any platform-specific Mixture of Experts kernels.
        """
        import vllm._moe_C  # noqa: F401

and then, to load them by:

# in _custom_ops.py
try:
    current_platform.import_general_kernels()
except ImportError as e:
    logger.warning("Failed to load kernels with %r", e)

supports_moe_ops = False
with contextlib.suppress(ImportError):
    current_platform.import_moe_kernels()
    supports_moe_ops = True

I've made a draft pr in #25823, any comments or help would be appreciated!

Feedback Period.

9.28 ~ 10.8

CC List.

@NickLucche @youkaichao @mgoin @ProExpertProg @MatthewBonanni

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[irasin]

LGTM

[hmellor]

You can use https://github.com/huggingface/kernels to load portable kernels from other places.

Here's a nice blog that explains how to use it https://huggingface.co/blog/hello-hf-kernels

[ILikeIneine]

You can use https://github.com/huggingface/kernels to load portable kernels from other places.

Here's a nice blog that explains how to use it https://huggingface.co/blog/hello-hf-kernels

I don't think we hit the same situation. Basically our plugin backend could not directly use the cuda kernel, we need recompile (nearly) all the cuda kernels on maca backends (with some compatibility related modifications).

You can simply think of it as a cuda_alike backend that intergrated by oot_plugin, kind of like making a oot ROCm plugin and lots of things could be reused with vllm (and I believe in the future there may be more and more plugins or backend that will behave the same).

In the other hand we'd also need to maintain the kernels by our own and may add some additional optimizations on kernels.

In this way, what we only need to do is to keep the signature of these kernels same as vllm's. The modifications make it easier to reuse the common interfaces in _custom_ops.py

[hmellor]

Ah I see, I misunderstood your use case. Thank you for considering my suggestion nonetheless!