[RFC]: Enhancing vLLM Plugin Architecture

[kzawora-intel]

Motivation.

This RFC proposes changes to the vLLM plugin infrastructure to address compatibility issues, simplify out-of-tree code maintenance, and improve overall user experience.

Problem Statement

1. Registering and Dispatching Semi-Custom Operators:

   • Vendors need a mechanism to register and dispatch custom implementations for generic layers like RotaryEmbedding and RMSNorm without modifying standard model code.
   • Current Limitations:
     • The existing forward_oot method in CustomOp requires monkey-patching, which makes plugin code vulnerable to upstream changes. For example, Ascend's RoPE implementation relies on layer properties that are not guaranteed, risking breakage with any refactor.
     • There is no straightforward way to dispatch new out-of-tree operators without altering model code, as seen with AscendFusedMoE and custom DeepSeek models.

2. Ensuring Out-Of-Tree Compatibility:

   • Upstream changes can break plugin compatibility, placing the burden on plugin maintainers to address these issues.
     • Platform workers and model runners consume vLLM APIs that are not guaranteed to be stable - e.g. vllm.v1.worker.gpu_input_batch.InputBatch, vllm.v1.sample.metadata.SamplingMetadata, vllm.sampling_params.SamplingType)
     • Additionally: Lack of basic functionality checks for plugin compatibility in upstream leads to uncertainty about whether the latest plugin will work with the latest vLLM, forcing users to mix and match versions.
   • Lack of clearly defined, immutable interfaces for communication between vLLM and out-of-tree plugins.
     • Both "vLLM-consumed" and "plugin-consumed" API definitions are crucial for stable integration.

3. No In-Tree Visibility:

   • Plugins live entirely out-of-tree, requiring users to clone, build, and install specific plugin repositories or packages.
   • Installing "plugin-less" vLLM via pip install vllm lacks any skeleton for out-of-tree platform support, and upstream development is also unaware of these backends.
     • If user chooses an out-of-tree backend without installing a proper plugin first, vLLM won't recognize that this platform is indeed supported via an out-of-tree plugin

Proposed Change.

1. Registering and Dispatching Semi-Custom Operators:

• Implement robust registration and automatic dispatch of out-of-tree operators. There is already a registration and dispatch logic for models with support for overriding (e.g. register_model) which can be leveraged by vendors.
• The integration could look as follows:

from vllm import OperatorRegistry
def register_operators():
    from .custom_ops import CustomRotaryEmbedding, CustomRMSNorm

    OperatorRegistry.register_operator(
        "rotary_embedding",
        "vllm_plugin.custom_ops:CustomRotaryEmbedding")

    OperatorRegistry.register_operator(
        "rms_norm",
        "vllm_plugin.custom_ops:CustomRMSNorm")

• The defined operators are expected to match APIs of their corresponding ops, and this can be easily achieved via inheritance, e.g:

from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
from vllm.logger import init_logger
logger = init_logger(__name__)

class CustomRMSNorm(RMSNorm):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Additional initialization for the dummy platform can go here.
        logger.warning("CustomRMSNorm initialized.")
    def forward(self, *args, **kwargs):
        logger.warning("CustomRMSNorm fwd pass.")
        return super().forward(*args, **kwargs)
    
class CustomRotaryEmbedding(RotaryEmbedding):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Additional initialization for the dummy platform can go here.
        logger.warning("CustomRotaryEmbedding initialized.")
    def forward(self, *args, **kwargs):
        logger.warning("CustomRotaryEmbedding fwd pass.")
        return super().forward(*args, **kwargs)

• The difference between this approach, and forward_oot is that vendors have freedom to either completely override the existing initialization logic, perform additional sanity checks, and explicitly define class attributes used in forward passes.

  • Drawbacks:
    - This will require refactor of the model code. Models should fetch the operator class by name from the registry rather than utilizing the class from vllm.model_executor.layers.
    - This will need to ensure that all in-tree operators are registered before any registration occurs out-of-tree.

2. Defined Platform Stubs and APIs:

• Provide clearly defined API stubs in upstream (e.g., vllm.platforms.hpu, vllm.platforms.npu) for plugin integration
  • Platform interface is a very good starting point, but vLLM should be aware of platforms that have established plugins and have lightweight "manifest" of them, extensible by out-of-tree code.
  • In-tree stub support with out-of-tree implementation.
• Require specific, versioned, strongly-typed definitions of exposed APIs and dataclasses, covering both consumed (e.g. vllm.v1.worker.gpu_input_batch.InputBatch, vllm.v1.sample.metadata.SamplingMetadata, vllm.sampling_params.SamplingType) and produced APIs (e.g. vllm_plugin.v1.worker.init_device, determine_available_memory)

3. CI/Code Scanning Checks:

• Implement basic CI/code scanning checks in upstream for plugin compatibility.
• Once platform stubs and produced/consumed APIs are defined, extensions can be covered by static code checks and mypy should detect incompatibilities.
• More ideal, long-term solution - actual unit test execution on the pluggable platforms in vLLM PRs

Feedback Period.

