model_executor/custom_op.py: CustomOp, PluggableLayer and their registries
Code version: vLLM @ 3e44c0d · 2026-10-22
Purpose#
vllm/model_executor/custom_op.py gives layers a way to pick a forward method per platform, and gives plugins a way to replace a layer class without a change to the model code.
| File | vllm/model_executor/custom_op.py, 356 lines of Python. Two module dicts, one helper function, two classes. |
|---|---|
| Main classes | PluggableLayer (vllm/model_executor/custom_op.py:32) and CustomOp (vllm/model_executor/custom_op.py:102), both subclasses of nn.Module. |
| Runs in | Every process that builds a model: each worker process, or the EngineCore process with the uni executor. The decorators also run at import time in any process that imports a layer module or a plugin. |
| Called by | The layer modules that subclass the two classes (50 files in vllm/ import from this module, from git grep), model constructors that build these layers, plugins that call register_oot, the LoRA layers (maybe_get_oot_by_class) and CompilationConfig.custom_op_log_check (op_registry). |
| Calls | get_cached_compilation_config() (vllm/model_executor/custom_op.py:9), current_platform (vllm/model_executor/custom_op.py:12), maybe_disable_graph_partition (vllm/model_executor/custom_op.py:11), torch.compile and torch._dynamo.mark_dynamic. |
| Not in this file | The default of the custom_ops list (config/vllm.py), its validation (config/compilation.py), the op implementations (the layer files), the second copy for hardware-agnostic layers (hw_agnostic/custom_op.py) and the IR op system (vllm/ir). |
The page CustomOp, out-of-tree ops and IR op providers explains the ideas, the default switches and the IR ops, with a worked example. This page reads the file from top to bottom.
Where it sits#
The file has three layers. At the top are the two registries (lines 16-29). PluggableLayer (lines 32-99) has only the class swap. CustomOp (lines 102-356) has the class swap and the per-platform dispatch.
Who calls it#
- Layer definitions. For example
@CustomOp.register("rms_norm")onRMSNorm(vllm/model_executor/layers/layernorm.py:38-39) and@PluggableLayer.register("column_parallel_linear")(vllm/model_executor/layers/linear.py:434-435). - Model constructors. They build the layers inside
set_current_vllm_config(vllm/model_executor/model_loader/utils.py:100-101). Construction runs__new__and__init__of this file. - Plugins.
@RotaryEmbedding.register_ootin the test platform (tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_custom_ops.py:10-11). - LoRA.
maybe_get_oot_by_class(QKVParallelLinear)and similar calls (vllm/lora/layers/column_parallel_linear.py:448). - Config check.
custom_op_log_checkimportsop_registryto tell a missing op from an op that is not in the model (vllm/config/compilation.py:1343-1352).
What it calls#
get_cached_compilation_config(): the currentCompilationConfig, cached until the config context changes (vllm/config/vllm.py:3570-3573).current_platform.is_rocm(),is_cpu(),is_tpu(),is_xpu(),is_out_of_tree()(vllm/model_executor/custom_op.py:194-205), andsimple_compile_backend(vllm/model_executor/custom_op.py:230-233).maybe_disable_graph_partition:{"graph_partition": False}for Inductor on torch 2.9 or newer (vllm/model_executor/utils.py:190-194).
Key classes and functions#
| Name | Line | What it does |
|---|---|---|
op_registry | vllm/model_executor/custom_op.py:21 | Registered op name → class, for CustomOp and PluggableLayer. |
op_registry_oot | vllm/model_executor/custom_op.py:22 | In-tree class name (or the given name) → out-of-tree class. |
maybe_get_oot_by_class(cls) | vllm/model_executor/custom_op.py:25-29 | The OOT replacement of a class, else the class. |
PluggableLayer.__new__ | vllm/model_executor/custom_op.py:46-65 | Builds the OOT class when the class name is in op_registry_oot. |
PluggableLayer.register(name) | vllm/model_executor/custom_op.py:68-76 | Sets name, adds to op_registry. |
PluggableLayer.register_oot | vllm/model_executor/custom_op.py:82-99 | Adds an OOT layer class. |
CustomOp.__new__ | vllm/model_executor/custom_op.py:107-126 | Same swap as above. |
CustomOp.__init__, forward | vllm/model_executor/custom_op.py:128-134 | Binds the method once; forward calls it. |
forward_native ... forward_oot | vllm/model_executor/custom_op.py:136-170 | Seven methods and their default bodies. |
dispatch_forward | vllm/model_executor/custom_op.py:172-205 | Enabled check + platform choice. |
maybe_compile | vllm/model_executor/custom_op.py:207-266 | Optional torch.compile of forward_native. |
enabled() | vllm/model_executor/custom_op.py:268-290 | +name / -name / base mode. |
default_on() | vllm/model_executor/custom_op.py:292-307 | Base mode: "all" → true, "none" → false. |
CustomOp.register(name, dynamic_arg_dims) | vllm/model_executor/custom_op.py:310-323 | Sets name and _dynamic_arg_dims, adds to op_registry. |
CustomOp.register_oot | vllm/model_executor/custom_op.py:334-356 | Adds an OOT op class; string positional → TypeError. |
Walk through the main methods#
CustomOp.__new__: the class swap#
def __new__(cls, *args, **kwargs):
try:
op_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"@CustomOp.register, or it's the CustomOp base class itself."
) from None
if op_name not in op_registry_oot:
op_cls_to_instantiate = cls
else:
op_cls_to_instantiate = op_registry_oot[op_name]
logger.debug(
"Instantiating custom op: %s using %s",
op_name,
str(op_cls_to_instantiate),
)
return super().__new__(op_cls_to_instantiate)- Lines 108-115. The
tryblock readscls.__name__. Every Python class has__name__, so theAttributeErrorbranch, with its text about thenameattribute, cannot run for a class. (Derived from Python semantics.) - Line 117. The key is the class name of the object that the code builds, for example
"RotaryEmbedding". A subclass with another name, such as a scaling variant of rotary embedding, is not swapped (derived). - Lines 120-125. The OOT class is used, with a DEBUG log.
- Line 126. The parent
__new__makes an object of the chosen class. Python calls__init__next only if that object is an instance of the called class, so the OOT class must subclass the in-tree class (general Python rule).
dispatch_forward: the decision#
def dispatch_forward(self, compile_native: bool):
# NOTE(woosuk): Here we assume that vLLM was built for only one
# specific backend. Currently, we do not support dynamic dispatching.
compilation_config = get_cached_compilation_config()
# NOTE(shen-shanshan): CustomOp object can be enforce enabled, e.g.,
# enable device-specific kernels in ViT models when enabling graph
# mode. By default, it will follow the compilation_config to determine
# whether enable itself.
# This enforce_enable mechanism will be removed after we adding a
# separate compilation_config for multi-modal part.
enabled = self._enforce_enable or self.enabled()
if enabled:
compilation_config.enabled_custom_ops.update([self.__class__.name])
else:
compilation_config.disabled_custom_ops.update([self.__class__.name])
if not enabled:
# Compile forward_native to avoid eager torch ops if inside
# opaque torch custom op (e.g. fused_moe, unified_attention, etc.)
return self.maybe_compile(self.forward_native, enable=compile_native)
if current_platform.is_rocm():
return self.forward_hip
elif current_platform.is_cpu():
return self.forward_cpu
elif current_platform.is_tpu():
return self.forward_tpu
elif current_platform.is_xpu():
return self.forward_xpu
elif current_platform.is_out_of_tree():
return self.forward_oot
else:
return self.forward_cuda- Line 175. The comment above says that vLLM is built for one backend and has no dynamic dispatch. The config comes from the current config context.
- Line 183.
enforce_enablefrom__init__overrides the list. The comment says it exists for the vision encoder of multimodal models and will be removed (lines 177-182). - Lines 184-187. The counters
enabled_custom_opsanddisabled_custom_opsrecordself.__class__.name. - Lines 189-192. Disabled:
forward_native, compiled only withcompile_native=True. - Lines 194-205. Enabled: the first matching platform test wins; CUDA is the final
else. AnUnspecifiedPlatformtherefore getsforward_cuda(derived).
enabled and default_on#
@classmethod
def enabled(cls) -> bool:
# if no name, then it was not registered
compilation_config = get_cached_compilation_config()
custom_ops = compilation_config.custom_ops
if not hasattr(cls, "name"):
logger.warning_once(
"Custom op %s was not registered, which means it won't appear "
"in the op registry. It will be enabled/disabled based on the "
"global settings.",
cls.__name__,
)
return CustomOp.default_on()
enabled = f"+{cls.name}" in custom_ops
disabled = f"-{cls.name}" in custom_ops
if enabled and disabled:
raise ValueError(
"custom_ops cannot both enable and disable the same operation: "
f"{cls.name}. Remove either the '+' or '-' directive"
)
return (CustomOp.default_on() or enabled) and not disabled
@staticmethod
def default_on() -> bool:
"""Behavior controlled by `CompilationConfig.custom_ops`: On by default if
'all', off by default if 'none'.
When PyTorch Inductor is used, 'none' is the default value,
otherwise 'all'.
"""
compilation_config = get_cached_compilation_config()
count_none = compilation_config.custom_ops.count("none")
count_all = compilation_config.custom_ops.count("all")
if count_none + count_all != 1:
raise ValueError(
"custom_ops must contain exactly one base mode: 'all' or 'none'"
)
return not count_none > 0 or count_all > 0- Line 273. A class with no
nameattribute was never registered: WARNING once, and the base mode decides. - Lines 282-289. Both
+nameand-name:ValueError.CompilationConfigalready rejects this earlier (vllm/config/compilation.py:1027-1034). - Line 290.
(default_on() or enabled) and not disabled. - Lines 300-305. Exactly one base mode, else
ValueError. - Line 307.
not count_none > 0 or count_all > 0: true for"all", false for"none".
register_oot: the two working forms#
@classmethod
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 op_registry_oot, f"Duplicate op name: {reg_name}"
op_cls.name = reg_name
op_registry_oot[reg_name] = op_cls
return op_cls
if _decorated_op_cls is None:
# Called with parentheses: @CustomOP.register_oot()
# or @CustomOP.register_oot(name="...")
# So, _decorated_op_cls is None.
# We return the actual decorator function.
return decorator
elif isinstance(_decorated_op_cls, type): # Check if it's a class
# Called without parentheses: @CustomOP.register_oot
# The first argument is the class itself.
# We call the 'decorator' function immediately with the class.
return decorator(_decorated_op_cls)
else:
# Handle other unexpected cases if necessary
raise TypeError("Decorator can only be applied to classes.")- Line 337. The key is
name, or the name of the class on which the decorator is called. OnCustomOp.register_oot()withoutname, the key would be"CustomOp"(derived). - Line 338. A duplicate key fails the
assert. - Line 339. The OOT class's
namebecomes the key. The+nameswitch for it is then+RotaryEmbedding, not the in-tree op name (derived; see CustomOp, out-of-tree ops and IR op providers). - Lines 343-356.
None→ return the decorator; a class → decorate it now; anything else (a string) →TypeError.
PluggableLayer.register_oot (lines 82-99) has the same structure, with the text "Duplicate layer name".
Deep dive#
The sections below give the full detail. Open them one at a time, or use Open all deep dives at the top of the page.
Deep diveData: the registries and the attributes this file sets
| Value | Type | Set at |
|---|---|---|
op_registry | dict[str, type], key = registered name such as "rms_norm" | vllm/model_executor/custom_op.py:73, vllm/model_executor/custom_op.py:320 |
op_registry_oot | dict[str, type], key = class name such as "RotaryEmbedding" | vllm/model_executor/custom_op.py:88, vllm/model_executor/custom_op.py:340 |
cls.name | str class attribute; inherited by subclasses | vllm/model_executor/custom_op.py:72, vllm/model_executor/custom_op.py:318, vllm/model_executor/custom_op.py:339 |
cls._dynamic_arg_dims | dict[str, int | list[int]] | None | vllm/model_executor/custom_op.py:319 |
self._enforce_enable, self._forward_method | bool, bound method or compiled function | vllm/model_executor/custom_op.py:130-131 |
Both dicts are module globals, so each process has its own. A plugin must register in every process that builds a model: a general plugin does this (How vLLM finds and loads plugins). Because name is a class attribute, a subclass without its own register call uses the parent's name. RotaryEmbedding is such a subclass: the decorator is on RotaryEmbeddingBase (vllm/model_executor/layers/rotary_embedding/base.py:14-15, vllm/model_executor/layers/rotary_embedding/base.py:139).
Deep divemaybe_compile, line by line
def maybe_compile(self, fn, *, enable: bool = True):
"""Compile fn if compilation enabled.
Useful for CustomOp instances called from within a torch custom op,
meaning the forward call is hidden from the model-level torch.compile.
NOTE: this does not enable fusion across ops, so opaque custom ops
should still be unwrapped wherever possible.
"""
from vllm.config.compilation import CompilationMode
# Do not compile if compilation disabled
if not enable:
return fn
# Do not compile if global compilation disabled
compilation_config = get_cached_compilation_config()
if compilation_config.mode == CompilationMode.NONE:
return fn
# If eager backend is used, do not compile either
if compilation_config.backend == "eager":
return fn
compile_options = maybe_disable_graph_partition(
current_platform.simple_compile_backend
)
backend = current_platform.simple_compile_backend
dynamic_arg_dims = getattr(self.__class__, "_dynamic_arg_dims", None)
if dynamic_arg_dims is not None:
compiled_fn = torch.compile(
fn,
dynamic=False,
backend=backend,
options=compile_options,
)
sig = inspect.signature(fn)
@functools.wraps(fn)
def wrapper(*args, **kwargs):
bound = sig.bind(*args, **kwargs)
bound.apply_defaults()
for name, dims in dynamic_arg_dims.items():
arg = bound.arguments.get(name)
if arg is not None and isinstance(arg, torch.Tensor):
dims_list = [dims] if isinstance(dims, int) else dims
for d in dims_list:
real_d = arg.ndim + d if d < 0 else d
torch._dynamo.mark_dynamic(arg, real_d)
return compiled_fn(*args, **kwargs)
return wrapper
# dynamic=True to avoid recompilations
return torch.compile(
fn,
dynamic=True,
backend=backend,
options=compile_options,
)- Lines 218-228. Return
fnunchanged when compilation is off for this call, for the whole config (modeNONE), or with the"eager"backend. - Lines 230-233. The backend is the platform's
simple_compile_backend("inductor"by default,vllm/platforms/interface.py:172), with graph partition off. - Lines 235-258. With
dynamic_arg_dims:dynamic=False, and each call marks the listed dimensions as dynamic. A negative index counts from the end (line 254). An argument that is not a tensor is skipped (line 251). - Lines 261-266. Else
dynamic=True.
Example: mla_decode_concat_quant_fp8 registers dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0} (vllm/model_executor/layers/attention/mla_attention.py:1561-1564): dimension 0 (tokens) of both tensors is dynamic.
Deep diveProcesses and timing
- Import time.
registerruns when Python imports a layer module.register_ootruns when the plugin module is imported. - Model build.
__new__,__init__anddispatch_forwardrun once per layer object, insideset_current_vllm_config. A CustomOp built when no config is set raisesAssertionError(vllm/config/vllm.py:3576-3586). - Forward.
forwardonly calls_forward_method. No config lookup happens per step. - Threads. The registries have no lock. Registration happens at import and plugin time, before the model build. (Derived; the file has no comment about threads.)
Deep diveError paths: what raises, what only logs
| Situation | Result | Where |
|---|---|---|
Duplicate register name | AssertionError "Duplicate op name" | vllm/model_executor/custom_op.py:317, vllm/model_executor/custom_op.py:71 |
| Duplicate OOT key | AssertionError | vllm/model_executor/custom_op.py:338, vllm/model_executor/custom_op.py:86 |
register_oot("X") | TypeError | vllm/model_executor/custom_op.py:354-356 |
| Unregistered CustomOp subclass | WARNING once; base mode decides | vllm/model_executor/custom_op.py:273-280 |
+x and -x both | ValueError | vllm/model_executor/custom_op.py:284-288 |
| No or two base modes | ValueError | vllm/model_executor/custom_op.py:302-305 |
A forward_cuda or forward_native that the subclass does not define, when chosen | NotImplementedError at the first forward | vllm/model_executor/custom_op.py:136-145 |
Deep diveConfiguration that this file reads
| Setting | Default | Set at | Effect here |
|---|---|---|---|
compilation_config.custom_ops | "none" is added with Inductor + compile mode, else "all"; some code adds +name entries (CustomOp, out-of-tree ops and IR op providers) | vllm/config/vllm.py:1819-1826 | enabled(), default_on() |
compilation_config.mode | VLLM_COMPILE above O0; NONE with --enforce-eager or breakable CUDA graphs | vllm/config/vllm.py:1806-1810, vllm/config/vllm.py:1745-1746, vllm/config/vllm.py:849-851 | maybe_compile line 223 |
compilation_config.backend | platform's compile backend | vllm/config/compilation.py:1080-1081 | maybe_compile line 227 |
No environment variable is read in this file.
Related concepts#
- CustomOp, out-of-tree ops and IR op providers: the switches, the default, the out-of-tree flow and the IR ops.
- The Platform interface:
is_out_of_tree()and the platform hooks that can changecustom_ops. - All the ways to extend vLLM:
register_ootamong the other extension points. - torch.compile and CUDA graphs in vLLM: how
torch.compilesees disabled ops.
Check yourself#
- 1Which dict doeseasy
__new__read, and with which key?Show answer
op_registry_oot, withcls.__name__(vllm/model_executor/custom_op.py:109,vllm/model_executor/custom_op.py:117-120). - 2An enabled op runs on ROCm. The subclass defines onlymedium
forward_nativeandforward_cuda. Which method runs?Show answer
dispatch_forwardreturnsforward_hip(vllm/model_executor/custom_op.py:194-195). Its default body callsforward_cuda(vllm/model_executor/custom_op.py:147-149). - 3What is the difference betweenmedium
PluggableLayerandCustomOp?Show answer
Both swap the class in
__new__throughop_registry_oot. OnlyCustomOphas theforward_*methods anddispatch_forward(vllm/model_executor/custom_op.py:33-44). - 4An op registered withhard
dynamic_arg_dims={"x": -1}is disabled and built withcompile_native=TrueunderVLLM_COMPILE. What does a call with a 2-D tensorxdo?Show answer
maybe_compilecompiledforward_nativewithdynamic=False. The wrapper binds the arguments, turns-1intox.ndim - 1 = 1, callsmark_dynamic(x, 1), then calls the compiled function (vllm/model_executor/custom_op.py:235-258).