[PluggableLayer][1/N] Define PluggableLayer

[whx-sjtu]

EDIT: Reverted in #32725, reapplied in #32744

Purpose

First implementation of RFC #23786: Define PluggableLayer and apply to MLA as an example.

Test Plan

New ut waitted to be added.

Test Result

All ci should pass.

CC List

@ProExpertProg @wangxiyuan @jgong5 @Yikun

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a new CustomLayer abstraction to handle out-of-tree layer replacements, refactoring this logic out of CustomOp. Existing layers like FusedMoE and MultiHeadLatentAttentionWrapper are updated to use CustomLayer. My review focuses on a critical syntax error and improving the clarity of the new CustomLayer implementation.

[ProExpertProg]

Are we really not going to break OOT overriding of custom ops by removing register_oot from CustomOp?

[ProExpertProg]

Let's port the register_oot functionality and op_registry_oot over to the new layer and make them share the registry, but keep the functionality in CustomOp. That way CustomOp-based layers can slowly migrate over to the new layer abstraction without breaking anything

[whx-sjtu]

Currently this PR indeed breaks the original register_oot function. My original intention was to further differentiate the definitions of CustomLayer and CustomOp, and to redefine the usage of CustomOp. However, it seems the RFC #32358 proposed by @ProExpertProg offers a better solution to this. Based on this, I believe my PR can focus on defining PluggableLayer while preserving as much existing functionality as possible.

[whx-sjtu]

Let's port the register_oot functionality and op_registry_oot over to the new layer and make them share the registry, but keep the functionality in CustomOp. That way CustomOp-based layers can slowly migrate over to the new layer abstraction without breaking anything

I will migrate register_oot functionality and op_registry_oot to PluggableLayer. But I think the enalbed and default_on mechanism which is introduced to control fine-grained CustomOp enabling, is not needed by PluggableLayer anymore, right? @ProExpertProg

[mergify]

Documentation preview: https://vllm--32331.org.readthedocs.build/en/32331/

[mergify]

Hi @whx-sjtu, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[ProExpertProg]

Looks good to me! Let's see that tests pass, and then could you make a tracking issue with the list of all custom ops that should be ported over (the ones that only have 1 in-tree implementation)? Those with multiple in-tree implementations will be ported over later as we transition them to vLLM IR.

[wangxiyuan]

nit: type should be CustomOp or PluggableLayer

[wangxiyuan]

Looks good to me! Let's see that tests pass, and then could you make a tracking issue with the list of all custom ops that should be ported over (the ones that only have 1 in-tree implementation)? Those with multiple in-tree implementations will be ported over later as we transition them to vLLM IR.

++, it'll make the workflow more clear

[whx-sjtu]

Looks good to me! Let's see that tests pass, and then could you make a tracking issue with the list of all custom ops that should be ported over (the ones that only have 1 in-tree implementation)? Those with multiple in-tree implementations will be ported over later as we transition them to vLLM IR.

Sure

[mergify]

Hi @whx-sjtu, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[whx-sjtu]

The tracker list is here: #32676. Note that for some custom ops like short_conv, it has forward_cuda and an empty forward_native. I'm not sure whether it might have other in-tree implementations and assign it to PluggableLayer for now. Could you please review and make sure of it with other related maintainers @ProExpertProg ? Thanks a lot.

[mergify]

⚠️ The sha of the head commit of this PR conflicts with #32744. Mergify cannot evaluate rules on this PR. ⚠️