24/7

CC List.

@simon-mo @WoosukKwon @xuechendi

Any Other Things.

The problem statements and proposed solutions are ordered by priority - it's most crucial to address the points in the beginning (the custom ops). The proposed solutions for later points may evolve over time based on feedback.

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.

[yarongmu-google]

Enhancing the flexibility of the vLLM framework is a great idea; however, are we constrained w/ Plug-ins? Plug-ins are a concrete implementation of the DPI (Dependency Inversion Principle), and there are other design patterns/implementations also belonging to DIP but was created later than Plug-ins to be a more advanced version of it.

@mgoin too

[yarongmu-google]

A related proposal: #19376

[xuechendi]

Enhancing the flexibility of the vLLM framework is a great idea; however, are we constrained w/ Plug-ins? Plug-ins are a concrete implementation of the DPI (Dependency Inversion Principle), and there are other design patterns/implementations also belonging to DIP but was created later than Plug-ins to be a more advanced version of it.

@mgoin too

Thanks, @yarongmu-google , I read through the RFC #19376 , that will definitely be a good way to fill the gaps for plugin, especially we are now struggling about the customer scheduler and custom quantization method.

I want to align with you on the scope and timeline:
For current phase: we hope to have all necessary changes in this RFC merged as soon as possible, so we can migrate HPU to plugin with better customer layer support and CI.
For a long term: with more and more modules becoming 'DI' modules as #19376 proposed, we can continue involving plugin impl to fully utilize the flexibility.

[yarongmu-google]

Thanks @xuechendi. Glad that we are aligned on the benefits of DI - totally make sense to continue building on plug-ins for immediate needs and move to a more comprehensive framework later. Let's see what the other community members think :)

[joerunde]

Re: Ensuring Out-Of-Tree Compatibility

Our plugin (vllm-spyre) assumes that all vllm apis are unstable, since vllm is still on a 0.x release stream. We have CI checks that regularly check our plugin against the latest vllm main branch, and make updates whenever that breaks. It's kinda painful, but we recognize the need to refactor internal interfaces of vllm sometimes, and I don't know if this needs to be solved immediately.

It would be nice though if all new field additions to existing internal types (like vllm.v1.sample.metadata.SamplingMetadata) came with default values, so that we didn't have to keep adding new_field=None in our code.

I agree though that nailing down some interfaces for when vllm makes a v1 release is important

[joerunde]

Re: CI Checks

More ideal, long-term solution - actual unit test execution on the pluggable platforms in vLLM PRs

We could do this today with the spyre plugin, since we support a cpu execution mode. (It'd be a bit harder to get real hardware setup for each backend to run real tests though.) Not really sure what that would accomplish today though, because we wouldn't expect the PR author or other vllm maintainers to come and fix the code in our plugin to be forward-compatible.

If there's commitment for well-versioned APIs that don't break consumers though, then this would be pretty valuable to have as a signal.

[wangxiyuan]

I‘d like to list the problem vllm-ascend hit during these few months:

1. forword_oot is not enough. Usually we want to change the __init__ or some other function inner the ops well.
2. CI/Function Test is usually break due to the change from vllm. It contains:
   1. the api is changed.
   2. vllm-ascend import some vllm inner code from here and there, which makes the code more coupled
   3. some new added common code in vllm is hard code to CUDA.

Other problem we hit have been fixed with our contribution to vllm already. Such as:

1. [Compile][Platform] Make PiecewiseBackend pluggable and extendable #18076
2. [Platform][Dist] Make torch distributed process group extendable #18763

So from our experience, this RFC solved the problems all we hit currently.

1. The custom operators register mechanism is great. I think it's good to go first.
2. The api between vllm and plugin should be more clear and well designed. We should improve the guide/principle/interface to make it clear. Platform module is a great one, but it's not enough. I think we can create a new RFC to describe it with real code design. It should contain something like: what interface/function is for plugin, how it bump versions and so on.
3. CI test in vllm is required. At least the API change should be noticed.

action 1 is easy. I think we can do it right now.
action 2 need more design. @kzawora-intel I can help to do it with you
action 3 can be done following with action 2 when the new code is landed.

---

For the RFC #19376 it's a nice solution for system level. I'd like to see it can be done in the future. But the most urgent thing at the moment is solve the problem that plugin really hit. After that, we can move the architecture to DI.

my two cents

@Yikun @youkaichao @simon-mo

---

Besides the tech design/problem, there is another important thing for plugin IMO. The release plan/cycle for vllm should be stable and clear. It can help to improve the quality of plugins for develop/test/release. For example, if the plugin know that vllm will release in the next Monday, then in this week, we'll focus on functional test, speed up feature develop, so that the plugin can be released at the same time with high quality

[zou3519]

@wangxiyuan do you have an example of wanting to change __init__? If vLLM changes the signature of any inits, (like RotaryEmbedding.__init__), would that break you? The thing I am trying to get at is if the signature of these are now FC/BC surfaces.

[wangxiyuan]

@zou3519 here is an case: https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/ops/rotary_embedding.py#L281