A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 23:00 UTC
[PluggableLayer][1/N] Define PluggableLayer (#32331)
Signed-off-by: whx-sjtu <2952154980@qq.com>
whx committed
Jan 21, 2026 at 00:19 UTC
4ca62a0dbdba8b8aaf40438c99d20cdb56db8d5d
4 files changed
+99
-39
docs/design/custom_op.md
-9
@@ -8,15 +8,6 @@ This document will introduce how CustomOp works in vLLM and how to implement a n
8
9
`CustomOp` manages two dictionaries of all custom ops (i.e., op classes, indexed by registered name) in its class, for vLLM and OOT plugins respectively.
10
11
-??? code
12
-
13
- ```python
14
- class CustomOp(nn.Module):
15
-
16
- op_registry: dict[str, type["CustomOp"]] = {}
17
- op_registry_oot: dict[str, type["CustomOp"]] = {}
18
- ```
19
-
11
We can use `@CustomOp.register("op_name")` to register an op class to the `CustomOp` system. After this, the `op_name` and its class will be added into the `op_registry` dictionary. In addition, We can also register an OOT op by `@CustomOp.register_oot("op_name")`. We will introduce this mechanism in detail later.
12
13
When a `CustomOp` is called (i.e., call its `forward()` method), if it is enabled (i.e., with `--compilation_config.custom_ops '["+op_name"]'`), it will automatically dispatch the forward method to the appropriate backend according to `current_platform`. Otherwise (i.e., it is disabled), it will only call the `forward_native()` method to use PyTorch-native implementation of this forward method.
tests/model_executor/test_enabled_custom_ops.py
+5
-5
@@ -11,7 +11,7 @@ from vllm.config import (
11
get_cached_compilation_config,
12
set_current_vllm_config,
13
)
14
-from vllm.model_executor.custom_op import CustomOp
14
+from vllm.model_executor.custom_op import CustomOp, op_registry
15
from vllm.model_executor.layers.activation import (
16
GeluAndMul,
17
ReLUSquaredActivation,
@@ -98,17 +98,17 @@ def test_enabled_ops(
98
ops_enabled = [bool(x) for x in ops_enabled]
99
100
assert RMSNorm(1024).enabled() == ops_enabled[0]
101
- assert CustomOp.op_registry["rms_norm"].enabled() == ops_enabled[0]
101
+ assert op_registry["rms_norm"].enabled() == ops_enabled[0]
102
103
assert SiluAndMul().enabled() == ops_enabled[1]
104
- assert CustomOp.op_registry["silu_and_mul"].enabled() == ops_enabled[1]
104
+ assert op_registry["silu_and_mul"].enabled() == ops_enabled[1]
105
106
assert GeluAndMul().enabled() == ops_enabled[2]
107
- assert CustomOp.op_registry["gelu_and_mul"].enabled() == ops_enabled[2]
107
+ assert op_registry["gelu_and_mul"].enabled() == ops_enabled[2]
108
109
# If registered, subclasses should follow their own name
110
assert Relu3().enabled() == ops_enabled[3]
111
- assert CustomOp.op_registry["relu3"].enabled() == ops_enabled[3]
111
+ assert op_registry["relu3"].enabled() == ops_enabled[3]
112
113
# Unregistered subclass
114
class SiluAndMul2(SiluAndMul):
vllm/model_executor/custom_op.py
+86
-14
@@ -11,6 +11,86 @@ from vllm.platforms import current_platform
11
logger = init_logger(__name__)
12
13
14
+# Dictionary of all custom ops (classes, indexed by registered name).
15
+# To check if an op with a name is enabled, call .enabled() on the class.
16
+# Examples:
17
+# - MyOp.enabled()
18
+# - op_registry["my_op"].enabled()
19
+op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
20
+op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
21
+
22
+
23
+class PluggableLayer(nn.Module):
24
+ """
25
+ Base class for pluggable layers.
26
+
27
+ A PluggableLayer is a *module-composing* abstraction: it may instantiate other
28
+ ``torch.nn.Module`` objects as sub-layers, and its functionality depends on
29
+ these sub-layers following a generalized invocation sequence. Also, it is stateful
30
+ and may hold parameters or buffers.
31
+
32
+ Unlike :class:`CustomOp`, PluggableLayer does NOT provide per-platform
33
+ ``forward_*`` dispatch. Instead, it supports out-of-tree (OOT) replacement
34
+ of the entire layer class at instantiation time, allowing customized
35
+ initialization and submodule composition.
36
+ """
37
+
38
+ def __new__(cls, *args, **kwargs):
39
+ try:
40
+ layer_class_name = cls.__name__
41
+ except AttributeError:
42
+ raise TypeError(
43
+ f"Cannot instantiate '{cls.__name__}': its 'name' attribute "
44
+ f"was not set, possibly because it was not decorated with "
45
+ f"@PluggableLayer.register, or it's the PluggableLayer itself."
46
+ ) from None
47
+
48
+ if layer_class_name not in op_registry_oot:
49
+ layer_cls_to_instantiate = cls
50
+ else:
51
+ layer_cls_to_instantiate = op_registry_oot[layer_class_name]
52
+ logger.debug(
53
+ "Instantiating pluggable layer: %s using %s",
54
+ layer_class_name,
55
+ str(layer_cls_to_instantiate),
56
+ )
57
+ return super().__new__(layer_cls_to_instantiate)
58
+
59
+ # Decorator to register pluggable layers.
60
+ @classmethod
61
+ def register(cls, name: str):
62
+ def decorator(op_cls):
63
+ assert name not in op_registry, f"Duplicate op name: {name}"
64
+ op_cls.name = name
65
+ op_registry[name] = op_cls
66
+ return op_cls
67
+
68
+ return decorator
69
+
70
+ # Decorator to register out-of-tree(oot) pluggable layers.
71
+ # For OOT pluggable layers:
72
+ # if in-tree layer class is registered with an oot_custom_layer,
73
+ # the oot_custom_layer will be used instead.
74
+ @classmethod
75
+ def register_oot(cls, _decorated_layer_cls=None, name: str | None = None):
76
+ def decorator(layer_cls):
77
+ reg_name = name if name is not None else cls.__name__
78
+ assert reg_name not in op_registry_oot, f"Duplicate layer name: {reg_name}"
79
+ layer_cls.name = reg_name
80
+ op_registry_oot[reg_name] = layer_cls
81
+ return layer_cls
82
+
83
+ if _decorated_layer_cls is None:
84
+ # Called with parentheses: @PluggableLayer.register_oot()
85
+ # or @PluggableLayer.register_oot(name="...")
86
+ return decorator
87
+ elif isinstance(_decorated_layer_cls, type): # Check if it's a class
88
+ # Called without parentheses: @PluggableLayer.register_oot
89
+ return decorator(_decorated_layer_cls)
90
+ else:
91
+ raise TypeError("Decorator can only be applied to classes.")
92
+
93
+
94
class CustomOp(nn.Module):
95
"""
96
Base class for custom ops.
@@ -27,10 +107,10 @@ class CustomOp(nn.Module):
107
f"@CustomOp.register, or it's the CustomOp base class itself."
108
) from None
109
30
- if op_name not in cls.op_registry_oot:
110
+ if op_name not in op_registry_oot:
111
op_cls_to_instantiate = cls
112
else:
33
- op_cls_to_instantiate = cls.op_registry_oot[op_name]
113
+ op_cls_to_instantiate = op_registry_oot[op_name]
114
logger.debug(
115
"Instantiating custom op: %s using %s",
116
op_name,
@@ -150,21 +230,13 @@ class CustomOp(nn.Module):
230
231
return not count_none > 0 or count_all > 0
232
153
- # Dictionary of all custom ops (classes, indexed by registered name).
154
- # To check if an op with a name is enabled, call .enabled() on the class.
155
- # Examples:
156
- # - MyOp.enabled()
157
- # - op_registry["my_op"].enabled()
158
- op_registry: dict[str, type["CustomOp"]] = {}
159
- op_registry_oot: dict[str, type["CustomOp"]] = {}
160
-
233
# Decorator to register custom ops.
234
@classmethod
235
def register(cls, name: str):
236
def decorator(op_cls):
165
- assert name not in cls.op_registry, f"Duplicate op name: {name}"
237
+ assert name not in op_registry, f"Duplicate op name: {name}"
238
op_cls.name = name
167
- cls.op_registry[name] = op_cls
239
+ op_registry[name] = op_cls
240
return op_cls
241
242
return decorator
@@ -182,9 +254,9 @@ class CustomOp(nn.Module):
254
def register_oot(cls, _decorated_op_cls=None, name: str | None = None):
255
def decorator(op_cls):
256
reg_name = name if name is not None else cls.__name__
185
- assert reg_name not in cls.op_registry_oot, f"Duplicate op name: {reg_name}"
257
+ assert reg_name not in op_registry_oot, f"Duplicate op name: {reg_name}"
258
op_cls.name = reg_name
187
- cls.op_registry_oot[reg_name] = op_cls
259
+ op_registry_oot[reg_name] = op_cls
260
return op_cls
261
262
if _decorated_op_cls is None:
vllm/model_executor/layers/mla.py
+8
-11
@@ -6,7 +6,7 @@ import torch
6
7
from vllm.attention.layer import MLAAttention
8
from vllm.config import CacheConfig
9
-from vllm.model_executor.custom_op import CustomOp
9
+from vllm.model_executor.custom_op import PluggableLayer
10
from vllm.model_executor.layers.quantization import QuantizationConfig
11
12
@@ -30,13 +30,13 @@ class MLAModules:
30
31
32
# --8<-- [start:multi_head_latent_attention]
33
-@CustomOp.register("multi_head_latent_attention")
34
-class MultiHeadLatentAttentionWrapper(CustomOp):
35
- """MLA layer registered as CustomOp to allow OOT backends to add
33
+@PluggableLayer.register("multi_head_latent_attention")
34
+class MultiHeadLatentAttentionWrapper(PluggableLayer):
35
+ """Pluggable MLA layer which allows OOT backends to add
36
custom implementations of the outer MLA layer (including rope & o_proj).
37
- Note that currently MLA ignores the enable/disable mechanism of CustomOp
38
- because there is only one in-tree implementation in forward_native.
39
- TODO: implement this with a new PluggableLayer mechanism.
37
+ Note that currently oot platforms can still use CustomOp.register_oot to
38
+ replace MLA layer entirly, although we use PluggableLayer to register
39
+ this layer now.
40
41
This class takes positions and hidden_states as input.
42
The input tensors can either contain prefill tokens or decode tokens.
@@ -110,7 +110,7 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
110
111
self.prefix = prefix
112
113
- def forward_native(
113
+ def forward(
114
self,
115
positions: torch.Tensor,
116
hidden_states: torch.Tensor,
@@ -174,6 +174,3 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
174
)
175
176
return self.o_proj(attn_out)[0]
177
-
178
- def forward_cuda(self, *args, **kwargs):
179
- return self.forward_native(*args, **kwargs)