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
Revert "[PluggableLayer][1/N] Define PluggableLayer" (#32725)
Robert Shaw committed
Jan 20, 2026 at 19:21 UTC
c78ee240b30d39af3b9343fb0b215b98588ee9bb
4 files changed
+39
-99
docs/design/custom_op.md
+9
@@ -8,6 +8,15 @@ 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
+
20
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.
21
22
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, op_registry
14
+from vllm.model_executor.custom_op import CustomOp
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 op_registry["rms_norm"].enabled() == ops_enabled[0]
101
+ assert CustomOp.op_registry["rms_norm"].enabled() == ops_enabled[0]
102
103
assert SiluAndMul().enabled() == ops_enabled[1]
104
- assert op_registry["silu_and_mul"].enabled() == ops_enabled[1]
104
+ assert CustomOp.op_registry["silu_and_mul"].enabled() == ops_enabled[1]
105
106
assert GeluAndMul().enabled() == ops_enabled[2]
107
- assert op_registry["gelu_and_mul"].enabled() == ops_enabled[2]
107
+ assert CustomOp.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 op_registry["relu3"].enabled() == ops_enabled[3]
111
+ assert CustomOp.op_registry["relu3"].enabled() == ops_enabled[3]
112
113
# Unregistered subclass
114
class SiluAndMul2(SiluAndMul):
vllm/model_executor/custom_op.py
+14
-86
@@ -11,86 +11,6 @@ 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
-
14
class CustomOp(nn.Module):
15
"""
16
Base class for custom ops.
@@ -107,10 +27,10 @@ class CustomOp(nn.Module):
27
f"@CustomOp.register, or it's the CustomOp base class itself."
28
) from None
29
110
- if op_name not in op_registry_oot:
30
+ if op_name not in cls.op_registry_oot:
31
op_cls_to_instantiate = cls
32
else:
113
- op_cls_to_instantiate = op_registry_oot[op_name]
33
+ op_cls_to_instantiate = cls.op_registry_oot[op_name]
34
logger.debug(
35
"Instantiating custom op: %s using %s",
36
op_name,
@@ -230,13 +150,21 @@ class CustomOp(nn.Module):
150
151
return not count_none > 0 or count_all > 0
152
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
+
161
# Decorator to register custom ops.
162
@classmethod
163
def register(cls, name: str):
164
def decorator(op_cls):
237
- assert name not in op_registry, f"Duplicate op name: {name}"
165
+ assert name not in cls.op_registry, f"Duplicate op name: {name}"
166
op_cls.name = name
239
- op_registry[name] = op_cls
167
+ cls.op_registry[name] = op_cls
168
return op_cls
169
170
return decorator
@@ -254,9 +182,9 @@ class CustomOp(nn.Module):
182
def register_oot(cls, _decorated_op_cls=None, name: str | None = None):
183
def decorator(op_cls):
184
reg_name = name if name is not None else cls.__name__
257
- assert reg_name not in op_registry_oot, f"Duplicate op name: {reg_name}"
185
+ assert reg_name not in cls.op_registry_oot, f"Duplicate op name: {reg_name}"
186
op_cls.name = reg_name
259
- op_registry_oot[reg_name] = op_cls
187
+ cls.op_registry_oot[reg_name] = op_cls
188
return op_cls
189
190
if _decorated_op_cls is None:
vllm/model_executor/layers/mla.py
+11
-8
@@ -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 PluggableLayer
9
+from vllm.model_executor.custom_op import CustomOp
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
-@PluggableLayer.register("multi_head_latent_attention")
34
-class MultiHeadLatentAttentionWrapper(PluggableLayer):
35
- """Pluggable MLA layer which allows OOT backends to add
33
+@CustomOp.register("multi_head_latent_attention")
34
+class MultiHeadLatentAttentionWrapper(CustomOp):
35
+ """MLA layer registered as CustomOp to allow OOT backends to add
36
custom implementations of the outer MLA layer (including rope & o_proj).
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.
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.
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(PluggableLayer):
110
111
self.prefix = prefix
112
113
- def forward(
113
+ def forward_native(
114
self,
115
positions: torch.Tensor,
116
hidden_states: torch.Tensor,
@@ -174,3 +174,6 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
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)