biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[PluggableLayer][1/N] Define PluggableLayer (Fix ci) (#32744)

Browse Source 
Signed-off-by: whx-sjtu <2952154980@qq.com>
This commit is contained in:
whx
2026-01-22 00:38:04 +08:00
committed by GitHub
parent 4e31b7f228
commit 1861ae8aae
7 changed files with 108 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/config/compilation.py
View File

@@ -1033,13 +1033,13 @@ class CompilationConfig:
# check if op name exists in model
op_name = op[1:]
if op_name not in all_ops_in_model:
from vllm.model_executor.custom_op import CustomOp
from vllm.model_executor.custom_op import op_registry
# Does op exist at all or is it just not present in this model?
# Note: Only imported op classes appear in the registry.
missing_str = (
"doesn't exist (or wasn't imported/registered)"
if op_name not in CustomOp.op_registry
if op_name not in op_registry
else "not present in model"
)

100
vllm/model_executor/custom_op.py
View File

@@ -11,6 +11,86 @@ from vllm.platforms import current_platform
logger = init_logger(__name__)
# Dictionary of all custom ops (classes, indexed by registered name).
# To check if an op with a name is enabled, call .enabled() on the class.
# Examples:
# - MyOp.enabled()
# - op_registry["my_op"].enabled()
op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
class PluggableLayer(nn.Module):
"""
Base class for pluggable layers.
A PluggableLayer is a *module-composing* abstraction: it may instantiate other
``torch.nn.Module`` objects as sub-layers, and its functionality depends on
these sub-layers following a generalized invocation sequence. Also, it is stateful
and may hold parameters or buffers.
Unlike :class:`CustomOp`, PluggableLayer does NOT provide per-platform
``forward_*`` dispatch. Instead, it supports out-of-tree (OOT) replacement
of the entire layer class at instantiation time, allowing customized
initialization and submodule composition.
"""
def __new__(cls, *args, **kwargs):
try:
layer_class_name = cls.__name__
except AttributeError:
raise TypeError(
f"Cannot instantiate '{cls.__name__}': its 'name' attribute "
f"was not set, possibly because it was not decorated with "
f"@PluggableLayer.register, or it's the PluggableLayer itself."
) from None
if layer_class_name not in op_registry_oot:
layer_cls_to_instantiate = cls
else:
layer_cls_to_instantiate = op_registry_oot[layer_class_name]
logger.debug(
"Instantiating pluggable layer: %s using %s",
layer_class_name,
str(layer_cls_to_instantiate),
)
return super().__new__(layer_cls_to_instantiate)
# Decorator to register pluggable layers.
@classmethod
def register(cls, name: str):
def decorator(op_cls):
assert name not in op_registry, f"Duplicate op name: {name}"
op_cls.name = name
op_registry[name] = op_cls
return op_cls
return decorator
# Decorator to register out-of-tree(oot) pluggable layers.
# For OOT pluggable layers:
# if in-tree layer class is registered with an oot_custom_layer,
# the oot_custom_layer will be used instead.
@classmethod
def register_oot(cls, _decorated_layer_cls=None, name: str | None = None):
def decorator(layer_cls):
reg_name = name if name is not None else cls.__name__
assert reg_name not in op_registry_oot, f"Duplicate layer name: {reg_name}"
layer_cls.name = reg_name
op_registry_oot[reg_name] = layer_cls
return layer_cls
if _decorated_layer_cls is None:
# Called with parentheses: @PluggableLayer.register_oot()
# or @PluggableLayer.register_oot(name="...")
return decorator
elif isinstance(_decorated_layer_cls, type): # Check if it's a class
# Called without parentheses: @PluggableLayer.register_oot
return decorator(_decorated_layer_cls)
else:
raise TypeError("Decorator can only be applied to classes.")
class CustomOp(nn.Module):
"""
Base class for custom ops.
@@ -27,10 +107,10 @@ class CustomOp(nn.Module):
f"@CustomOp.register, or it's the CustomOp base class itself."
) from None
if op_name not in cls.op_registry_oot:
if op_name not in op_registry_oot:
op_cls_to_instantiate = cls
else:
op_cls_to_instantiate = cls.op_registry_oot[op_name]
op_cls_to_instantiate = op_registry_oot[op_name]
logger.debug(
"Instantiating custom op: %s using %s",
op_name,
@@ -150,21 +230,13 @@ class CustomOp(nn.Module):
return not count_none > 0 or count_all > 0
# Dictionary of all custom ops (classes, indexed by registered name).
# To check if an op with a name is enabled, call .enabled() on the class.
# Examples:
# - MyOp.enabled()
# - op_registry["my_op"].enabled()
op_registry: dict[str, type["CustomOp"]] = {}
op_registry_oot: dict[str, type["CustomOp"]] = {}
# Decorator to register custom ops.
@classmethod
def register(cls, name: str):
def decorator(op_cls):
assert name not in cls.op_registry, f"Duplicate op name: {name}"
assert name not in op_registry, f"Duplicate op name: {name}"
op_cls.name = name
cls.op_registry[name] = op_cls
op_registry[name] = op_cls
return op_cls
return decorator
@@ -182,9 +254,9 @@ class CustomOp(nn.Module):
def register_oot(cls, _decorated_op_cls=None, name: str | None = None):
def decorator(op_cls):
reg_name = name if name is not None else cls.__name__
assert reg_name not in cls.op_registry_oot, f"Duplicate op name: {reg_name}"
assert reg_name not in op_registry_oot, f"Duplicate op name: {reg_name}"
op_cls.name = reg_name
cls.op_registry_oot[reg_name] = op_cls
op_registry_oot[reg_name] = op_cls
return op_cls
if _decorated_op_cls is None:

19
vllm/model_executor/layers/mla.py
View File

@@ -6,7 +6,7 @@ import torch
from vllm.attention.layer import MLAAttention
from vllm.config import CacheConfig
from vllm.model_executor.custom_op import CustomOp
from vllm.model_executor.custom_op import PluggableLayer
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -30,13 +30,13 @@ class MLAModules:
# --8<-- [start:multi_head_latent_attention]
@CustomOp.register("multi_head_latent_attention")
class MultiHeadLatentAttentionWrapper(CustomOp):
"""MLA layer registered as CustomOp to allow OOT backends to add
@PluggableLayer.register("multi_head_latent_attention")
class MultiHeadLatentAttentionWrapper(PluggableLayer):
"""Pluggable MLA layer which allows OOT backends to add
custom implementations of the outer MLA layer (including rope & o_proj).
Note that currently MLA ignores the enable/disable mechanism of CustomOp
because there is only one in-tree implementation in forward_native.
TODO: implement this with a new PluggableLayer mechanism.
Note that currently oot platforms can still use CustomOp.register_oot to
replace MLA layer entirly, although we use PluggableLayer to register
this layer now.
This class takes positions and hidden_states as input.
The input tensors can either contain prefill tokens or decode tokens.
@@ -110,7 +110,7 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
self.prefix = prefix
def forward_native(
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
@@ -174,6 +174,3 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
)
return self.o_proj(attn_out)[0]
def forward_cuda(self, *args, **kwargs):
return self.forward_native(*args, **kwargs)